GODRIVER-1538 Add support for server-side hedged reads (mongodb#363)

mongo/options/clientoptions_test.go

@@ -200,7 +200,7 @@ func TestClientOptions(t *testing.T) {
 				"ReadPreference Primary With Options",
 				"mongodb://localhost/?readPreference=Primary&maxStaleness=200",
 				&ClientOptions{
-					err:   errors.New("can not specify tags or max staleness on primary"),
+					err:   errors.New("can not specify tags, max staleness, or hedge with mode primary"),
 					Hosts: []string{"localhost"},
 				},
 			},

mongo/readpref/options.go

@@ -58,3 +58,14 @@ func WithTagSets(tagSets ...tag.Set) Option {
 		return nil
 	}
 }
+
+// WithHedgeEnabled specifies whether or not hedged reads should be enabled in the server. This feature requires MongoDB
+// server version 4.4 or higher. For more information about hedged reads, see
+// https://docs.mongodb.com/master/core/sharded-cluster-query-router/#mongos-hedged-reads. If not specified, the default
+// is to not send a value to the server, which will result in the server defaults being used.
+func WithHedgeEnabled(hedgeEnabled bool) Option {
+	return func(rp *ReadPref) error {
+		rp.hedgeEnabled = &hedgeEnabled
+		return nil
+	}
+}

mongo/readpref/readpref.go

@@ -14,7 +14,7 @@ import (
 )
 
 var (
-	errInvalidReadPreference = errors.New("can not specify tags or max staleness on primary")
+	errInvalidReadPreference = errors.New("can not specify tags, max staleness, or hedge with mode primary")
 )
 
 var primary = ReadPref{mode: PrimaryMode}
@@ -78,6 +78,7 @@ type ReadPref struct {
 	maxStalenessSet bool
 	mode            Mode
 	tagSets         []tag.Set
+	hedgeEnabled    *bool
 }
 
 // MaxStaleness is the maximum amount of time to allow
