cont.

Author: prestonvasquez

internal/integration/cursor_test.go

@@ -394,6 +394,74 @@ func TestCursor_tailableAwaitData(t *testing.T) {
 		assert.LessOrEqual(mt, parseMaxAwaitTime(mt, getMoreStartedEvents[1]), int64(71))
 	})
 
+	// Use a 200ms timeout that caps the lifetime of cursor.Next. The underlying
+	// getMore loop should run at least two times: the first getMore will block
+	// for 30ms on the getMore and then an additional 100ms for the
+	// maxAwaitTimeMS. The second getMore will then use the remaining ~70ms
+	// left on the timeout.
+	clientMtOpts := mtOpts.ClientOptions(options.Client().SetTimeout(200 * time.Millisecond))
+
+	mt.RunOpts("apply remaining client-level timeoutMS if less than maxAwaitTimeMS", clientMtOpts, func(mt *mtest.T) {
+		initCollection(mt, mt.Coll)
+
+		// Create a 30ms failpoint for getMore.
+		mt.SetFailPoint(failpoint.FailPoint{
+			ConfigureFailPoint: "failCommand",
+			Mode: failpoint.Mode{
+				Times: 1,
+			},
+			Data: failpoint.Data{
+				FailCommands:    []string{"getMore"},
+				BlockConnection: true,
+				BlockTimeMS:     30,
+			},
+		})
+
+		// Create a find cursor with a 100ms maxAwaitTimeMS and a tailable awaitData
+		// cursor type.
+		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)
+
+		defer cursor.Close(context.Background())
+
+		// Iterate twice to force a getMore
+		cursor.Next(context.Background())
+
+		// We expect 2 calls to getMore. Since batchSize=1 the first call will
+		mt.ClearEvents()
+		cursor.Next(context.Background())
+
+		m := make(map[string]any)
+
+		err = cursor.Decode(&m)
+		require.NoError(t, err, "expected to decode a document, got error: %v", err)
+
+		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)
+			}
+		}
+
+		// 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))
+	})
+
 	mtOpts.Topologies(mtest.ReplicaSet, mtest.Sharded, mtest.LoadBalanced, mtest.Single)
 
 	mt.RunOpts("apply maxAwaitTimeMS if less than remaining timeout", mtOpts, func(mt *mtest.T) {

mongo/collection.go

@@ -1567,7 +1567,9 @@ func (coll *Collection) find(
 	if err != nil {
 		return nil, wrapErrors(err)
 	}
-	return newCursorWithSession(bc, coll.bsonOpts, coll.registry, sess)
+
+	return newCursorWithSession(bc, coll.bsonOpts, coll.registry, sess,
+		withCursorOptionClientTimeout(coll.client.timeout))
 }
 
 func newFindArgsFromFindOneArgs(args *options.FindOneOptions) *options.FindOptions {

mongo/cursor.go

@@ -37,10 +37,25 @@ type Cursor struct {
 	bsonOpts      *options.BSONOptions
 	registry      *bson.Registry
 	clientSession *session.Client
+	clientTimeout *time.Duration
 
 	err error
 }
 
+type cursorOptions struct {
+	clientTimeout *time.Duration
+}
+
+type cursorOption func(*cursorOptions)
+
+func withCursorOptionClientTimeout(dur *time.Duration) cursorOption {
+	return func(opts *cursorOptions) {
+		if dur != nil {
+			opts.clientTimeout = dur
+		}
+	}
+}
+
 func newCursor(
 	bc batchCursor,
 	bsonOpts *options.BSONOptions,
@@ -54,18 +69,26 @@ func newCursorWithSession(
 	bsonOpts *options.BSONOptions,
 	registry *bson.Registry,
 	clientSession *session.Client,
+	opts ...cursorOption,
 ) (*Cursor, error) {
 	if registry == nil {
 		registry = defaultRegistry
 	}
 	if bc == nil {
 		return nil, errors.New("batch cursor must not be nil")
 	}
+
+	cursorOpts := &cursorOptions{}
+	for _, opt := range opts {
+		opt(cursorOpts)
+	}
+
 	c := &Cursor{
 		bc:            bc,
 		bsonOpts:      bsonOpts,
 		registry:      registry,
 		clientSession: clientSession,
+		clientTimeout: cursorOpts.clientTimeout,
 	}
 	if bc.ID() == 0 {
 		c.closeImplicitSession()
@@ -177,6 +200,15 @@ func (c *Cursor) next(ctx context.Context, nonBlocking bool) bool {
 		ctx = context.Background()
 	}
 
+	// If the context does not have a deadline we defer to a client-level timeout,
+	// if one is set.
+	if _, ok := ctx.Deadline(); !ok && c.clientTimeout != nil {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(context.Background(), *c.clientTimeout)
+
+		defer cancel()
+	}
+
 	// To avoid unnecessary socket timeouts, we attempt to short-circuit tailable
 	// awaitData "getMore" operations by ensuring that the maxAwaitTimeMS is less
 	// than the operation timeout.