(POC) GODRIVER-3444 Adjust getMore maxTimeMS Calculation for tailable awaitData Cursors

internal/driverutil/operation.go

@@ -6,6 +6,13 @@
 
 package driverutil
 
+import (
+	"context"
+	"fmt"
+	"math"
+	"time"
+)
+
 // Operation Names should be sourced from the command reference documentation:
 // https://www.mongodb.com/docs/manual/reference/command/
 const (
@@ -30,3 +37,35 @@ const (
 	UpdateOp            = "update"            // UpdateOp is the name for updating
 	BulkWriteOp         = "bulkWrite"         // BulkWriteOp is the name for client-level bulk write
 )
+
+func CalculateMaxTimeMS(ctx context.Context, rttMin time.Duration, rttStats string, err error) (int64, error) {
+	deadline, ok := ctx.Deadline()
+	if !ok {
+		return 0, nil
+	}
+
+	remainingTimeout := time.Until(deadline)
+
+	// Always round up to the next millisecond value so we never truncate the calculated
+	// maxTimeMS value (e.g. 400 microseconds evaluates to 1ms, not 0ms).
+	maxTimeMS := int64((remainingTimeout - rttMin) / time.Millisecond)
+	if maxTimeMS <= 0 {
+		return 0, fmt.Errorf(
+			"remaining time %v until context deadline is less than or equal to min network round-trip time %v (%v): %w",
+			remainingTimeout,
+			rttMin,
+			rttStats,
+			err)
+	}
+
+	// The server will return a "BadValue" error if maxTimeMS is greater
+	// than the maximum positive int32 value (about 24.9 days). If the
+	// user specified a timeout value greater than that,  omit maxTimeMS
+	// and let the client-side timeout handle cancelling the op if the
+	// timeout is ever reached.
+	if maxTimeMS > math.MaxInt32 {
+		return 0, nil
+	}
+
+	return maxTimeMS, nil
+}

internal/integration/cursor_test.go

@@ -17,6 +17,7 @@ import (
 	"go.mongodb.org/mongo-driver/v2/internal/assert"
 	"go.mongodb.org/mongo-driver/v2/internal/failpoint"
 	"go.mongodb.org/mongo-driver/v2/internal/integration/mtest"
+	"go.mongodb.org/mongo-driver/v2/internal/require"
 	"go.mongodb.org/mongo-driver/v2/mongo"
 	"go.mongodb.org/mongo-driver/v2/mongo/options"
 )
@@ -303,6 +304,75 @@ func TestCursor(t *testing.T) {
 		batchSize = sizeVal.Int32()
 		assert.Equal(mt, int32(4), batchSize, "expected batchSize 4, got %v", batchSize)
 	})
