cleanup

Author: prestonvasquez

internal/integration/cursor_test.go

@@ -349,7 +349,6 @@ func TestCursor_tailableAwaitData(t *testing.T) {
 		// Create a find cursor with a 100ms maxAwaitTimeMS and a tailable awaitData
 		// cursor type.
 		opts := options.Find().
-			SetBatchSize(1).
 			SetMaxAwaitTime(100 * time.Millisecond).
 			SetCursorType(options.TailableAwait)
 
@@ -367,10 +366,10 @@ func TestCursor_tailableAwaitData(t *testing.T) {
 		defer cancel()
 
 		// Iterate twice to force a getMore
-		cursor.Next(ctx)
+		assert.True(mt, cursor.Next(ctx), "expected Next to return true, got false")
 
 		mt.ClearEvents()
-		cursor.Next(ctx)
+		assert.False(mt, cursor.Next(ctx), "expected Next to return true, got false")
 
 		require.Error(mt, cursor.Err(), "expected error from cursor.Next")
 		assert.ErrorIs(mt, cursor.Err(), context.DeadlineExceeded, "expected context deadline exceeded error")
@@ -385,6 +384,8 @@ func TestCursor_tailableAwaitData(t *testing.T) {
 			}
 		}
 
+		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))
 
@@ -431,14 +432,9 @@ func TestCursor_tailableAwaitData(t *testing.T) {
 		// Iterate twice to force a getMore
 		cursor.Next(context.Background())
 
-		// We expect 2 calls to getMore. Since batchSize=1 the first call will
+		// We expect 2 calls to getMore.
 		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)
+		assert.False(mt, cursor.Next(context.Background()), "expected Next to return false, got true")
 
 		require.Error(mt, cursor.Err(), "expected error from cursor.Next")
 		assert.ErrorIs(mt, cursor.Err(), context.DeadlineExceeded, "expected context deadline exceeded error")
@@ -453,6 +449,8 @@ func TestCursor_tailableAwaitData(t *testing.T) {
 			}
 		}
 
+		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))
 

mongo/client_bulk_write.go

@@ -476,7 +476,11 @@ func (mb *modelBatches) processResponse(ctx context.Context, resp bsoncore.Docum
 		return err
 	}
 	var cursor *Cursor
-	cursor, err = newCursor(bCursor, mb.client.bsonOpts, mb.client.registry)
+	cursor, err = newCursor(bCursor, mb.client.bsonOpts, mb.client.registry,
+
+		// This op doesn't return a cursor to the user, so setting the client
+		// timeout should be a no-op.
+		withCursorOptionClientTimeout(mb.client.timeout))
 	if err != nil {
 		return err
 	}

mongo/collection.go

