add maxAdaptiveRetries and enableOverloadRetargeting

Author: qingyang-hu

internal/integration/backpressure_prose_test.go

@@ -40,7 +40,16 @@ func TestBackpressureProse(t *testing.T) {
 		mt.ResetClient(options.Client())
 
 		transWithJitter := func(t *mtest.T, ratio float64) time.Duration {
-			defer randutil.SetJitterForTesting(ratio)()
+			defer randutil.SetJitterForTesting(func(n int64) int64 {
+				val := int64(ratio * float64(n))
+				if val < 0 {
+					return 0
+				}
+				if val > n {
+					return n
+				}
+				return val
+			})()
 
 			startTime := time.Now()
 			_, err := t.Coll.InsertOne(context.Background(), bson.D{{"a", 1}})
@@ -49,9 +58,9 @@ func TestBackpressureProse(t *testing.T) {
 		}
 		noBackoffTime := transWithJitter(mt, 0)
 		withBackoffTime := transWithJitter(mt, 1)
-		assert.GreaterOrEqualf(
+		assert.InDelta(
 			mt,
-			withBackoffTime, noBackoffTime+2_100*time.Millisecond,
+			withBackoffTime, noBackoffTime+300*time.Millisecond, float64(300*time.Millisecond),
 			"with backoff time: %v, no backoff time: %v", withBackoffTime, noBackoffTime,
 		)
 	})
@@ -81,6 +90,34 @@ func TestBackpressureProse(t *testing.T) {
 		require.Truef(mt, errors.As(err, &cmdErr), "expected a CommandError, got %T: %v", err, err)
 		assert.True(mt, cmdErr.HasErrorLabel("RetryableError"), `expected error has "RetryableError" label`)
 		assert.True(mt, cmdErr.HasErrorLabel("SystemOverloadedError"), `expected error has "SystemOverloadedError" label`)
-		assert.Equalf(mt, 6, opsCnt, "expected 6 attempts (1 original + 5 retries), got %d", opsCnt)
+		assert.Equalf(mt, 3, opsCnt, "expected 3 attempts (1 original + 2 retries), got %d", opsCnt)
+	})
+	mt.Run("4. Overload Errors are Retried a Maximum of maxAdaptiveRetries times when configured", func(mt *mtest.T) {
+		mt.SetFailPoint(failpoint.FailPoint{
+			ConfigureFailPoint: "failCommand",
+			Mode:               failpoint.ModeAlwaysOn,
+			Data: failpoint.Data{
+				FailCommands: []string{"find"},
+				ErrorCode:    462,
+				ErrorLabels:  &[]string{"SystemOverloadedError", "RetryableError"},
+			},
+		})
+
+		var opsCnt int
+		monitor := &event.CommandMonitor{
+			Started: func(_ context.Context, e *event.CommandStartedEvent) {
+				if e.CommandName == "find" {
+					opsCnt++
+				}
+			},
+		}
+		mt.ResetClient(options.Client().SetMonitor(monitor).SetMaxAdaptiveRetries(1))
+
+		_, err := mt.Coll.Find(context.Background(), bson.D{})
+		var cmdErr mongo.CommandError
+		require.Truef(mt, errors.As(err, &cmdErr), "expected a CommandError, got %T: %v", err, err)
+		assert.True(mt, cmdErr.HasErrorLabel("RetryableError"), `expected error has "RetryableError" label`)
+		assert.True(mt, cmdErr.HasErrorLabel("SystemOverloadedError"), `expected error has "SystemOverloadedError" label`)
+		assert.Equalf(mt, 2, opsCnt, "expected 2 attempts (1 original + 1 retry), got %d", opsCnt)
 	})
 }

internal/integration/retryable_reads_prose_test.go

