DRIVERS-2868 Guard server round trips from timeout

Author: prestonvasquez

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 - 1) / 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/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,22 @@ func createFindCursor(ctx context.Context, operation *operation) (*cursorResult,
 			opts.SetSkip(int64(val.Int32()))
 		case "sort":
 			opts.SetSort(val.Document())
+		case "timeoutMode":
+			return nil, newSkipTestError(fmt.Sprintf("timeoutMode is not supported"))
+		case "cursorType":
+			fmt.Println("cursorType check", val.String(), strings.ToLower(val.String()), val.String() == "tailableAwait")
+			switch strings.ToLower(val.StringValue()) {
+			case "tailable":
+				opts.SetCursorType(options.Tailable)
+			case "tailableawait":
+				fmt.Println("gottem")
+				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

@@ -19,17 +19,24 @@ func executeIterateOnce(ctx context.Context, operation *operation) (*operationRe
 		return nil, err
 	}
 
+	fmt.Println("iterate once")
+
 	// TryNext will attempt to get the next document, potentially issuing a single 'getMore'.
 	if cursor.TryNext(ctx) {
 		// We don't expect the server to return malformed documents, so any errors from Decode here are treated
 		// as fatal.
 		var res bson.Raw
 		if err := cursor.Decode(&res); err != nil {
+			fmt.Println("err:", err)
 			return nil, fmt.Errorf("error decoding cursor result: %w", err)
 		}
 
+		fmt.Println("err:", cursor.Err())
+
 		return newDocumentResult(res, nil), nil
 	}
+
+	fmt.Println("err:", cursor.Err())
 	return newErrorResult(cursor.Err()), nil
 }
 

internal/integration/unified/unified_spec_runner.go

@@ -78,6 +78,10 @@ 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.",
 
 		// DRIVERS-2953: This test requires that the driver sends a "getMore"
 		// with "maxTimeMS" set. However, "getMore" can only include "maxTimeMS"
@@ -178,14 +182,14 @@ func runTestFile(t *testing.T, filepath string, expectValidFail bool, opts ...*O
 				mt.Skip("Skipping CSOT spec test because SKIP_CSOT_TESTS=true")
 			}
 
-			defer func() {
-				// catch panics from looking up elements and fail if it's unexpected
-				if r := recover(); r != nil {
-					if !expectValidFail {
-						mt.Fatal(r)
-					}
-				}
-			}()
+			//defer func() {
+			//	// catch panics from looking up elements and fail if it's unexpected
+			//	if r := recover(); r != nil {
+			//		if !expectValidFail {
+			//			mt.Fatal(r)
+			//		}
+			//	}
+			//}()
 			err := testCase.Run(mt)
 			if expectValidFail {
 				if err != nil {