@@ -97,3 +98,9 @@ func (r *ReadPref) Mode() Mode {
 func (r *ReadPref) TagSets() []tag.Set {
 	return r.tagSets
 }
+
+// HedgeEnabled returns whether or not hedged reads are enabled for this read preference. If this option was not
+// specified during read preference construction, nil is returned.
+func (r *ReadPref) HedgeEnabled() *bool {
+	return r.hedgeEnabled
+}

mongo/readpref/readpref_test.go

@@ -4,14 +4,14 @@
 // 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 readpref_test
+package readpref
 
 import (
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/require"
-	. "go.mongodb.org/mongo-driver/mongo/readpref"
+	"go.mongodb.org/mongo-driver/internal/testutil/assert"
 	"go.mongodb.org/mongo-driver/tag"
 )
 
@@ -120,3 +120,17 @@ func TestNearest_with_options(t *testing.T) {
 	require.Equal(time.Duration(10), ms)
 	require.Equal([]tag.Set{{tag.Tag{Name: "a", Value: "1"}, tag.Tag{Name: "b", Value: "2"}}}, subject.TagSets())
 }
+
+func TestHedge(t *testing.T) {
+	t.Run("hedge specified with primary mode errors", func(t *testing.T) {
+		_, err := New(PrimaryMode, WithHedgeEnabled(true))
+		assert.Equal(t, errInvalidReadPreference, err, "expected error %v, got %v", errInvalidReadPreference, err)
+	})
+	t.Run("valid hedge document and mode succeeds", func(t *testing.T) {
+		rp, err := New(SecondaryMode, WithHedgeEnabled(true))
+		assert.Nil(t, err, "expected no error, got %v", err)
+		enabled := rp.HedgeEnabled()
+		assert.NotNil(t, enabled, "expected HedgeEnabled to return a non-nil value, got nil")
+		assert.True(t, *enabled, "expected HedgeEnabled to return true, got false")
+	})
+}

x/mongo/driver/operation.go

@@ -1061,7 +1061,7 @@ func (op Operation) createReadPref(serverKind description.ServerKind, topologyKi
 		doc = bsoncore.AppendStringElement(doc, "mode", "primaryPreferred")
 	case readpref.SecondaryPreferredMode:
 		_, ok := rp.MaxStaleness()
-		if serverKind == description.Mongos && isOpQuery && !ok && len(rp.TagSets()) == 0 {
+		if serverKind == description.Mongos && isOpQuery && !ok && len(rp.TagSets()) == 0 && rp.HedgeEnabled() == nil {
 			return nil, nil
 		}
 		doc = bsoncore.AppendStringElement(doc, "mode", "secondaryPreferred")
@@ -1096,6 +1096,16 @@ func (op Operation) createReadPref(serverKind description.ServerKind, topologyKi
 		doc = bsoncore.AppendInt32Element(doc, "maxStalenessSeconds", int32(d.Seconds()))
 	}
 
+	if hedgeEnabled := rp.HedgeEnabled(); hedgeEnabled != nil {
+		var hedgeIdx int32
+		hedgeIdx, doc = bsoncore.AppendDocumentElementStart(doc, "hedge")
+		doc = bsoncore.AppendBooleanElement(doc, "enabled", *hedgeEnabled)
+		doc, err = bsoncore.AppendDocumentEnd(doc, hedgeIdx)
+		if err != nil {
+			return nil, fmt.Errorf("error creating hedge document: %v", err)
+		}
+	}
+
 	doc, _ = bsoncore.AppendDocumentEnd(doc, idx)
 	return doc, nil
 }

x/mongo/driver/operation_test.go

@@ -346,6 +346,28 @@ func TestOperation(t *testing.T) {
 			bsoncore.AppendStringElement(nil, "mode", "secondaryPreferred"),
 			bsoncore.AppendInt32Element(nil, "maxStalenessSeconds", 25),
 		)
+		// Hedged read preference: {mode: "secondaryPreferred", hedge: {enabled: true}}
+		rpWithHedge := bsoncore.BuildDocumentFromElements(nil,
+			bsoncore.AppendStringElement(nil, "mode", "secondaryPreferred"),
+			bsoncore.AppendDocumentElement(nil, "hedge", bsoncore.BuildDocumentFromElements(nil,
+				bsoncore.AppendBooleanElement(nil, "enabled", true),
+			)),
+		)
+		rpWithAllOptions := bsoncore.BuildDocumentFromElements(nil,
+			bsoncore.AppendStringElement(nil, "mode", "secondaryPreferred"),
+			bsoncore.BuildArrayElement(nil, "tags",
+				bsoncore.Value{Type: bsontype.EmbeddedDocument,
+					Data: bsoncore.BuildDocumentFromElements(nil,
+						bsoncore.AppendStringElement(nil, "disk", "ssd"),
+						bsoncore.AppendStringElement(nil, "use", "reporting"),
+					),
+				},
+			),
+			bsoncore.AppendInt32Element(nil, "maxStalenessSeconds", 25),
+			bsoncore.AppendDocumentElement(nil, "hedge", bsoncore.BuildDocumentFromElements(nil,
+				bsoncore.AppendBooleanElement(nil, "enabled", false),
+			)),
+		)
 
 		rpPrimaryPreferred := bsoncore.BuildDocumentFromElements(nil, bsoncore.AppendStringElement(nil, "mode", "primaryPreferred"))
 		rpPrimary := bsoncore.BuildDocumentFromElements(nil, bsoncore.AppendStringElement(nil, "mode", "primary"))
@@ -381,6 +403,27 @@ func TestOperation(t *testing.T) {
 				readpref.SecondaryPreferred(readpref.WithMaxStaleness(25 * time.Second)),
 				description.RSSecondary, description.ReplicaSet, false, rpWithMaxStaleness,
 			},
+			{
+				// A read preference document is generated for SecondaryPreferred if the hedge document is non-nil.
+				"secondaryPreferred with hedge to mongos using OP_QUERY",
+				readpref.SecondaryPreferred(readpref.WithHedgeEnabled(true)),
+				description.Mongos,
+				description.Sharded,
+				true,
+				rpWithHedge,
+			},
+			{
+				"secondaryPreferred with all options",
+				readpref.SecondaryPreferred(
+					readpref.WithTags("disk", "ssd", "use", "reporting"),
+					readpref.WithMaxStaleness(25*time.Second),
+					readpref.WithHedgeEnabled(false),
+				),
+				description.RSSecondary,
+				description.ReplicaSet,
+				false,
+				rpWithAllOptions,
+			},
 		}
 
 		for _, tc := range testCases {