Add RC test

Author: prestonvasquez

internal/integration/cursor_test.go

@@ -319,6 +319,179 @@ func parseMaxAwaitTime(mt *mtest.T, evt *event.CommandStartedEvent) int64 {
 	return got
 }
 
+func TestCursor_ClientTimeout_tailableAwaitData(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().CreateClient(false))
+
+	// Helper to set a long block to force client timeout to fire.
+	setGetMoreBlock := func(mt *mtest.T, times, ms int32) {
+		mt.Helper()
+		mt.SetFailPoint(failpoint.FailPoint{
+			ConfigureFailPoint: "failCommand",
+			Mode:               failpoint.Mode{Times: times},
+			Data: failpoint.Data{
+				FailCommands:    []string{"getMore"},
+				BlockConnection: true,
+				BlockTimeMS:     ms,
+			},
+		})
+	}
+
+	cappedOpts := options.CreateCollection().SetCapped(true).
+		SetSizeInBytes(1024 * 64)
+
+	// TODO(SERVER-96344): mongos doesn't honor a failpoint's full blockTimeMS.
+	mtOpts := mtest.NewOptions().
+		ClientOptions(options.Client().SetTimeout(200*time.Millisecond)).
+		CollectionCreateOptions(cappedOpts).
+		Topologies(mtest.ReplicaSet, mtest.LoadBalanced, mtest.Single)
+
+	mt.RunOpts("find", mtOpts, func(mt *mtest.T) {
+		initCollection(mt, mt.Coll)
+		setGetMoreBlock(mt, 1, 30)
+
+		opts := options.Find().
+			SetMaxAwaitTime(100 * time.Millisecond).
+			SetCursorType(options.TailableAwait)
+
+		cursor, err := mt.Coll.Find(context.Background(), bson.D{{"x", 1}}, opts)
+		require.NoError(mt, err, "Find error: %v", err)
+
+		defer func() {
+			err := cursor.Close(context.Background())
+			assert.NoError(mt, err, "Close error: %v", err)
+		}()
+
+		// Iterate twice to force a getMore
+		assert.True(mt, cursor.Next(context.Background()))
+
+		mt.ClearEvents()
+
+		assert.False(mt, cursor.Next(context.Background()))
+
+		require.Error(mt, cursor.Err(), "expected error from cursor.Next")
+		assert.ErrorIs(mt, cursor.Err(), context.DeadlineExceeded, "expected context deadline exceeded error")
+
+		// Collect all started events to find the getMore commands.
+		startedEvents := mt.GetAllStartedEvents()
+
+		var getMoreStartedEvents []*event.CommandStartedEvent
+		for _, evt := range startedEvents {
+			if evt.CommandName == "getMore" {
+				getMoreStartedEvents = append(getMoreStartedEvents, evt)
+			}
+		}
+
+		require.Len(mt, getMoreStartedEvents, 2, "expected 2 getMore commands, got %d", len(getMoreStartedEvents))
+
+		// The first getMore should have a maxTimeMS of <= 100ms.
+		assert.LessOrEqual(mt, parseMaxAwaitTime(mt, getMoreStartedEvents[0]), int64(100))
+
+		// The second getMore should have a maxTimeMS of <=71, indicating that we
+		// are using the time remaining in the context rather than the
+		// maxAwaitTimeMS.
+		assert.LessOrEqual(mt, parseMaxAwaitTime(mt, getMoreStartedEvents[1]), int64(71))
+	})
+
+	mt.RunOpts("aggregate with $changeStream", mtOpts, func(mt *mtest.T) {
+		initCollection(mt, mt.Coll)
+		setGetMoreBlock(mt, 1, 30)
+
+		opts := options.Aggregate().SetMaxAwaitTime(100 * time.Millisecond)
+		pipe := mongo.Pipeline{{{"$changeStream", bson.D{}}}}
+
+		cursor, err := mt.Coll.Aggregate(context.Background(), pipe, opts)
+		require.NoError(mt, err, "Aggregate error: %v", err)
+
+		defer func() {
+			err := cursor.Close(context.Background())
+			assert.NoError(mt, err, "Close error: %v", err)
+		}()
+
+		assert.False(mt, cursor.Next(context.Background()))
+
+		require.Error(mt, cursor.Err(), "expected error from cursor.Next")
+		assert.ErrorIs(mt, cursor.Err(), context.DeadlineExceeded, "expected context deadline exceeded error")
+
+		// Collect all started events to find the getMore commands.
+		startedEvents := mt.GetAllStartedEvents()
+
+		var getMoreStartedEvents []*event.CommandStartedEvent
+		for _, evt := range startedEvents {
+			if evt.CommandName == "getMore" {
+				getMoreStartedEvents = append(getMoreStartedEvents, evt)
+			}
+		}
+
+		require.Len(mt, getMoreStartedEvents, 2, "expected 2 getMore commands, got %d", len(getMoreStartedEvents))
+
+		// The first getMore should have a maxTimeMS of <= 100ms.
+		assert.LessOrEqual(mt, parseMaxAwaitTime(mt, getMoreStartedEvents[0]), int64(100))
+
+		// The second getMore should have a maxTimeMS of <=71, indicating that we
+		// are using the time remaining in the context rather than the
+		// maxAwaitTimeMS.
+		assert.LessOrEqual(mt, parseMaxAwaitTime(mt, getMoreStartedEvents[1]), int64(71))
+	})
+
+	mt.RunOpts("runCommandCursor", mtOpts, func(mt *mtest.T) {
+		initCollection(mt, mt.Coll)
+		mt.SetFailPoint(failpoint.FailPoint{
+			ConfigureFailPoint: "failCommand",
+			Mode:               failpoint.Mode{Times: 1},
+			Data: failpoint.Data{
+				FailCommands:    []string{"find"},
+				BlockConnection: true,
+				BlockTimeMS:     30,
+			},
+		})
+
+		cur, err := mt.DB.RunCommandCursor(context.Background(), bson.D{
+			{"find", mt.Coll.Name()},
+			{"filter", bson.D{{"x", 1}}},
+			{"tailable", true},
+			{"awaitData", true},
+			{"batchSize", int32(1)},
+		})
+
+		require.NoError(mt, err, "RunCommandCursor error: %v", err)
+
+		defer func() {
+			err := cur.Close(context.Background())
+			assert.NoError(mt, err, "Close error: %v", err)
+		}()
+
+		// Iterate twice to force a getMore
+		assert.True(mt, cur.Next(context.Background()))
+
+		mt.ClearEvents()
+
+		assert.False(mt, cur.Next(context.Background()))
+
+		require.Error(mt, cur.Err(), "expected error from cursor.Next")
+		assert.ErrorIs(mt, cur.Err(), context.DeadlineExceeded, "expected context deadline exceeded error")
+
+		// Collect all started events to find the getMore commands.
+		startedEvents := mt.GetAllStartedEvents()
+
+		var getMoreStartedEvents []*event.CommandStartedEvent
+		for _, evt := range startedEvents {
+			if evt.CommandName == "getMore" {
+				getMoreStartedEvents = append(getMoreStartedEvents, evt)
+			}
+		}
+
+		require.Len(mt, getMoreStartedEvents, 2, "expected 2 getMore commands, got %d", len(getMoreStartedEvents))
+
+		//// The first getMore should have a maxTimeMS of <= 100ms.
+		//assert.LessOrEqual(mt, parseMaxAwaitTime(mt, getMoreStartedEvents[0]), int64(100))
+
+		//// The second getMore should have a maxTimeMS of <=71, indicating that we
+		//// are using the time remaining in the context rather than the
+		//// maxAwaitTimeMS.
+		//assert.LessOrEqual(mt, parseMaxAwaitTime(mt, getMoreStartedEvents[1]), int64(71))
+	})
+}
+
 func TestCursor_tailableAwaitData(t *testing.T) {
 	mt := mtest.New(t, mtest.NewOptions().CreateClient(false))
 

mongo/collection.go

@@ -1093,6 +1093,11 @@ func aggregate(a aggregateParams, opts ...options.Lister[options.AggregateOption
 		return nil, wrapErrors(err)
 	}
 	cursor, err := newCursorWithSession(bc, a.client.bsonOpts, a.registry, sess,
+
+		// The only way the server will return a tailable/awaitData cursor for an
+		// aggregate operation is for the first stage in the pipeline to
+		// be $changeStream, this is the only time maxAwaitTimeMS should be applied.
+		// For this reason, we pass the client timeout to the cursor.
 		withCursorOptionClientTimeout(a.client.timeout))
 	return cursor, wrapErrors(err)
 }