+
+	tailableAwaitDataCursorOpts := mtest.NewOptions().MinServerVersion("4.4").
+		Topologies(mtest.ReplicaSet, mtest.Sharded, mtest.LoadBalanced, mtest.Single)
+
+	mt.RunOpts("tailable awaitData cursor", tailableAwaitDataCursorOpts, func(mt *mtest.T) {
+		mt.Run("apply remaining timeoutMS if less than maxAwaitTimeMS", func(mt *mtest.T) {
+			initCollection(mt, mt.Coll)
+			mt.ClearEvents()
+
+			// Create a find cursor
+			opts := options.Find().SetBatchSize(1).SetMaxAwaitTime(100 * time.Millisecond)
+
+			cursor, err := mt.Coll.Find(context.Background(), bson.D{}, opts)
+			require.NoError(mt, err)
+
+			_ = mt.GetStartedEvent() // Empty find from started list.
+
+			defer cursor.Close(context.Background())
+
+			ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+			defer cancel()
+
+			// Iterate twice to force a getMore
+			cursor.Next(ctx)
+			cursor.Next(ctx)
+
+			cmd := mt.GetStartedEvent().Command
+
+			maxTimeMSRaw, err := cmd.LookupErr("maxTimeMS")
+			require.NoError(mt, err)
+
+			got, ok := maxTimeMSRaw.AsInt64OK()
+			require.True(mt, ok)
+
+			assert.LessOrEqual(mt, got, int64(50))
+		})
+
+		mt.RunOpts("apply maxAwaitTimeMS if less than remaining timeout", tailableAwaitDataCursorOpts, func(mt *mtest.T) {
+			initCollection(mt, mt.Coll)
+			mt.ClearEvents()
+
+			// Create a find cursor
+			opts := options.Find().SetBatchSize(1).SetMaxAwaitTime(50 * time.Millisecond)
+
+			cursor, err := mt.Coll.Find(context.Background(), bson.D{}, opts)
+			require.NoError(mt, err)
+
+			_ = mt.GetStartedEvent() // Empty find from started list.
+
+			defer cursor.Close(context.Background())
+
+			ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+			defer cancel()
+
+			// Iterate twice to force a getMore
+			cursor.Next(ctx)
+			cursor.Next(ctx)
+
+			cmd := mt.GetStartedEvent().Command
+
+			maxTimeMSRaw, err := cmd.LookupErr("maxTimeMS")
+			require.NoError(mt, err)
+
+			got, ok := maxTimeMSRaw.AsInt64OK()
+			require.True(mt, ok)
+
+			assert.LessOrEqual(mt, got, int64(50))
+		})
+	})
 }
 
 type tryNextCursor interface {

internal/integration/unified/collection_operation_execution.go

@@ -10,6 +10,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"strings"
 	"time"
 
 	"go.mongodb.org/mongo-driver/v2/bson"
@@ -1485,6 +1486,20 @@ func createFindCursor(ctx context.Context, operation *operation) (*cursorResult,
 			opts.SetSkip(int64(val.Int32()))
 		case "sort":
 			opts.SetSort(val.Document())
+		case "timeoutMode":
+			return nil, newSkipTestError("timeoutMode is not supported")
+		case "cursorType":
+			switch strings.ToLower(val.StringValue()) {
+			case "tailable":
+				opts.SetCursorType(options.Tailable)
+			case "tailableawait":
+				opts.SetCursorType(options.TailableAwait)
+			case "nontailable":
+				opts.SetCursorType(options.NonTailable)
+			}
+		case "maxAwaitTimeMS":
+			maxAwaitTimeMS := time.Duration(val.Int32()) * time.Millisecond
+			opts.SetMaxAwaitTime(maxAwaitTimeMS)
 		default:
 			return nil, fmt.Errorf("unrecognized find option %q", key)
 		}

internal/integration/unified/cursor_operation_execution.go

@@ -30,6 +30,7 @@ func executeIterateOnce(ctx context.Context, operation *operation) (*operationRe
 
 		return newDocumentResult(res, nil), nil
 	}
+
 	return newErrorResult(cursor.Err()), nil
 }
 

internal/integration/unified/unified_spec_runner.go

@@ -78,6 +78,24 @@ var (
 		"operation is retried multiple times for non-zero timeoutMS - aggregate on collection": "maxTimeMS is disabled on find and aggregate. See DRIVERS-2722.",
 		"operation is retried multiple times for non-zero timeoutMS - aggregate on database":   "maxTimeMS is disabled on find and aggregate. See DRIVERS-2722.",
 		"timeoutMS applied to find command":                                                    "maxTimeMS is disabled on find and aggregate. See DRIVERS-2722.",
+		"timeoutMS applied to find":                                                            "maxTimeMS is disabled on find and aggregate. See DRIVERS-2722.",
+		"timeoutMS is refreshed for getMore if maxAwaitTimeMS is not set":                      "maxTimeMS is disabled on find and aggregate. See DRIVERS-2722.",
+		"timeoutMS is refreshed for getMore if maxAwaitTimeMS is set":                          "maxTimeMS is disabled on find and aggregate. See DRIVERS-2722.",
+		"timeoutMS is refreshed for getMore - failure":                                         "maxTimeMS is disabled on find and aggregate. See DRIVERS-2722.",
+
+		// GODRIVER-3473: the implementation of DRIVERS-2868 makes it clear that the
+		// Go Driver does not correctly implement the following validation for
+		// tailable awaitData cursors:
+		//
+		//     Drivers MUST error if this option is set, timeoutMS is set to a
+		//     non-zero value, and maxAwaitTimeMS is greater than or equal to
+		//     timeoutMS.
+		//
+		// Once GODRIVER-3473 is completed, we can continue running these tests.
+		"error if maxAwaitTimeMS is equal to timeoutMS":         "Go Driver does not implement this behavior. See GODRIVER-3473",
+		"error if maxAwaitTimeMS is greater than timeoutMS":     "Go Driver does not implement this behavior. See GODRIVER-3473",
+		"apply remaining timeoutMS if less than maxAwaitTimeMS": "Go Driver does not implement this behavior. See GODRIVER-3473",
+		"apply maxAwaitTimeMS if less than remaining timeout":   "Go Driver does not implement this behavior. See GODRIVER-3473",
 
 		// DRIVERS-2953: This test requires that the driver sends a "getMore"
 		// with "maxTimeMS" set. However, "getMore" can only include "maxTimeMS"