Call Next() on the underlying cursor after a resumeable change stream error.

GODRIVER-746

Change-Id: I4b62daffa4443a3af1775f51bdb63487dd9f571a

Author: Divjot Arora

mongo/change_stream.go

@@ -205,42 +205,13 @@ func (cs *changeStream) runCommand(ctx context.Context, replaceOptions bool) err
 	}
 	cs.cursor = cursor
 
-	// can get resume token from initial aggregate command if non-empty batch
-	// operationTime from aggregate saved in the session
 	cursorValue, err := rdr.LookupErr("cursor")
 	if err != nil {
 		return err
 	}
 	cursorDoc := cursorValue.Document()
-
 	cs.ns = command.ParseNamespace(cursorDoc.Lookup("ns").StringValue())
 
-	batchVal := cursorDoc.Lookup("firstBatch")
-	if err != nil {
-		return err
-	}
-
-	batch := batchVal.Array()
-	elements, err := batch.Elements()
-	if err != nil {
-		return err
-	}
-
-	if len(elements) == 0 {
-		return nil // no resume token
-	}
-
-	firstElem, err := batch.IndexErr(0)
-	if err != nil {
-		return err
-	}
-
-	tokenDoc, err := bsonx.ReadDoc(firstElem.Value().Document().Lookup("_id").Document())
-	if err != nil {
-		return err
-	}
-
-	cs.resumeToken = tokenDoc
 	return nil
 }
 
@@ -398,6 +369,34 @@ func newClientChangeStream(ctx context.Context, client *Client, pipeline interfa
 	return cs, nil
 }
 
