Merge pull request #188 from squareup/prudhvi.remove_embedded_type

Revert "Merge pull request #183 from squareup/prudhvi.watermark_prior…

Author: prudhvi

pkg/table/chunker.go

@@ -34,7 +34,7 @@ type Chunker interface {
 	IsRead() bool
 	Close() error
 	Next() (*Chunk, error)
-	Feedback(chunk *Chunk, d time.Duration)
+	Feedback(*Chunk, time.Duration)
 	GetLowWatermark() (string, error)
 	KeyAboveHighWatermark(interface{}) bool
 }
@@ -47,12 +47,10 @@ func NewChunker(t *TableInfo, chunkerTarget time.Duration, logger loggers.Advanc
 	// tables with a single column key.
 	if len(t.KeyColumns) == 1 && t.KeyIsAutoInc {
 		return &chunkerOptimistic{
-			coreChunker: &coreChunker{
-				Ti:                     t,
-				ChunkerTarget:          chunkerTarget,
-				lowerBoundWatermarkMap: make(map[string]*Chunk, 0),
-				logger:                 logger,
-			},
+			Ti:                     t,
+			ChunkerTarget:          chunkerTarget,
+			lowerBoundWatermarkMap: make(map[string]*Chunk, 0),
+			logger:                 logger,
 		}, nil
 	}
 	return NewCompositeChunker(t, chunkerTarget, logger, "", "")
@@ -62,12 +60,10 @@ func NewChunker(t *TableInfo, chunkerTarget time.Duration, logger loggers.Advanc
 // setting its Key if keyName and where conditions are provided
 func NewCompositeChunker(t *TableInfo, chunkerTarget time.Duration, logger loggers.Advanced, keyName string, whereCondition string) (Chunker, error) {
 	c := chunkerComposite{
-		coreChunker: &coreChunker{
-			Ti:                     t,
-			ChunkerTarget:          chunkerTarget,
-			lowerBoundWatermarkMap: make(map[string]*Chunk, 0),
-			logger:                 logger,
-		},
+		Ti:                     t,
+		ChunkerTarget:          chunkerTarget,
+		lowerBoundWatermarkMap: make(map[string]*Chunk, 0),
+		logger:                 logger,
 	}
 	var err error
 	if keyName != "" && whereCondition != "" {

pkg/table/chunker_composite.go

@@ -6,19 +6,37 @@ import (
 	"fmt"
 	"reflect"
 	"strings"
+	"sync"
 	"time"
 
+	"github.com/siddontang/loggers"
 	"golang.org/x/exp/slices"
 )
 
 type chunkerComposite struct {
-	*coreChunker
+	sync.Mutex
+	Ti             *TableInfo
+	chunkSize      uint64
+	chunkPtrs      []Datum  // a list of Ptrs for each of the keys.
+	chunkKeys      []string // all the keys to chunk on (usually all the col names of the PK)
+	keyName        string   // the name of the key we are chunking on
+	where          string   // any additional WHERE conditions.
+	finalChunkSent bool
+	isOpen         bool
 
-	chunkSize uint64
-	chunkPtrs []Datum  // a list of Ptrs for each of the keys.
-	chunkKeys []string // all the keys to chunk on (usually all the col names of the PK)
-	keyName   string   // the name of the key we are chunking on
-	where     string   // any additional WHERE conditions.
+	// Dynamic Chunking is time based instead of row based.
+	// It uses *time* to determine the target chunk size.
+	chunkTimingInfo []time.Duration
+	ChunkerTarget   time.Duration // i.e. 500ms for target
+
+	// This is used for restore.
+	watermark *Chunk
+	// Map from lowerbound value of a chunk -> chunk,
+	// Used to update the watermark by applying stored chunks,
+	// by comparing their lowerBound with current watermark upperBound.
+	lowerBoundWatermarkMap map[string]*Chunk
+
+	logger loggers.Advanced
 }
 
 var _ Chunker = &chunkerComposite{}
@@ -227,6 +245,91 @@ func (t *chunkerComposite) Feedback(chunk *Chunk, d time.Duration) {
 	}
 }
 
+// GetLowWatermark returns the highest known value that has been safely copied,
+// which (due to parallelism) could be significantly behind the high watermark.
+// The value is discovered via ChunkerFeedback(), and when retrieved from this func
+// can be used to write a checkpoint for restoration.
+func (t *chunkerComposite) GetLowWatermark() (string, error) {
+	t.Lock()
+	defer t.Unlock()
+	if t.watermark == nil || t.watermark.UpperBound == nil || t.watermark.LowerBound == nil {
+		return "", errors.New("watermark not yet ready")
+	}
+
+	return t.watermark.JSON(), nil
+}
+
+// isSpecialRestoredChunk is used to test for the first chunk after restore-from-checkpoint.
+// The restored chunk is a really special beast because the lowerbound
+// will be repeated by the first chunk that is applied post restore.
+// This is called under a mutex.
+func (t *chunkerComposite) isSpecialRestoredChunk(chunk *Chunk) bool {
+	if chunk.LowerBound == nil || chunk.UpperBound == nil || t.watermark == nil || t.watermark.LowerBound == nil || t.watermark.UpperBound == nil {
+		return false // restored checkpoints always have both.
+	}
+	return chunk.LowerBound.comparesTo(t.watermark.LowerBound)
+}
+
+// bumpWatermark updates the minimum value that is known to be safely copied,
+// and is called under a mutex.
+// Because of parallelism, it is possible that a chunk is copied out of order,
+// so this func needs to account for that.
+// Basically:
+//   - If the chunk does not "align" to the current low watermark, it's stored in a map keyed by its lowerBound valuesString() value.
+//   - If it does align, the watermark is bumped to the chunk's max value. Then
+//     stored chunk map is checked to see if an existing chunk lowerBound aligns with the new watermark.
+//   - If any stored chunk aligns, it is deleted off the map and the watermark is bumped.
+//   - This process repeats until there is no more alignment from the stored map *or* the map is empty.
+func (t *chunkerComposite) bumpWatermark(chunk *Chunk) {
+	if chunk.UpperBound == nil {
+		return
+	}
+	// Check if this is the first chunk or it's the special restored chunk.
+	// If so, set the watermark and then go on to applying any stored chunks.
+	if (t.watermark == nil && chunk.LowerBound == nil) || t.isSpecialRestoredChunk(chunk) {
+		t.watermark = chunk
+		goto applyStoredChunks
+	}
+
+	// Validate that chunk has lower bound before moving on
+	if chunk.LowerBound == nil {
+		errMsg := fmt.Sprintf("coreChunker.bumpWatermark: nil lowerBound value encountered more than once: %v", chunk)
+		t.logger.Fatal(errMsg)
+	}
+
+	// We haven't set the first chunk yet, or it's not aligned with the
+	// previous watermark. Store it in the map keyed by its lowerBound, and move on.
+
+	// We only need to store by lowerBound because, when updating watermark
+	// we always compare the upperBound of current watermark to lowerBound of stored chunks.
+	// Key can never be nil, because first chunk will not hit this code path and all remaining chunks will have lowerBound.
+	if t.watermark == nil || !t.watermark.UpperBound.comparesTo(chunk.LowerBound) {
+		t.lowerBoundWatermarkMap[chunk.LowerBound.valuesString()] = chunk
+		return
+	}
+
+	// The remaining case is:
+	// t.watermark.UpperBound.Value == chunk.LowerBound.Value
+	// Replace the current watermark with the chunk.
+	t.watermark = chunk
+
+applyStoredChunks:
+
+	// Check the waterMarkMap for any chunks that align with the new watermark.
+	// If there are any, bump the watermark and delete from the map.
+	// If there are none, we're done.
+	for t.waterMarkMapNotEmpty() && t.watermark.UpperBound != nil && t.lowerBoundWatermarkMap[t.watermark.UpperBound.valuesString()] != nil {
+		key := t.watermark.UpperBound.valuesString()
+		nextWatermark := t.lowerBoundWatermarkMap[key]
+		t.watermark = nextWatermark
+		delete(t.lowerBoundWatermarkMap, key)
+	}
+}
+
+func (t *chunkerComposite) waterMarkMapNotEmpty() bool {
+	return t.lowerBoundWatermarkMap != nil && len(t.lowerBoundWatermarkMap) != 0
+}
+
 func (t *chunkerComposite) open() (err error) {
 	if t.isOpen {
 		// This prevents an error where open is re-called
@@ -316,17 +419,3 @@ func (t *chunkerComposite) SetKey(keyName string, where string) error {
 	t.where = where
 	return nil
 }
-
-// GetLowWatermark returns the highest known value that has been safely copied,
-// which (due to parallelism) could be significantly behind the high watermark.
-// The value is discovered via Chunker Feedback(), and when retrieved from this func
-// can be used to write a checkpoint for restoration.
-func (t *chunkerComposite) GetLowWatermark() (string, error) {
-	t.Lock()
-	defer t.Unlock()
-	if t.watermark == nil || t.watermark.UpperBound == nil || t.watermark.LowerBound == nil {
-		return "", errors.New("watermark not yet ready")
-	}
-
-	return t.watermark.JSON(), nil
-}

pkg/table/chunker_composite_test.go

@@ -270,12 +270,10 @@ func TestCompositeLowWatermark(t *testing.T) {
 	assert.NoError(t, t1.SetInfo(context.Background()))
 
 	chunker := &chunkerComposite{
-		coreChunker: &coreChunker{
-			Ti:                     t1,
-			ChunkerTarget:          ChunkerDefaultTarget,
-			lowerBoundWatermarkMap: make(map[string]*Chunk, 0),
-			logger:                 logrus.New(),
-		},
+		Ti:                     t1,
+		ChunkerTarget:          ChunkerDefaultTarget,
+		lowerBoundWatermarkMap: make(map[string]*Chunk, 0),
+		logger:                 logrus.New(),
 	}
 	_, err = chunker.Next()
 	assert.Error(t, err) // not open yet
@@ -430,11 +428,9 @@ func TestSetKey(t *testing.T) {
 	t1 := NewTableInfo(db, "test", "setkey_t1")
 	assert.NoError(t, t1.SetInfo(context.Background()))
 	chunker := &chunkerComposite{
-		coreChunker: &coreChunker{
-			Ti:            t1,
-			ChunkerTarget: 100 * time.Millisecond,
-			logger:        logrus.New(),
-		},
+		Ti:            t1,
+		ChunkerTarget: 100 * time.Millisecond,
+		logger:        logrus.New(),
 	}
 	err = chunker.SetKey("s", "status = 'ARCHIVED' AND updated_at < NOW() - INTERVAL 1 DAY")
 	assert.NoError(t, err)
@@ -449,11 +445,9 @@ func TestSetKey(t *testing.T) {
 
 	// If I reset again with a different condition it should range as chunks.
 	chunker = &chunkerComposite{
-		coreChunker: &coreChunker{
-			logger:        logrus.New(),
-			ChunkerTarget: 100 * time.Millisecond,
-			Ti:            t1,
-		},
+		Ti:            t1,
+		ChunkerTarget: 100 * time.Millisecond,
+		logger:        logrus.New(),
 	}
 	err = chunker.SetKey("s", "status = 'PENDING' AND updated_at > NOW() - INTERVAL 1 DAY")
 	assert.NoError(t, err)
@@ -484,11 +478,9 @@ func TestSetKey(t *testing.T) {
 	// Test other index types.
 	for _, index := range []string{"u", "su", "ui"} {
 		chunker = &chunkerComposite{
-			coreChunker: &coreChunker{
-				Ti:            t1,
-				ChunkerTarget: 100 * time.Millisecond,
-				logger:        logrus.New(),
-			},
+			Ti:            t1,
+			ChunkerTarget: 100 * time.Millisecond,
+			logger:        logrus.New(),
 		}
 		err = chunker.SetKey(index, "updated_at < NOW() - INTERVAL 1 DAY")
 		assert.NoError(t, err)
@@ -554,11 +546,9 @@ func TestSetKeyCompositeKeyMerge(t *testing.T) {
 	t1 := NewTableInfo(db, "test", "setkeycomposite_t1")
 	assert.NoError(t, t1.SetInfo(context.Background()))
 	chunker := &chunkerComposite{
-		coreChunker: &coreChunker{
-			logger:        logrus.New(),
-			ChunkerTarget: 100 * time.Millisecond,
-			Ti:            t1,
-		},
+		Ti:            t1,
+		ChunkerTarget: 100 * time.Millisecond,
+		logger:        logrus.New(),
 	}
 	err = chunker.SetKey("dnc", "")
 	assert.NoError(t, err)