GODRIVER-2348 Make CSOT feature-gated behavior the default (#1515)

Co-authored-by: Matt Dale <9760375+matthewdale@users.noreply.github.com>
Co-authored-by: Qingyang Hu <103950869+qingyang-hu@users.noreply.github.com>

internal/csot/csot.go

@@ -11,26 +11,74 @@ import (
 	"time"
 )
 
-type timeoutKey struct{}
+type clientLevel struct{}
 
-// MakeTimeoutContext returns a new context with Client-Side Operation Timeout (CSOT) feature-gated behavior
-// and a Timeout set to the passed in Duration. Setting a Timeout on a single operation is not supported in
-// public API.
-//
-// TODO(GODRIVER-2348) We may be able to remove this function once CSOT feature-gated behavior becomes the
-// TODO default behavior.
-func MakeTimeoutContext(ctx context.Context, to time.Duration) (context.Context, context.CancelFunc) {
-	// Only use the passed in Duration as a timeout on the Context if it
-	// is non-zero.
-	cancelFunc := func() {}
-	if to != 0 {
-		ctx, cancelFunc = context.WithTimeout(ctx, to)
+func isClientLevel(ctx context.Context) bool {
+	val := ctx.Value(clientLevel{})
+	if val == nil {
+		return false
 	}
-	return context.WithValue(ctx, timeoutKey{}, true), cancelFunc
+
+	return val.(bool)
 }
 
+// IsTimeoutContext checks if the provided context has been assigned a deadline
+// or has unlimited retries.
 func IsTimeoutContext(ctx context.Context) bool {
-	return ctx.Value(timeoutKey{}) != nil
+	_, ok := ctx.Deadline()
+
+	return ok || isClientLevel(ctx)
+}
+
+// WithTimeout will set the given timeout on the context, if no deadline has
+// already been set.
+//
+// This function assumes that the timeout field is static, given that the
+// timeout should be sourced from the client. Therefore, once a timeout function
+// parameter has  been applied to the context, it will remain for the lifetime
+// of the context.
+func WithTimeout(parent context.Context, timeout *time.Duration) (context.Context, context.CancelFunc) {
+	cancel := func() {}
+
+	if timeout == nil || IsTimeoutContext(parent) {
+		// In the following conditions, do nothing:
+		//	1. The parent already has a deadline
+		//	2. The parent does not have a deadline, but a client-level timeout has
+		//		 been applied.
+		//	3. The parent does not have a deadline, there is not client-level
+		//     timeout, and the timeout parameter DNE.
+		return parent, cancel
+	}
+
+	// If a client-level timeout has not been applied, then apply it.
+	parent = context.WithValue(parent, clientLevel{}, true)
+
+	dur := *timeout
+
+	if dur == 0 {
+		// If the parent does not have a deadline and the timeout is zero, then
+		// do nothing.
+		return parent, cancel
+	}
+
+	// If the parent does not have a dealine and the timeout is non-zero, then
+	// apply the timeout.
+	return context.WithTimeout(parent, dur)
+}
+
+// WithServerSelectionTimeout creates a context with a timeout that is the
+// minimum of serverSelectionTimeoutMS and context deadline. The usage of
+// non-positive values for serverSelectionTimeoutMS are an anti-pattern and are
+// not considered in this calculation.
+func WithServerSelectionTimeout(
+	parent context.Context,
+	serverSelectionTimeout time.Duration,
+) (context.Context, context.CancelFunc) {
+	if serverSelectionTimeout <= 0 {
+		return parent, func() {}
+	}
+
+	return context.WithTimeout(parent, serverSelectionTimeout)
 }
 
 // ZeroRTTMonitor implements the RTTMonitor interface and is used internally for testing. It returns 0 for all

internal/csot/csot_test.go

@@ -0,0 +1,249 @@
+// Copyright (C) MongoDB, Inc. 2024-present.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License. You may obtain
+// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+
+package csot
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"go.mongodb.org/mongo-driver/internal/assert"
+	"go.mongodb.org/mongo-driver/internal/ptrutil"
+)
+
+func newTestContext(t *testing.T, timeout time.Duration) context.Context {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	t.Cleanup(cancel)
+
+	return ctx
+}
+
+func TestWithServerSelectionTimeout(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name                   string
+		parent                 context.Context
+		serverSelectionTimeout time.Duration
+		wantTimeout            time.Duration
+		wantOk                 bool
+	}{
+		{
+			name:                   "no context deadine and ssto is zero",
+			parent:                 context.Background(),
+			serverSelectionTimeout: 0,
+			wantTimeout:            0,
+			wantOk:                 false,
+		},
+		{
+			name:                   "no context deadline and ssto is positive",
+			parent:                 context.Background(),
+			serverSelectionTimeout: 1,
+			wantTimeout:            1,
+			wantOk:                 true,
+		},
+		{
+			name:                   "no context deadline and ssto is negative",
+			parent:                 context.Background(),
+			serverSelectionTimeout: -1,
+			wantTimeout:            0,
+			wantOk:                 false,
+		},
+		{
+			name:                   "context deadline is zero and ssto is positive",
+			parent:                 newTestContext(t, 0),
+			serverSelectionTimeout: 1,
+			wantTimeout:            1,
+			wantOk:                 true,
+		},
+		{
+			name:                   "context deadline is zero and ssto is negative",
+			parent:                 newTestContext(t, 0),
+			serverSelectionTimeout: -1,
+			wantTimeout:            0,
+			wantOk:                 true,
+		},
+		{
+			name:                   "context deadline is negative and ssto is zero",
+			parent:                 newTestContext(t, -1),
+			serverSelectionTimeout: 0,
+			wantTimeout:            -1,
+			wantOk:                 true,
+		},
+		{
+			name:                   "context deadline is negative and ssto is positive",
+			parent:                 newTestContext(t, -1),
+			serverSelectionTimeout: 1,
+			wantTimeout:            1,
+			wantOk:                 true,
+		},
+		{
+			name:                   "context deadline is negative and ssto is negative",
+			parent:                 newTestContext(t, -1),
+			serverSelectionTimeout: -1,
+			wantTimeout:            -1,
+			wantOk:                 true,
+		},
+		{
+			name:                   "context deadline is positive and ssto is zero",
+			parent:                 newTestContext(t, 1),
+			serverSelectionTimeout: 0,
+			wantTimeout:            1,
+			wantOk:                 true,
+		},
+		{
+			name:                   "context deadline is positive and equal to ssto",
+			parent:                 newTestContext(t, 1),
+			serverSelectionTimeout: 1,
+			wantTimeout:            1,
+			wantOk:                 true,
+		},
+		{
+			name:                   "context deadline is positive lt ssto",
+			parent:                 newTestContext(t, 1),
+			serverSelectionTimeout: 2,
+			wantTimeout:            2,
+			wantOk:                 true,
+		},
+		{
+			name:                   "context deadline is positive gt ssto",
+			parent:                 newTestContext(t, 2),
+			serverSelectionTimeout: 1,
+			wantTimeout:            2,
+			wantOk:                 true,
+		},
+		{
+			name:                   "context deadline is positive and ssto is negative",
+			parent:                 newTestContext(t, -1),
+			serverSelectionTimeout: -1,
+			wantTimeout:            1,
+			wantOk:                 true,
+		},
+	}
+
+	for _, test := range tests {
+		test := test // Capture the range variable
+
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			ctx, cancel := WithServerSelectionTimeout(test.parent, test.serverSelectionTimeout)
+			t.Cleanup(cancel)
+
+			deadline, gotOk := ctx.Deadline()
+			assert.Equal(t, test.wantOk, gotOk)
+
+			if gotOk {
+				delta := time.Until(deadline) - test.wantTimeout
+				tolerance := 10 * time.Millisecond
+
+				assert.True(t, delta > -1*tolerance, "expected delta=%d > %d", delta, -1*tolerance)
+				assert.True(t, delta <= tolerance, "expected delta=%d <= %d", delta, tolerance)
+			}
+		})
+	}
+}
+
+func TestWithTimeout(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name         string
+		parent       context.Context
+		timeout      *time.Duration
+		wantTimeout  time.Duration
+		wantDeadline bool
+		wantValues   []interface{}
+	}{
+		{
+			name:         "deadline set with non-zero timeout",
+			parent:       newTestContext(t, 1),
+			timeout:      ptrutil.Ptr(time.Duration(2)),
+			wantTimeout:  1,
+			wantDeadline: true,
+			wantValues:   []interface{}{},
+		},
+		{
+			name:         "deadline set with zero timeout",
+			parent:       newTestContext(t, 1),
+			timeout:      ptrutil.Ptr(time.Duration(0)),
+			wantTimeout:  1,
+			wantDeadline: true,
+			wantValues:   []interface{}{},
+		},
+		{
+			name:         "deadline set with nil timeout",
+			parent:       newTestContext(t, 1),
+			timeout:      nil,
+			wantTimeout:  1,
+			wantDeadline: true,
+			wantValues:   []interface{}{},
+		},
+		{
+			name:         "deadline unset with non-zero timeout",
+			parent:       context.Background(),
+			timeout:      ptrutil.Ptr(time.Duration(1)),
+			wantTimeout:  1,
+			wantDeadline: true,
+			wantValues:   []interface{}{},
+		},
+		{
+			name:         "deadline unset with zero timeout",
+			parent:       context.Background(),
+			timeout:      ptrutil.Ptr(time.Duration(0)),
+			wantTimeout:  0,
+			wantDeadline: false,
+			wantValues:   []interface{}{clientLevel{}},
+		},
+		{
+			name:         "deadline unset with nil timeout",
+			parent:       context.Background(),
+			timeout:      nil,
+			wantTimeout:  0,
+			wantDeadline: false,
+			wantValues:   []interface{}{},
+		},
+		{
+			// If "clientLevel" has been set, but a new timeout is applied
+			// to the context, then the constructed context should retain the old
+			// timeout. To simplify the code, we assume the first timeout is static.
+			name:         "deadline unset with non-zero timeout at clientLevel",
+			parent:       context.WithValue(context.Background(), clientLevel{}, true),
+			timeout:      ptrutil.Ptr(time.Duration(1)),
+			wantTimeout:  0,
+			wantDeadline: false,
+			wantValues:   []interface{}{},
+		},
+	}
+
+	for _, test := range tests {
+		test := test // Capture the range variable
+
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			ctx, cancel := WithTimeout(test.parent, test.timeout)
+			t.Cleanup(cancel)
+
+			deadline, gotDeadline := ctx.Deadline()
+			assert.Equal(t, test.wantDeadline, gotDeadline)
+
+			if gotDeadline {
+				delta := time.Until(deadline) - test.wantTimeout
+				tolerance := 10 * time.Millisecond
+
+				assert.True(t, delta > -1*tolerance, "expected delta=%d > %d", delta, -1*tolerance)
+				assert.True(t, delta <= tolerance, "expected delta=%d <= %d", delta, tolerance)
+			}
+
+			for _, wantValue := range test.wantValues {
+				assert.NotNil(t, ctx.Value(wantValue), "expected context to have value %v", wantValue)
+			}
+		})
+	}
+
+}

internal/docexamples/examples.go

@@ -1978,7 +1978,6 @@ func WithTransactionExample(ctx context.Context) error {
 
 	// Prereq: Create collections.
 	wcMajority := writeconcern.Majority()
-	wcMajority.WTimeout = 1 * time.Second
 	wcMajorityCollectionOpts := options.Collection().SetWriteConcern(wcMajority)
 	fooColl := client.Database("mydb1").Collection("foo", wcMajorityCollectionOpts)
 	barColl := client.Database("mydb1").Collection("bar", wcMajorityCollectionOpts)
@@ -2559,7 +2558,6 @@ func CausalConsistencyExamples(client *mongo.Client) error {
 
 	rc := readconcern.Majority()
 	wc := writeconcern.Majority()
-	wc.WTimeout = 1000
 	// Use a causally-consistent session to run some operations
 	opts := options.Session().SetDefaultReadConcern(rc).SetDefaultWriteConcern(wc)
 	session1, err := client.StartSession(opts)