rosmar changes prior to CBG-5041

feeds.go

@@ -11,6 +11,7 @@ package rosmar
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"expvar"
 	"fmt"
 	"sync/atomic"
@@ -25,32 +26,30 @@ var activeFeedCount int32 // for tests
 
 func (bucket *Bucket) StartDCPFeed(ctx context.Context, args sgbucket.FeedArguments, callback sgbucket.FeedEventCallbackFunc, dbStats *expvar.Map) error {
 	traceEnter("StartDCPFeed", "bucket=%s, args=%+v", bucket.GetName(), args)
-	// If no scopes are specified, return feed for the default collection, if it exists
+	// Validate requested collections exist before starting feeds. No scopes means the default
+	// collection, if it exists. A collection named twice gets one feed.
+	requestedCollections := map[*Collection]struct{}{}
 	if len(args.Scopes) == 0 {
 		collection, err := bucket.getCollection(defaultDataStoreName)
 		if err != nil {
 			return err
 		}
-		_, err = collection.startDCPFeed(ctx, args, callback, dbStats)
-		return err
+		requestedCollections[collection] = struct{}{}
 	}
-
-	// Validate requested collections exist before starting feeds
-	requestedCollections := make([]*Collection, 0)
 	for scopeName, collections := range args.Scopes {
 		for _, collectionName := range collections {
 			collection, err := bucket.getCollection(sgbucket.DataStoreNameImpl{Scope: scopeName, Collection: collectionName})
 			if err != nil {
 				return fmt.Errorf("couldn't open collection %s:%s for DCP feed: %w", scopeName, collectionName, err)
 			}
-			requestedCollections = append(requestedCollections, collection)
+			requestedCollections[collection] = struct{}{}
 		}
 	}
 
 	doneChan := args.DoneChan
-	doneChans := map[*Collection]chan struct{}{}
+	doneChans := make(map[*Collection]chan struct{}, len(requestedCollections))
 	startedFeeds := make([]*dcpFeed, 0, len(requestedCollections))
-	for _, collection := range requestedCollections {
+	for collection := range requestedCollections {
 		// Not bothering to remove scopes from args for the single collection feeds
 		// here because it's ignored by Collection.startDCPFeed
 		collectionID := collection.GetCollectionID()
@@ -77,8 +76,12 @@ func (bucket *Bucket) StartDCPFeed(ctx context.Context, args sgbucket.FeedArgume
 
 	// coalesce doneChans
 	go func() {
-		for _, collection := range requestedCollections {
-			<-doneChans[collection]
+		for _, done := range doneChans {
+			<-done
+		}
+		// The collections share one checkpoint document, so a dump saves it once they have all stopped.
+		if args.Dump {
+			_ = saveDumpCheckpoint(startedFeeds)
 		}
 		if doneChan != nil {
 			close(doneChan)
@@ -113,6 +116,9 @@ func (c *Collection) startDCPFeed(ctx context.Context, args sgbucket.FeedArgumen
 			if args.CheckpointPrefix == "" {
 				return nil, fmt.Errorf("feed's Backfill is FeedResume but no CheckpointPrefix given")
 			}
+			if args.MetadataStore == nil {
+				return nil, fmt.Errorf("feed's Backfill is FeedResume but no MetadataStore given")
+			}
 			if err := feed.readCheckpoint(); err != nil {
 				return nil, fmt.Errorf("couldn't read DCP feed checkpoint: %w", err)
 			}
@@ -217,6 +223,8 @@ type dcpFeed struct {
 	events         eventQueue
 	lastCas        CAS
 	lastCasChanged bool
+	// completed reports that the feed drained to its end-of-feed marker, rather than being terminated.
+	completed atomic.Bool
 }
 
 func (feed *dcpFeed) String() string {
@@ -227,15 +235,20 @@ func (feed *dcpFeed) checkpointKey() string {
 	return feed.args.CheckpointPrefix
 }
 
+// usesCheckpoint reports whether the feed has somewhere to keep a checkpoint.
+func (feed *dcpFeed) usesCheckpoint() bool {
+	return feed.args.CheckpointPrefix != "" && feed.metadataStore != nil
+}
+
 // Reads the feed's lastCas from the checkpoint document, if there is one.
 func (feed *dcpFeed) readCheckpoint() (err error) {
-	if feed.args.CheckpointPrefix == "" {
+	if !feed.usesCheckpoint() {
 		return
 	}
 	key := feed.checkpointKey()
 	var checkpt checkpoint
 	if _, err = feed.metadataStore.Get(feed.ctx, key, &checkpt); err != nil {
-		if _, ok := err.(sgbucket.MissingError); ok {
+		if _, ok := errors.AsType[sgbucket.MissingError](err); ok {
 			err = nil
 			debug("%s checkpoint %q missing", feed, key)
 		} else {
@@ -252,14 +265,32 @@ func (feed *dcpFeed) readCheckpoint() (err error) {
 }
 
 // Writes the feed's lastCas to the checkpoint document, if there is one.
-func (feed *dcpFeed) writeCheckpoint() (err error) {
-	if feed.args.CheckpointPrefix == "" || !feed.lastCasChanged {
-		return
+func (feed *dcpFeed) writeCheckpoint() error {
+	if !feed.lastCasChanged {
+		return nil
 	}
-	key := feed.checkpointKey()
-	colID := feed.collection.GetCollectionID()
+	return updateCheckpoint([]*dcpFeed{feed}, false)
+}
 
-	_, err = feed.metadataStore.Update(feed.ctx, key, 0, func(current []byte) (updated []byte, expiry *uint32, delete bool, err error) {
+// saveDumpCheckpoint records a dump's progress. A resumed dump whose collections all streamed to the end has
+// nothing left to resume, so it removes their entries instead.
+func saveDumpCheckpoint(feeds []*dcpFeed) error {
+	if len(feeds) == 0 {
+		return nil
+	}
+	remove := feeds[0].args.Backfill == sgbucket.FeedResume && allFeedsCompleted(feeds)
+	return updateCheckpoint(feeds, remove)
+}
+
+// updateCheckpoint writes each feed's lastCas to their shared checkpoint document in one update. With remove,
+// it drops the feeds' entries instead, and deletes the document once none are left.
+func updateCheckpoint(feeds []*dcpFeed, remove bool) error {
+	feed := feeds[0]
+	if !feed.usesCheckpoint() {
+		return nil
+	}
+	key := feed.checkpointKey()
+	_, err := feed.metadataStore.Update(feed.ctx, key, 0, func(current []byte) (updated []byte, expiry *uint32, del bool, err error) {
 		var checkpt checkpoint
 		if current != nil {
 			if err := json.Unmarshal(current, &checkpt); err != nil {
@@ -269,8 +300,24 @@ func (feed *dcpFeed) writeCheckpoint() (err error) {
 		if checkpt.LastCas == nil {
 			checkpt.LastCas = make(map[uint32]uint64)
 		}
-		checkpt.LastCas[colID] = feed.lastCas
-
+		changed := false
+		for _, f := range feeds {
+			colID := f.collection.GetCollectionID()
+			if remove {
+				_, found := checkpt.LastCas[colID]
+				changed = changed || found
+				delete(checkpt.LastCas, colID)
+			} else if f.lastCasChanged {
+				checkpt.LastCas[colID] = f.lastCas
+				changed = true
+			}
+		}
+		if !changed {
+			return nil, nil, false, nil // cancels the update
+		}
+		if len(checkpt.LastCas) == 0 {
+			return nil, nil, true, nil
+		}
 		updated, err = json.Marshal(checkpt)
 		if err != nil {
 			return nil, nil, false, fmt.Errorf("failed to marshal checkpoint: %w", err)
@@ -279,9 +326,9 @@ func (feed *dcpFeed) writeCheckpoint() (err error) {
 	})
 
 	if err == nil {
-		debug("%s wrote lastCas 0x%x to %q", feed, feed.lastCas, key)
+		debug("%s updated checkpoint %q, remove=%v", feed, key, remove)
 	} else {
-		logError("%s failed to write checkpoint to %q: %v", feed, key, err)
+		logError("%s failed to update checkpoint %q: %v", feed, key, err)
 	}
 	return err
 }
@@ -305,30 +352,58 @@ func (feed *dcpFeed) run() {
 	collectionID := feed.collection.GetCollectionID()
 	feedContent := feed.args.FeedContent
 	for {
-		if e := feed.events.pull(); e != nil {
-			feedEvent, err := e.asFeedEvent(collectionID, feedContent)
-			if err != nil {
-				logError("Fatal error converting %s event to feed event: %v", feed, err)
-				break
-			}
-			feed.callback(*feedEvent)
-			if feedEvent.Cas > feed.lastCas {
-				feed.lastCas = feedEvent.Cas
-				feed.lastCasChanged = true
-				debug("%s lastCas = 0x%x", feed, feed.lastCas)
-				// TODO: Set a timer to write the checkpoint "soon"
-			}
-		} else {
+		// Stop before the next event once the Terminator is closed, without waiting for the goroutine above.
+		if feed.terminated() {
+			break
+		}
+		e, ok := feed.events.pull()
+		if !ok {
+			break // terminated
+		}
+		if e == nil {
+			// Stored before the deferred close of DoneChan, so a waiter sees a settled value.
+			feed.completed.Store(true)
+			break
+		}
+		feedEvent, err := e.asFeedEvent(collectionID, feedContent)
+		if err != nil {
+			logError("Fatal error converting %s event to feed event: %v", feed, err)
 			break
 		}
+		feed.callback(*feedEvent)
+		if feedEvent.Cas > feed.lastCas {
+			feed.lastCas = feedEvent.Cas
+			feed.lastCasChanged = true
+			debug("%s lastCas = 0x%x", feed, feed.lastCas)
+			// TODO: Set a timer to write the checkpoint "soon"
+		}
 	}
 	debug("%s stopping", feed)
 
-	if feed.lastCasChanged {
-		if err := feed.writeCheckpoint(); err != nil {
-			logError("Error saving %s checkpoint: %v", feed, err)
+	// A dump saves its checkpoint in StartDCPFeed, once all of its collections stop.
+	if !feed.args.Dump {
+		_ = feed.writeCheckpoint()
+	}
+}
+
+// terminated reports whether the feed's Terminator is closed. A feed with no Terminator is never terminated.
+func (feed *dcpFeed) terminated() bool {
+	select {
+	case <-feed.args.Terminator:
+		return true
+	default:
+		return false
+	}
+}
+
+// allFeedsCompleted reports whether every feed drained to its end-of-feed marker.
+func allFeedsCompleted(feeds []*dcpFeed) bool {
+	for _, feed := range feeds {
+		if !feed.completed.Load() {
+			return false
 		}
 	}
+	return true
 }
 
 func (feed *dcpFeed) close() {