+func (cs *changeStream) storeResumeToken() error {
+	br, err := cs.cursor.DecodeBytes()
+	if err != nil {
+		return err
+	}
+
+	idVal, err := br.LookupErr("_id")
+	if err != nil {
+		_ = cs.Close(context.Background())
+		return ErrMissingResumeToken
+	}
+
+	var idDoc bson.Raw
+	idDoc, ok := idVal.DocumentOK()
+	if !ok {
+		_ = cs.Close(context.Background())
+		return ErrMissingResumeToken
+	}
+	tokenDoc, err := bsonx.ReadDoc(idDoc)
+	if err != nil {
+		_ = cs.Close(context.Background())
+		return ErrMissingResumeToken
+	}
+
+	cs.resumeToken = tokenDoc
+	return nil
+}
+
 func (cs *changeStream) ID() int64 {
 	if cs.cursor == nil {
 		return 0
@@ -407,38 +406,45 @@ func (cs *changeStream) ID() int64 {
 }
 
 func (cs *changeStream) Next(ctx context.Context) bool {
-	if cs.cursor == nil {
-		return false
-	}
+	// execute in a loop to retry resume-able errors and advance the underlying cursor
+	for {
+		if cs.cursor == nil {
+			return false
+		}
 
-	if cs.cursor.Next(ctx) {
-		return true
-	}
+		if cs.cursor.Next(ctx) {
+			err := cs.storeResumeToken()
+			if err != nil {
+				cs.err = err
+				return false
+			}
 
-	err := cs.cursor.Err()
-	if err == nil {
-		return false
-	}
+			return true
+		}
 
-	switch t := err.(type) {
-	case command.Error:
-		if t.Code == errorInterrupted || t.Code == errorCappedPositionLost || t.Code == errorCursorKilled {
+		err := cs.cursor.Err()
+		if err == nil {
 			return false
 		}
-	}
 
-	killCursors := command.KillCursors{
-		NS:  cs.ns,
-		IDs: []int64{cs.ID()},
-	}
+		switch t := err.(type) {
+		case command.Error:
+			if t.Code == errorInterrupted || t.Code == errorCappedPositionLost || t.Code == errorCursorKilled {
+				return false
+			}
+		}
 
-	_, _ = driver.KillCursors(ctx, killCursors, cs.client.topology, cs.db.writeSelector)
-	cs.err = cs.runCommand(ctx, true)
-	if cs.err != nil {
-		return false
-	}
+		killCursors := command.KillCursors{
+			NS:  cs.ns,
+			IDs: []int64{cs.ID()},
+		}
 
-	return true
+		_, _ = driver.KillCursors(ctx, killCursors, cs.client.topology, cs.db.writeSelector)
+		cs.err = cs.runCommand(ctx, true)
+		if cs.err != nil {
+			return false
+		}
+	}
 }
 
 func (cs *changeStream) Decode(out interface{}) error {
@@ -458,32 +464,11 @@ func (cs *changeStream) DecodeBytes() (bson.Raw, error) {
 	if cs.cursor == nil {
 		return nil, ErrNilCursor
 	}
-
-	br, err := cs.cursor.DecodeBytes()
-	if err != nil {
-		return nil, err
-	}
-
-	idVal, err := br.LookupErr("_id")
-	if err != nil {
-		_ = cs.Close(context.Background())
-		return nil, ErrMissingResumeToken
-	}
-
-	var idDoc bson.Raw
-	idDoc, ok := idVal.DocumentOK()
-	if !ok {
-		_ = cs.Close(context.Background())
-		return nil, ErrMissingResumeToken
-	}
-	tokenDoc, err := bsonx.ReadDoc(idDoc)
-	if err != nil {
-		_ = cs.Close(context.Background())
-		return nil, ErrMissingResumeToken
+	if cs.err != nil {
+		return nil, cs.err
 	}
 
-	cs.resumeToken = tokenDoc
-	return br, nil
+	return cs.cursor.DecodeBytes()
 }
 
 func (cs *changeStream) Err() error {

mongo/change_stream_test.go

@@ -336,13 +336,14 @@ func TestChangeStream_ReplicaSet(t *testing.T) {
 		coll.writeConcern = wcMajority
 		_, err := coll.InsertOne(ctx, doc1)
 		testhelpers.RequireNil(t, err, "error running insertOne: %s", err)
-		if !stream.Next(ctx) {
-			t.Fatal("no change found")
-		}
 
-		_, err = stream.DecodeBytes()
-		if err == nil || err != ErrMissingResumeToken {
-			t.Fatalf("expected ErrMissingResumeToken, got %s", err)
+		// Next should set the change stream error and return false if a document is missing the resume token
+		if stream.Next(ctx) {
+			t.Fatal("Next returned true, expected false")
+		}
+		err = stream.Err()
+		if err != ErrMissingResumeToken {
+			t.Fatalf("expected ErrMissingResumeToken, got %v", err)
 		}
 	})
 
@@ -355,6 +356,8 @@ func TestChangeStream_ReplicaSet(t *testing.T) {
 		startCmd := (<-startedChan).Command
 		startPipeline := startCmd.Lookup("pipeline").Array()
 
+		// make sure resume token is recorded by the change stream because the resume process will hang otherwise
+		ensureResumeToken(t, coll, stream)
 		cs := stream.(*changeStream)
 
 		kc := command.KillCursors{
@@ -674,4 +677,34 @@ func TestChangeStream_ReplicaSet(t *testing.T) {
 			testhelpers.RequireNil(t, stream.Err(), "error while reading stream: %v", err)
 		})
 	})
+
+	t.Run("ResumeErrorCallsNext", func(t *testing.T) {
+		// Test that the underlying cursor is advanced after a resumeable error occurs.
+
+		coll, stream := createCollectionStream(t, "ResumeNextDB", "ResumeNextColl", nil)
+		defer closeCursor(stream)
+		ensureResumeToken(t, coll, stream)
+
+		// kill the stream's underlying cursor to force a resumeable error
+		cs := stream.(*changeStream)
+		kc := command.KillCursors{
+			NS:  cs.ns,
+			IDs: []int64{cs.ID()},
+		}
+
+		_, err := driver.KillCursors(ctx, kc, cs.client.topology, cs.db.writeSelector)
+		testhelpers.RequireNil(t, err, "error running killCursors cmd: %s", err)
+
+		ensureResumeToken(t, coll, stream)
+	})
+}
+
+// ensure that a resume token has been recorded by a change stream
+func ensureResumeToken(t *testing.T, coll *Collection, cs Cursor) {
+	_, err := coll.InsertOne(ctx, bsonx.Doc{{"ensureResumeToken", bsonx.Int32(1)}})
+	testhelpers.RequireNil(t, err, "error inserting doc: %v", err)
+
+	if !cs.Next(ctx) {
+		t.Fatal("Next returned false, expected true")
+	}
 }