GO-7373: reconnect the mongo change stream instead of dying silently

db/db.go

@@ -26,6 +26,12 @@ const CName = "consensus.db"
 const (
 	settingsColl = "settings"
 	payloadColl  = "payload"
+
+	// maxReconnectWait is the upper bound of the wait between change stream reconnect attempts
+	maxReconnectWait = time.Second * 30
+	// healthyStreamTime is how long a change stream must live to be considered healthy, a stream that dies
+	// sooner keeps growing the reconnect backoff instead of resetting it
+	healthyStreamTime = time.Second * 10
 )
 
 var log = logger.NewNamed(CName)
@@ -36,6 +42,11 @@ func New() Service {
 
 type ChangeReceiver func(logId string, records []consensus.Record)
 
+// ResetReceiver is called when the change stream has been re-established after an interruption.
+// Changes made while the stream was down are not replayed, so the receiver must resync from the db.
+// The ctx is cancelled when the service is closing, the receiver must not outlive it.
+type ResetReceiver func(ctx context.Context)
+
 type Service interface {
 	// AddLog adds new log db
 	AddLog(ctx context.Context, log consensus.Log) (err error)
@@ -48,6 +59,9 @@ type Service interface {
 	FetchLog(ctx context.Context, logId, afterRecordId string) (log consensus.Log, err error)
 	// SetChangeReceiver sets the receiver for updates, it must be called before app.Run stage
 	SetChangeReceiver(receiver ChangeReceiver) (err error)
+	// SetResetReceiver sets the receiver that is called when the change stream was interrupted
+	// and updates could have been missed, it must be called before app.Run stage
+	SetResetReceiver(receiver ResetReceiver) (err error)
 	// SetDeletionId sets the last deleted log id
 	SetDeletionId(ctx context.Context, lastId string) (err error)
 	// GetDeletionId gets the last deletion log id
@@ -63,6 +77,7 @@ type service struct {
 	running      bool
 
 	changeReceiver ChangeReceiver
+	resetReceiver  ResetReceiver
 	streamCtx      context.Context
 	streamCancel   context.CancelFunc
 	listenerDone   chan struct{}
@@ -97,7 +112,7 @@ func (s *service) Run(ctx context.Context) (err error) {
 		return err
 	}
 	if s.changeReceiver != nil {
-		if err = s.runStreamListener(ctx); err != nil {
+		if err = s.runStreamListener(); err != nil {
 			return err
 		}
 	}
@@ -277,25 +292,39 @@ func (s *service) SetChangeReceiver(receiver ChangeReceiver) (err error) {
 	return
 }
 
+func (s *service) SetResetReceiver(receiver ResetReceiver) (err error) {
+	if s.running {
+		return fmt.Errorf("set receiver must be called before Run")
+	}
+	s.resetReceiver = receiver
+	return
+}
+
 type matchPipeline struct {
 	Match struct {
 		OT string `bson:"operationType"`
 	} `bson:"$match"`
 }
 
-func (s *service) runStreamListener(ctx context.Context) (err error) {
-	var mp matchPipeline
-	mp.Match.OT = "update"
-	stream, err := s.logColl.Watch(ctx, []matchPipeline{mp})
+func (s *service) runStreamListener() (err error) {
+	s.streamCtx, s.streamCancel = context.WithCancel(context.Background())
+	// the change stream must outlive the Run context, so it is opened with the stream context
+	stream, err := s.watchLog(s.streamCtx)
 	if err != nil {
+		s.streamCancel()
 		return
 	}
 	s.listenerDone = make(chan struct{})
-	s.streamCtx, s.streamCancel = context.WithCancel(context.Background())
 	go s.streamListener(stream)
 	return
 }
 
+func (s *service) watchLog(ctx context.Context) (*mongo.ChangeStream, error) {
+	var mp matchPipeline
+	mp.Match.OT = "update"
+	return s.logColl.Watch(ctx, []matchPipeline{mp})
+}
+
 type streamResult struct {
 	DocumentKey struct {
 		Id string `bson:"_id"`
@@ -308,23 +337,103 @@ type streamResult struct {
 }
 
 func (s *service) streamListener(stream *mongo.ChangeStream) {
-	defer close(s.listenerDone)
-	for stream.Next(s.streamCtx) {
-		var res streamResult
-		if err := stream.Decode(&res); err != nil {
-			// mongo driver maintains connections and handles reconnects so that the stream will work as usual in these cases
-			// here we have an unexpected error and should stop any operations to avoid an inconsistent state between db and cache
-			log.Fatal("stream decode error:", zap.Error(err))
+	defer func() {
+		// the listener is the only thing that keeps the cache in sync with the db: if it ever stops while
+		// the service is running, the cache silently serves stale records forever, so fail loudly instead
+		if s.streamCtx.Err() == nil {
+			log.Fatal("change stream listener stopped unexpectedly")
 		}
-		logWithRecords, err := s.injectPayloads(s.streamCtx, consensus.Log{Id: res.DocumentKey.Id, Records: res.UpdateDescription.UpdateFields.Records}, "")
-		if err != nil {
-			log.Error("failed to add payloads to log", zap.Error(err))
-			continue
+		close(s.listenerDone)
+	}()
+	var attempt int
+	for {
+		openedAt := time.Now()
+		for stream.Next(s.streamCtx) {
+			var res streamResult
+			if err := stream.Decode(&res); err != nil {
+				// mongo driver maintains connections and handles reconnects so that the stream will work as usual in these cases
+				// here we have an unexpected error and should stop any operations to avoid an inconsistent state between db and cache
+				log.Fatal("stream decode error:", zap.Error(err))
+			}
+			logWithRecords, err := s.injectPayloads(s.streamCtx, consensus.Log{Id: res.DocumentKey.Id, Records: res.UpdateDescription.UpdateFields.Records}, "")
+			if err != nil {
+				log.Error("failed to add payloads to log", zap.Error(err))
+				continue
+			}
+			s.changeReceiver(res.DocumentKey.Id, logWithRecords.Records)
+		}
+		// the driver retries resumable errors itself, so getting here means the stream is done for good:
+		// either we are closing, or it ended with a non-resumable error, a failed resume or an invalidate event.
+		// Err must be read before Close, Close overwrites it
+		err := stream.Err()
+		_ = stream.Close(s.streamCtx)
+		if s.streamCtx.Err() != nil {
+			return
+		}
+		log.Warn("change stream is closed, reconnecting", zap.Error(err))
+		// a stream that lived long enough is reconnected without a delay, but one that dies right after
+		// being opened must not be reopened in a tight loop
+		if time.Since(openedAt) >= healthyStreamTime {
+			attempt = 0
 		}
-		s.changeReceiver(res.DocumentKey.Id, logWithRecords.Records)
+		// the stream is not resumed from the last token on purpose: the pipeline filters out the invalidate
+		// event, so the token can never move past it and a resumed stream would die on it again and again.
+		// A new stream starts from now, so the changes made in between are never delivered and the cache
+		// has to be resynced with the db
+		if stream = s.reopenStream(&attempt); stream == nil {
+			return
+		}
+		if s.resetReceiver != nil {
+			s.resetReceiver(s.streamCtx)
+		}
+	}
+}
+
+// reopenStream reopens the change stream, retrying until it succeeds. It returns nil if the service is closing.
+func (s *service) reopenStream(attempt *int) (stream *mongo.ChangeStream) {
+	for {
+		if !s.waitBeforeReconnect(*attempt) {
+			return nil
+		}
+		*attempt++
+		var err error
+		if stream, err = s.watchLog(s.streamCtx); err == nil {
+			log.Info("change stream is reconnected")
+			return stream
+		}
+		if s.streamCtx.Err() != nil {
+			return nil
+		}
+		log.Error("can't reopen change stream", zap.Error(err))
+	}
+}
+
+// waitBeforeReconnect returns false if the service is closing
+func (s *service) waitBeforeReconnect(attempt int) bool {
+	waitTime := reconnectWait(attempt)
+	if waitTime == 0 {
+		return true
+	}
+	log.Debug("waiting before the next change stream reconnect", zap.Duration("waitTime", waitTime))
+	select {
+	case <-time.After(waitTime):
+		return true
+	case <-s.streamCtx.Done():
+		return false
 	}
 }
 
+// reconnectWait is a linear backoff: the first attempt is immediate, the following ones are capped
+func reconnectWait(attempt int) time.Duration {
+	if attempt <= 0 {
+		return 0
+	}
+	if wait := time.Duration(attempt) * time.Second; wait < maxReconnectWait {
+		return wait
+	}
+	return maxReconnectWait
+}
+
 func (s *service) SetDeletionId(ctx context.Context, lastId string) (err error) {
 	updateOpts := options.Update().SetUpsert(true)
 	_, err = s.settingsColl.UpdateOne(ctx, bson.D{{"_id", "settings"}}, bson.D{{"$set", bson.D{{"logId", lastId}}}}, updateOpts)
@@ -343,13 +452,14 @@ func (s *service) GetDeletionId(ctx context.Context) (lastId string, err error)
 }
 
 func (s *service) Close(ctx context.Context) (err error) {
-	if s.client != nil {
-		err = s.client.Disconnect(ctx)
-		s.client = nil
-	}
+	// stop the listener before disconnecting, otherwise it would try to reconnect to a closed client
 	if s.listenerDone != nil {
 		s.streamCancel()
 		<-s.listenerDone
 	}
+	if s.client != nil {
+		err = s.client.Disconnect(ctx)
+		s.client = nil
+	}
 	return
 }

db/db_test.go

@@ -204,6 +204,66 @@ func TestService_ChangeReceive(t *testing.T) {
 	})
 }
 
+// TestService_StreamListener_Reconnect: an interrupted change stream must not take the listener down with it
+func TestService_StreamListener_Reconnect(t *testing.T) {
+	var (
+		logs   = make(chan consensus.Log, 10)
+		resets = make(chan struct{}, 10)
+	)
+	fx := newFixtureReset(t, func(logId string, records []consensus.Record) {
+		logs <- consensus.Log{Id: logId, Records: records}
+	}, func(ctx context.Context) {
+		resets <- struct{}{}
+	})
+	defer fx.Finish(t)
+
+	addLogWithRecord := func(logId string) {
+		require.NoError(t, fx.AddLog(ctx, consensus.Log{
+			Id:      logId,
+			Records: []consensus.Record{{Id: logId + "1", Payload: []byte("payload1")}},
+		}))
+		require.NoError(t, fx.AddRecord(ctx, logId, consensus.Record{
+			Id:      logId + "2",
+			PrevId:  logId + "1",
+			Payload: []byte("payload2"),
+		}))
+	}
+	expectChange := func(logId string) {
+		t.Helper()
+		select {
+		case l := <-logs:
+			assert.Equal(t, logId, l.Id)
+		case <-time.After(time.Second * 10):
+			t.Fatalf("no change received for %s", logId)
+		}
+	}
+
+	// the change stream delivers updates
+	addLogWithRecord("logBefore")
+	expectChange("logBefore")
+
+	// dropping the watched collection ends the change stream with an invalidate event
+	require.NoError(t, fx.Service.(*service).logColl.Drop(ctx))
+
+	// the listener must notice it, reopen the stream, and report that updates could have been missed
+	select {
+	case <-resets:
+	case <-time.After(time.Second * 10):
+		t.Fatal("the change stream was not reconnected")
+	}
+
+	// and it must keep delivering updates on the reopened stream
+	addLogWithRecord("logAfter")
+	expectChange("logAfter")
+}
+
+func TestReconnectWait(t *testing.T) {
+	assert.Equal(t, time.Duration(0), reconnectWait(0), "the first reconnect must be immediate")
+	assert.Equal(t, time.Second, reconnectWait(1))
+	assert.Equal(t, time.Second*2, reconnectWait(2))
+	assert.Equal(t, maxReconnectWait, reconnectWait(100), "the wait must be capped")
+}
+
 func TestService_SetDeletionId(t *testing.T) {
 	t.Run("empty", func(t *testing.T) {
 		fx := newFixture(t, nil)
@@ -224,6 +284,10 @@ func TestService_SetDeletionId(t *testing.T) {
 }
 
 func newFixture(t *testing.T, cr ChangeReceiver) *fixture {
+	return newFixtureReset(t, cr, nil)
+}
+
+func newFixtureReset(t *testing.T, cr ChangeReceiver, rr ResetReceiver) *fixture {
 	ctx, cancel := context.WithTimeout(ctx, time.Second)
 	fx := &fixture{
 		Service: New(),
@@ -233,6 +297,7 @@ func newFixture(t *testing.T, cr ChangeReceiver) *fixture {
 	fx.a.Register(&testConfig{})
 	fx.a.Register(fx.Service)
 	require.NoError(t, fx.Service.SetChangeReceiver(cr))
+	require.NoError(t, fx.Service.SetResetReceiver(rr))
 	err := fx.a.Start(ctx)
 	if err != nil {
 		fx.cancel()