GODRIVER-1535 Fix session IDs batching in Disconnect

Author: Divjot Arora

mongo/client.go

@@ -10,7 +10,6 @@ import (
 	"context"
 	"crypto/tls"
 	"errors"
-	"strconv"
 	"strings"
 	"time"
 
@@ -21,6 +20,7 @@ import (
 	"go.mongodb.org/mongo-driver/mongo/readconcern"
 	"go.mongodb.org/mongo-driver/mongo/readpref"
 	"go.mongodb.org/mongo-driver/mongo/writeconcern"
+	"go.mongodb.org/mongo-driver/x/bsonx"
 	"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
 	"go.mongodb.org/mongo-driver/x/mongo/driver"
 	"go.mongodb.org/mongo-driver/x/mongo/driver/auth"
@@ -33,11 +33,14 @@ import (
 )
 
 const defaultLocalThreshold = 15 * time.Millisecond
-const batchSize = 10000
 
-// keyVaultCollOpts specifies options used to communicate with the key vault collection
-var keyVaultCollOpts = options.Collection().SetReadConcern(readconcern.Majority()).
-	SetWriteConcern(writeconcern.New(writeconcern.WMajority()))
+var (
+	// keyVaultCollOpts specifies options used to communicate with the key vault collection
+	keyVaultCollOpts = options.Collection().SetReadConcern(readconcern.Majority()).
+				SetWriteConcern(writeconcern.New(writeconcern.WMajority()))
+
+	endSessionsBatchSize = 10000
+)
 
 // Client is a handle representing a pool of connections to a MongoDB deployment. It is safe for concurrent use by
 // multiple goroutines.
@@ -287,31 +290,27 @@ func (c *Client) endSessions(ctx context.Context) {
 		return
 	}
 
-	ids := c.sessionPool.IDSlice()
-	idx, idArray := bsoncore.AppendArrayStart(nil)
-	for i, id := range ids {
-		idDoc, _ := id.MarshalBSON()
-		idArray = bsoncore.AppendDocumentElement(idArray, strconv.Itoa(i), idDoc)
-	}
-	idArray, _ = bsoncore.AppendArrayEnd(idArray, idx)
-
-	op := operation.NewEndSessions(idArray).ClusterClock(c.clock).Deployment(c.deployment).
+	sessionIDs := c.sessionPool.IDSlice()
+	op := operation.NewEndSessions(nil).ClusterClock(c.clock).Deployment(c.deployment).
 		ServerSelector(description.ReadPrefSelector(readpref.PrimaryPreferred())).CommandMonitor(c.monitor).
 		Database("admin").Crypt(c.crypt)
 
-	idx, idArray = bsoncore.AppendArrayStart(nil)
-	totalNumIDs := len(ids)
+	totalNumIDs := len(sessionIDs)
+	var currentBatch []bsonx.Doc
 	for i := 0; i < totalNumIDs; i++ {
-		idDoc, _ := ids[i].MarshalBSON()
-		idArray = bsoncore.AppendDocumentElement(idArray, strconv.Itoa(i), idDoc)
-		if ((i+1)%batchSize) == 0 || i == totalNumIDs-1 {
-			idArray, _ = bsoncore.AppendArrayEnd(idArray, idx)
-			_ = op.SessionIDs(idArray).Execute(ctx)
-			idArray = idArray[:0]
-			idx = 0
+		currentBatch = append(currentBatch, sessionIDs[i])
+
+		// If we are at the end of a batch or the end of the overall IDs array, execute the operation.
+		if ((i+1)%endSessionsBatchSize) == 0 || i == totalNumIDs-1 {
+			// Ignore all errors when ending sessions.
+			_, marshalVal, err := bson.MarshalValue(currentBatch)
+			if err == nil {
+				_ = op.SessionIDs(marshalVal).Execute(ctx)
+			}
+
+			currentBatch = currentBatch[:0]
 		}
 	}
-
 }
 
 func (c *Client) configure(opts *options.ClientOptions) error {

mongo/client_test.go

@@ -9,10 +9,13 @@ package mongo
 import (
 	"context"
 	"errors"
+	"math"
 	"testing"
 	"time"
 
 	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/event"
+	"go.mongodb.org/mongo-driver/internal/testutil"
 	"go.mongodb.org/mongo-driver/internal/testutil/assert"
 	"go.mongodb.org/mongo-driver/mongo/options"
 	"go.mongodb.org/mongo-driver/mongo/readconcern"
@@ -257,4 +260,87 @@ func TestClient(t *testing.T) {
 			assert.Equal(t, uri, got, "expected GetURI to return %v, got %v", uri, got)
 		})
 	})
+	t.Run("endSessions", func(t *testing.T) {
+		cs := testutil.ConnString(t)
+		originalBatchSize := endSessionsBatchSize
+		endSessionsBatchSize = 2
+		defer func() {
+			endSessionsBatchSize = originalBatchSize
+		}()
+
+		testCases := []struct {
+			name            string
+			numSessions     int
+			eventBatchSizes []int
+		}{
+			{"number of sessions divides evenly", endSessionsBatchSize * 2, []int{endSessionsBatchSize, endSessionsBatchSize}},
+			{"number of sessions does not divide evenly", endSessionsBatchSize + 1, []int{endSessionsBatchSize, 1}},
+		}
+		for _, tc := range testCases {
+			t.Run(tc.name, func(t *testing.T) {
+				// Setup a client and skip the test based on server version.
+				var started []*event.CommandStartedEvent
+				var failureReasons []string
+				cmdMonitor := &event.CommandMonitor{
+					Started: func(_ context.Context, evt *event.CommandStartedEvent) {
+						if evt.CommandName == "endSessions" {
+							started = append(started, evt)
+						}
+					},
+					Failed: func(_ context.Context, evt *event.CommandFailedEvent) {
+						if evt.CommandName == "endSessions" {
+							failureReasons = append(failureReasons, evt.Failure)
+						}
+					},
+				}
+				clientOpts := options.Client().ApplyURI(cs.Original).SetReadPreference(readpref.Primary()).
+					SetWriteConcern(writeconcern.New(writeconcern.WMajority())).SetMonitor(cmdMonitor)
+				client, err := Connect(bgCtx, clientOpts)
+				assert.Nil(t, err, "Connect error: %v", err)
+				defer func() {
+					_ = client.Disconnect(bgCtx)
+				}()
+
+				serverVersion, err := getServerVersion(client.Database("admin"))
+				assert.Nil(t, err, "getServerVersion error: %v", err)
+				if compareVersions(t, serverVersion, "3.6.0") < 1 {
+					t.Skip("skipping server version < 3.6")
+				}
+
+				coll := client.Database("foo").Collection("bar")
+				defer func() {
+					_ = coll.Drop(bgCtx)
+				}()
+
+				// Do an application operation and create the number of sessions specified by the test.
+				_, err = coll.CountDocuments(bgCtx, bson.D{})
+				assert.Nil(t, err, "CountDocuments error: %v", err)
+				var sessions []Session
+				for i := 0; i < tc.numSessions; i++ {
+					sess, err := client.StartSession()
+					assert.Nil(t, err, "StartSession error at index %d: %v", i, err)
+					sessions = append(sessions, sess)
+				}
+				for _, sess := range sessions {
+					sess.EndSession(bgCtx)
+				}
+
+				client.endSessions(bgCtx)
+				divisionResult := float64(tc.numSessions) / float64(endSessionsBatchSize)
+				numEventsExpected := int(math.Ceil(divisionResult))
+				assert.Equal(t, len(started), numEventsExpected, "expected %d started events, got %d", numEventsExpected,
+					len(started))
+				assert.Equal(t, len(failureReasons), 0, "endSessions errors: %v", failureReasons)
+
+				for i := 0; i < numEventsExpected; i++ {
+					sentArray := started[i].Command.Lookup("endSessions").Array()
+					values, _ := sentArray.Values()
+					expectedNumValues := tc.eventBatchSizes[i]
+					assert.Equal(t, len(values), expectedNumValues,
+						"batch size mismatch at index %d; expected %d sessions in batch, got %d", i, expectedNumValues,
+						len(values))
+				}
+			})
+		}
+	})
 }