@@ -8,6 +8,7 @@ package integration
 
 import (
 	"context"
+	"errors"
 	"sync"
 	"testing"
 	"time"
@@ -19,7 +20,9 @@ import (
 	"go.mongodb.org/mongo-driver/v2/internal/failpoint"
 	"go.mongodb.org/mongo-driver/v2/internal/integration/mtest"
 	"go.mongodb.org/mongo-driver/v2/internal/mongoutil"
+	"go.mongodb.org/mongo-driver/v2/internal/randutil"
 	"go.mongodb.org/mongo-driver/v2/internal/require"
+	"go.mongodb.org/mongo-driver/v2/mongo"
 	"go.mongodb.org/mongo-driver/v2/mongo/options"
 	"go.mongodb.org/mongo-driver/v2/mongo/readpref"
 )
@@ -206,24 +209,34 @@ func TestRetryableReadsProse(t *testing.T) {
 	mtOpts = mtest.NewOptions().Topologies(mtest.ReplicaSet).MinServerVersion("4.4").CreateClient(false)
 	mt.RunOpts("retrying reads in a replica set", mtOpts, func(mt *mtest.T) {
 		tests := []struct {
-			name              string
-			errLabels         []string
-			isServerIdentical bool
+			name                      string
+			errLabels                 []string
+			enableOverloadRetargeting bool
+			isServerIdentical         bool
 		}{
 			{
-				name:              "overload errors",
-				errLabels:         []string{"RetryableError", "SystemOverloadedError"},
-				isServerIdentical: false,
+				name:                      "overload errors retried on a different replicaset server",
+				errLabels:                 []string{"RetryableError", "SystemOverloadedError"},
+				enableOverloadRetargeting: true,
+				isServerIdentical:         false,
 			},
 			{
-				name:              "non-overload errors",
+				name:              "non-overload errors retried on the same replicaset server",
 				errLabels:         []string{"RetryableError"},
 				isServerIdentical: true,
 			},
+			{
+				name:              "overload errors retried on the same replicaset server",
+				errLabels:         []string{"RetryableError", "SystemOverloadedError"},
+				isServerIdentical: true,
+			},
 		}
 
-		clientOpts := options.Client().SetRetryReads(true).SetReadPreference(readpref.PrimaryPreferred())
 		for _, tc := range tests {
+			clientOpts := options.Client().SetRetryReads(true).SetReadPreference(readpref.PrimaryPreferred())
+			if tc.enableOverloadRetargeting {
+				clientOpts = clientOpts.SetEnableOverloadRetargeting(true)
+			}
 			mt.RunOpts(tc.name, mtest.NewOptions().ClientOptions(clientOpts), func(mt *mtest.T) {
 				failPoint := failpoint.FailPoint{
 					ConfigureFailPoint: "failCommand",
@@ -254,4 +267,108 @@ func TestRetryableReadsProse(t *testing.T) {
 			})
 		}
 	})
+
+	errorCodesContains := func(err error, code int) bool {
+		for _, ec := range mongo.ErrorCodes(err) {
+			if ec == code {
+				return true
+			}
+		}
+		return false
+	}
+
+	mtOpts = mtest.NewOptions().MinServerVersion("4.4").ClientType(mtest.Pinned).AllowFailPointsOnSharded()
+	mt.RunOpts("set the maximum number of retries for all retryable read errors", mtOpts, func(mt *mtest.T) {
+		mt.SetFailPoint(failpoint.FailPoint{
+			ConfigureFailPoint: "failCommand",
+			Mode: failpoint.Mode{
+				Times: 1,
+			},
+			Data: failpoint.Data{
+				FailCommands: []string{"find"},
+				ErrorLabels:  &[]string{"RetryableError", "SystemOverloadedError"},
+				ErrorCode:    91,
+			},
+		})
+
+		var opsCnt int
+		monitor := &event.CommandMonitor{
+			Started: func(_ context.Context, e *event.CommandStartedEvent) {
+				if e.CommandName == "find" {
+					opsCnt++
+				}
+			},
+			Failed: func(_ context.Context, event *event.CommandFailedEvent) {
+				if event.CommandName != "find" {
+					return
+				}
+				if errorCodesContains(event.Failure, 91) {
+					mt.SetFailPoint(failpoint.FailPoint{
+						ConfigureFailPoint: "failCommand",
+						Mode:               failpoint.ModeAlwaysOn,
+						Data: failpoint.Data{
+							FailCommands: []string{"find"},
+							ErrorLabels:  &[]string{"RetryableError"},
+							ErrorCode:    91,
+						},
+					})
+				}
+			},
+		}
+		mt.ResetClient(options.Client().SetRetryReads(true).SetMonitor(monitor))
+		_, err := mt.Coll.Find(context.Background(), bson.D{})
+		var cmdErr mongo.CommandError
+		require.Truef(mt, errors.As(err, &cmdErr), "expected a CommandError, got %T: %v", err, err)
+		assert.Equalf(mt, 3, opsCnt, "expected 3 attempts (1 original + 2 retries), got %d", opsCnt)
+	})
+
+	mt.RunOpts("do not apply backoff to non-overload errors", mtOpts, func(mt *mtest.T) {
+		mt.SetFailPoint(failpoint.FailPoint{
+			ConfigureFailPoint: "failCommand",
+			Mode: failpoint.Mode{
+				Times: 1,
+			},
+			Data: failpoint.Data{
+				FailCommands: []string{"find"},
+				ErrorLabels:  &[]string{"RetryableError", "SystemOverloadedError"},
+				ErrorCode:    91,
+			},
+		})
+
+		var ops []bool
+		monitor := &event.CommandMonitor{
+			Started: func(_ context.Context, e *event.CommandStartedEvent) {
+				if e.CommandName == "find" {
+					ops = append(ops, false)
+				}
+			},
+			Failed: func(_ context.Context, event *event.CommandFailedEvent) {
+				if event.CommandName != "find" {
+					return
+				}
+				if errorCodesContains(event.Failure, 91) {
+					mt.SetFailPoint(failpoint.FailPoint{
+						ConfigureFailPoint: "failCommand",
+						Mode:               failpoint.ModeAlwaysOn,
+						Data: failpoint.Data{
+							FailCommands: []string{"find"},
+							ErrorLabels:  &[]string{"RetryableError"},
+							ErrorCode:    91,
+						},
+					})
+				}
+			},
+		}
+
+		defer randutil.SetJitterForTesting(func(int64) int64 {
+			ops[len(ops)-1] = true
+			return 0
+		})()
+
+		mt.ResetClient(options.Client().SetRetryReads(true).SetMonitor(monitor))
+		_, err := mt.Coll.Find(context.Background(), bson.D{})
+		var cmdErr mongo.CommandError
+		require.Truef(mt, errors.As(err, &cmdErr), "expected a CommandError, got %T: %v", err, err)
+		assert.Equal(t, []bool{true, false, false}, ops, "expected backoff to be applied on the first attempt only, got %v", ops)
+	})
 }

internal/integration/retryable_writes_prose_test.go

@@ -23,6 +23,7 @@ import (
 	"go.mongodb.org/mongo-driver/v2/internal/failpoint"
 	"go.mongodb.org/mongo-driver/v2/internal/integration/mtest"
 	"go.mongodb.org/mongo-driver/v2/internal/mongoutil"
+	"go.mongodb.org/mongo-driver/v2/internal/randutil"
 	"go.mongodb.org/mongo-driver/v2/internal/require"
 	"go.mongodb.org/mongo-driver/v2/mongo"
 	"go.mongodb.org/mongo-driver/v2/mongo/options"
@@ -511,6 +512,102 @@ func TestErrorPropagationAfterEncounteringMultipleErrors(t *testing.T) {
 		require.True(mt, errors.As(err, &labeledError), "expected error to be a labeled error")
 		require.NotContains(mt, labeledError.Labels, "NoWritesPerformed", "expected error labels to not contain NoWritesPerformed")
 	})
+
+	mt.Run("Case 4: Test that drivers set the maximum number of retries for all retryable write errors when an overload error is encountered", func(mt *mtest.T) {
+		mt.SetFailPoint(failpoint.FailPoint{
+			ConfigureFailPoint: "failCommand",
+			Mode: failpoint.Mode{
+				Times: 1,
+			},
+			Data: failpoint.Data{
+				FailCommands: []string{"insert"},
+				ErrorLabels:  &[]string{"RetryableError", "SystemOverloadedError"},
+				ErrorCode:    91,
+			},
+		})
+
+		var opsCnt int
+		monitor := &event.CommandMonitor{
+			Started: func(_ context.Context, e *event.CommandStartedEvent) {
+				if e.CommandName == "insert" {
+					opsCnt++
+				}
+			},
+			Failed: func(_ context.Context, event *event.CommandFailedEvent) {
+				if event.CommandName != "insert" {
+					return
+				}
+				if errorCodesContains(event.Failure, 91) {
+					mt.SetFailPoint(failpoint.FailPoint{
+						ConfigureFailPoint: "failCommand",
+						Mode:               failpoint.ModeAlwaysOn,
+						Data: failpoint.Data{
+							FailCommands: []string{"insert"},
+							ErrorLabels:  &[]string{"RetryableError", "RetryableWriteError"},
+							ErrorCode:    91,
+						},
+					})
+				}
+			},
+		}
+		mt.ResetClient(options.Client().SetRetryWrites(true).SetMonitor(monitor))
+		_, err := mt.Coll.InsertOne(context.Background(), bson.D{})
+		require.True(mt, errorCodesContains(err, 91), "Expect the error code of the server error to be 91")
+		var labeledError driver.Error
+		require.True(mt, errors.As(err, &labeledError), "expected error to be a labeled error")
+		assert.Equalf(mt, 3, opsCnt, "expected 3 attempts (1 original + 2 retries), got %d", opsCnt)
+	})
+
+	mt.Run("Case 5: Test that drivers do not apply backoff to non-overload errors", func(mt *mtest.T) {
+		mt.SetFailPoint(failpoint.FailPoint{
+			ConfigureFailPoint: "failCommand",
+			Mode: failpoint.Mode{
+				Times: 1,
+			},
+			Data: failpoint.Data{
+				FailCommands: []string{"insert"},
+				ErrorLabels:  &[]string{"RetryableError", "SystemOverloadedError"},
+				ErrorCode:    91,
+			},
+		})
+
+		var ops []bool
+		monitor := &event.CommandMonitor{
+			Started: func(_ context.Context, e *event.CommandStartedEvent) {
+				if e.CommandName == "insert" {
+					ops = append(ops, false)
+				}
+			},
+			Failed: func(_ context.Context, event *event.CommandFailedEvent) {
+				if event.CommandName != "insert" {
+					return
+				}
+				if errorCodesContains(event.Failure, 91) {
+					mt.SetFailPoint(failpoint.FailPoint{
+						ConfigureFailPoint: "failCommand",
+						Mode:               failpoint.ModeAlwaysOn,
+						Data: failpoint.Data{
+							FailCommands: []string{"insert"},
+							ErrorLabels:  &[]string{"RetryableError", "RetryableWriteError"},
+							ErrorCode:    91,
+						},
+					})
+				}
+			},
+		}
+
+		defer randutil.SetJitterForTesting(func(int64) int64 {
+			ops[len(ops)-1] = true
+			return 0
+		})()
+
+		mt.ResetClient(options.Client().SetRetryWrites(true).SetMonitor(monitor))
+		_, err := mt.Coll.InsertOne(context.Background(), bson.D{})
+		require.True(mt, errorCodesContains(err, 91), "Expect the error code of the server error to be 91")
+		var labeledError driver.Error
+		require.True(mt, errors.As(err, &labeledError), "expected error to be a labeled error")
+		assert.Equal(t, []bool{true, false, false}, ops, "expected backoff to be applied on the first attempt only, got %v", ops)
+	})
 }
 
 type crudOperation struct {

internal/randutil/jitter.go

@@ -6,31 +6,26 @@
 
 package randutil
 
-var presetRatio *float64
+import "sync"
 
 var globalRand = NewLockedRand()
 
 // JitterInt63n returns, as an int64, a non-negative pseudo-random number in
 // the half-open interval [0,n). It panics if n <= 0.
 //
-// If a static jitter ratio is set by calling SetJitterForTesting, JitterInt63n
-// returns int64(n*ratio) in [0,n].
-func JitterInt63n(n int64) int64 {
-	if presetRatio == nil {
-		return globalRand.Int63n(n)
-	}
-	val := int64(*presetRatio * float64(n))
-	if val < 0 {
-		return 0
-	}
-	if val > n {
-		return n
-	}
-	return int64(val)
-}
+// If a test jitter function is set by calling SetJitterForTesting, JitterInt63n
+// returns the value from the custom function.
+var JitterInt63n func(int64) int64 = globalRand.Int63n
 
-// SetJitterForTesting sets a preset jitter ratio for testing and returns a restore function.
-func SetJitterForTesting(ratio float64) func() {
-	presetRatio = &ratio
-	return func() { presetRatio = nil }
+var testLock sync.Mutex
+
+// SetJitterForTesting sets a custom jitter function for testing and returns a restore function.
+func SetJitterForTesting(f func(int64) int64) func() {
+	testLock.Lock()
+
+	JitterInt63n = f
+	return func() {
+		JitterInt63n = globalRand.Int63n
+		testLock.Unlock()
+	}
 }