@@ -1092,7 +1092,8 @@ func aggregate(a aggregateParams, opts ...options.Lister[options.AggregateOption
 	if err != nil {
 		return nil, wrapErrors(err)
 	}
-	cursor, err := newCursorWithSession(bc, a.client.bsonOpts, a.registry, sess)
+	cursor, err := newCursorWithSession(bc, a.client.bsonOpts, a.registry, sess,
+		withCursorOptionClientTimeout(a.client.timeout))
 	return cursor, wrapErrors(err)
 }
 

mongo/cursor.go

@@ -31,27 +31,30 @@ type Cursor struct {
 	// to Next or TryNext. If continued access is required, a copy must be made.
 	Current bson.Raw
 
-	bc            batchCursor
-	batch         *bsoncore.Iterator
-	batchLength   int
-	bsonOpts      *options.BSONOptions
-	registry      *bson.Registry
-	clientSession *session.Client
-	clientTimeout *time.Duration
+	bc               batchCursor
+	batch            *bsoncore.Iterator
+	batchLength      int
+	bsonOpts         *options.BSONOptions
+	registry         *bson.Registry
+	clientSession    *session.Client
+	clientTimeout    time.Duration
+	hasClientTimeout bool
 
 	err error
 }
 
 type cursorOptions struct {
-	clientTimeout *time.Duration
+	clientTimeout    time.Duration
+	hasClientTimeout bool
 }
 
 type cursorOption func(*cursorOptions)
 
 func withCursorOptionClientTimeout(dur *time.Duration) cursorOption {
 	return func(opts *cursorOptions) {
-		if dur != nil {
-			opts.clientTimeout = dur
+		if dur != nil && *dur > 0 {
+			opts.clientTimeout = *dur
+			opts.hasClientTimeout = true
 		}
 	}
 }
@@ -60,8 +63,9 @@ func newCursor(
 	bc batchCursor,
 	bsonOpts *options.BSONOptions,
 	registry *bson.Registry,
+	opts ...cursorOption,
 ) (*Cursor, error) {
-	return newCursorWithSession(bc, bsonOpts, registry, nil)
+	return newCursorWithSession(bc, bsonOpts, registry, nil, opts...)
 }
 
 func newCursorWithSession(
@@ -84,11 +88,12 @@ func newCursorWithSession(
 	}
 
 	c := &Cursor{
-		bc:            bc,
-		bsonOpts:      bsonOpts,
-		registry:      registry,
-		clientSession: clientSession,
-		clientTimeout: cursorOpts.clientTimeout,
+		bc:               bc,
+		bsonOpts:         bsonOpts,
+		registry:         registry,
+		clientSession:    clientSession,
+		clientTimeout:    cursorOpts.clientTimeout,
+		hasClientTimeout: cursorOpts.hasClientTimeout,
 	}
 	if bc.ID() == 0 {
 		c.closeImplicitSession()
@@ -163,11 +168,17 @@ func NewCursorFromDocuments(documents []any, preloadedErr error, registry *bson.
 // ID returns the ID of this cursor, or 0 if the cursor has been closed or exhausted.
 func (c *Cursor) ID() int64 { return c.bc.ID() }
 
-// Next gets the next document for this cursor. It returns true if there were no errors and the cursor has not been
-// exhausted.
+// Next gets the next document for this cursor. It returns true if there were no
+// errors and the cursor has not been exhausted.
+//
+// Next blocks until a document is available or an error occurs. If the context
+// expires, the cursor's error will be set to ctx.Err(). In case of an error,
+// Next will return false.
 //
-// Next blocks until a document is available or an error occurs. If the context expires, the cursor's error will
-// be set to ctx.Err(). In case of an error, Next will return false.
+// If MaxAwaitTime is set, the operation will be bound by the Context's
+// deadline. If the context does not have a deadline, the operation will be
+// bound by the client-level timeout, if one is set. If MaxAwaitTime is greater
+// than the user-provided timeout, Next will return false.
 //
 // If Next returns false, subsequent calls will also return false.
 func (c *Cursor) Next(ctx context.Context) bool {
@@ -202,9 +213,9 @@ func (c *Cursor) next(ctx context.Context, nonBlocking bool) bool {
 
 	// 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 {
+	if _, ok := ctx.Deadline(); !ok && c.hasClientTimeout {
 		var cancel context.CancelFunc
-		ctx, cancel = context.WithTimeout(context.Background(), *c.clientTimeout)
+		ctx, cancel = context.WithTimeout(ctx, c.clientTimeout)
 
 		defer cancel()
 	}

mongo/database.go

@@ -306,7 +306,8 @@ func (db *Database) RunCommandCursor(
 		closeImplicitSession(sess)
 		return nil, wrapErrors(err)
 	}
-	cursor, err := newCursorWithSession(bc, db.bsonOpts, db.registry, sess)
+	cursor, err := newCursorWithSession(bc, db.bsonOpts, db.registry, sess,
+		withCursorOptionClientTimeout(db.client.timeout))
 	return cursor, wrapErrors(err)
 }
 
@@ -511,7 +512,8 @@ func (db *Database) ListCollections(
 		closeImplicitSession(sess)
 		return nil, wrapErrors(err)
 	}
-	cursor, err := newCursorWithSession(bc, db.bsonOpts, db.registry, sess)
+	cursor, err := newCursorWithSession(bc, db.bsonOpts, db.registry, sess,
+		withCursorOptionClientTimeout(db.client.timeout))
 	return cursor, wrapErrors(err)
 }
 

mongo/index_view.go

@@ -131,7 +131,9 @@ func (iv IndexView) List(ctx context.Context, opts ...options.Lister[options.Lis
 		closeImplicitSession(sess)
 		return nil, wrapErrors(err)
 	}
-	cursor, err := newCursorWithSession(bc, iv.coll.bsonOpts, iv.coll.registry, sess)
+	cursor, err := newCursorWithSession(bc, iv.coll.bsonOpts, iv.coll.registry, sess,
+		withCursorOptionClientTimeout(iv.coll.client.timeout))
+
 	return cursor, wrapErrors(err)
 }