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"
 
@@ -34,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.
@@ -288,29 +290,27 @@ func (c *Client) endSessions(ctx context.Context) {
 		return
 	}
 
-	ids := c.sessionPool.IDSlice()
-	idx, idArray := bsoncore.AppendArrayStart(nil)
-	for i, id := range ids {
-		idArray = bsoncore.AppendDocumentElement(idArray, strconv.Itoa(i), id)
-	}
-	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 []bsoncore.Document
 	for i := 0; i < totalNumIDs; i++ {
-		idArray = bsoncore.AppendDocumentElement(idArray, strconv.Itoa(i), ids[i])
-		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

@@ -13,6 +13,8 @@ import (
 	"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 +259,64 @@ 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)
+
+		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)
+		coll := client.Database("foo").Collection("bar")
+
+		// Lower the batch size to force multiple batches.
+		originalBatchSize := endSessionsBatchSize
+		endSessionsBatchSize = 2
+		defer func() {
+			endSessionsBatchSize = originalBatchSize
+			_ = coll.Drop(bgCtx)
+			_ = client.Disconnect(bgCtx)
+		}()
+
+		// Do an application operation and create four sessions so endSessions will execute in two batches.
+		_, err = coll.CountDocuments(bgCtx, bson.D{})
+		assert.Nil(t, err, "CountDocuments error: %v", err)
+		numSessions := 4
+		var sessions []Session
+		for i := 0; i < 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)
+		numEventsExpected := numSessions / endSessionsBatchSize
+		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()
+			assert.Equal(t, len(values), endSessionsBatchSize,
+				"batch size mismatch at index %d; expected %d sessions in batch, got %d", i, endSessionsBatchSize,
+				len(values))
+		}
+	})
 }