Extend mongo.batchCursor with MaxAwaitTime for validation

Author: prestonvasquez

mongo/batch_cursor.go

@@ -51,6 +51,11 @@ type batchCursor interface {
 	// SetComment will set a user-configurable comment that can be used to
 	// identify the operation in server logs.
 	SetComment(interface{})
+
+	// MaxAwaitTime returns the maximum amount of time the server will allow
+	// the operations to execute. This is only valid for tailable awaitData
+	// cursors.
+	MaxAwaitTime() *time.Duration
 }
 
 // changeStreamCursor is the interface implemented by batch cursors that also provide the functionality for retrieving

mongo/cursor.go

@@ -16,6 +16,7 @@ import (
 	"time"
 
 	"go.mongodb.org/mongo-driver/v2/bson"
+	"go.mongodb.org/mongo-driver/v2/internal/mongoutil"
 	"go.mongodb.org/mongo-driver/v2/mongo/options"
 	"go.mongodb.org/mongo-driver/v2/x/bsonx/bsoncore"
 	"go.mongodb.org/mongo-driver/v2/x/mongo/driver"
@@ -175,6 +176,32 @@ func (c *Cursor) next(ctx context.Context, nonBlocking bool) bool {
 	if ctx == nil {
 		ctx = context.Background()
 	}
+
+	// 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.
+	//
+	// The specifications assume that drivers iteratively apply the timeout
+	// provided at the constructor level (e.g., (*collection).Find) for tailable
+	// awaitData cursors:
+	//
+	//   If set, drivers MUST apply the timeoutMS option to the initial aggregate
+	//   operation. Drivers MUST also apply the original timeoutMS value to each
+	//   next call on the change stream but MUST NOT use it to derive a maxTimeMS
+	//   field for getMore commands.
+	//
+	// The Go Driver might decide to support the above behavior with DRIVERS-2722.
+	// The principal concern is that it would be unexpected for users to apply an
+	// operation-level timeout via contexts to a constructor and then that timeout
+	// later be applied while working with a resulting cursor. Instead, it is more
+	// idiomatic to apply the timeout to the context passed to Next or TryNext.
+	maxAwaitTime := c.bc.MaxAwaitTime() //
+	if maxAwaitTime != nil && !nonBlocking && !mongoutil.TimeoutWithinContext(ctx, *maxAwaitTime) {
+		c.err = fmt.Errorf("MaxAwaitTime must be less than the operation timeout")
+
+		return false
+	}
+
 	val, err := c.batch.Next()
 	switch {
 	case err == nil:

mongo/cursor_test.go

@@ -27,6 +27,8 @@ type testBatchCursor struct {
 	closed  bool
 }
 
+var _ batchCursor = (*testBatchCursor)(nil)
+
 func newTestBatchCursor(numBatches, batchSize int) *testBatchCursor {
 	batches := make([]*bsoncore.Iterator, 0, numBatches)
 
@@ -99,6 +101,7 @@ func (tbc *testBatchCursor) Close(context.Context) error {
 func (tbc *testBatchCursor) SetBatchSize(int32)            {}
 func (tbc *testBatchCursor) SetComment(interface{})        {}
 func (tbc *testBatchCursor) SetMaxAwaitTime(time.Duration) {}
+func (tbc *testBatchCursor) MaxAwaitTime() *time.Duration  { return nil }
 
 func TestCursor(t *testing.T) {
 	t.Run("TestAll", func(t *testing.T) {

x/mongo/driver/batch_cursor.go

@@ -545,6 +545,12 @@ func (bc *BatchCursor) getOperationDeployment() Deployment {
 	return SingleServerDeployment{bc.server}
 }
 
+// MaxAwaitTime returns the maximum amount of time the server will allow
+// the operations to execute. This is only valid for tailable awaitData cursors.
+func (bc *BatchCursor) MaxAwaitTime() *time.Duration {
+	return bc.maxAwaitTime
+}
+
 // loadBalancedCursorDeployment is used as a Deployment for getMore and killCursors commands when pinning to a
 // connection in load balanced mode. This type also functions as an ErrorProcessor to ensure that SDAM errors are
 // handled for these commands in this mode.