Change driver.KillCursors to not do server selection and fix tests.

GODRIVER-817
GODRIVER-819

Co-authored-by: Divjot Arora <[email protected]>

Change-Id: I78b457f32df9ea53ec650e1de3fb6aa75fc78174

Author: Isabella Siu

mongo/batch_cursor.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/mongodb/mongo-go-driver/x/bsonx/bsoncore"
+	"github.com/mongodb/mongo-go-driver/x/mongo/driver/topology"
 )
 
 // batchCursor is the interface implemented by types that can provide batches of document results.
@@ -19,6 +20,9 @@ type batchCursor interface {
 	// DocumentSequence is only valid until the next call to Next or Close.
 	Batch() *bsoncore.DocumentSequence
 
+	// Server returns a pointer to the cursor's server.
+	Server() *topology.Server
+
 	// Err returns the last error encountered.
 	Err() error
 

mongo/change_stream.go

@@ -453,12 +453,8 @@ func (cs *ChangeStream) Next(ctx context.Context) bool {
 			}
 		}
 
-		killCursors := command.KillCursors{
-			NS:  cs.ns,
-			IDs: []int64{cs.ID()},
-		}
+		_, _ = driver.KillCursors(ctx, cs.ns, cs.cursor.bc.Server(), cs.ID())
 
-		_, _ = driver.KillCursors(ctx, killCursors, cs.client.topology, cs.db.writeSelector)
 		cs.err = cs.runCommand(ctx, true)
 		if cs.err != nil {
 			return false

mongo/change_stream_test.go

@@ -63,6 +63,13 @@ func (er *errorCursor) Close(ctx context.Context) error {
 	return nil
 }
 
+func killChangeStreamCursor(t *testing.T, cs *ChangeStream) {
+	_, err := driver.KillCursors(context.Background(), cs.ns, cs.cursor.bc.Server(), cs.ID())
+	if err != nil {
+		t.Fatalf("error killing cursor: %v", err)
+	}
+}
+
 func skipIfBelow36(t *testing.T) {
 	serverVersion, err := getServerVersion(createTestDatabase(t, nil))
 	require.NoError(t, err)
@@ -355,15 +362,8 @@ func TestChangeStream_ReplicaSet(t *testing.T) {
 		ensureResumeToken(t, coll, stream)
 		cs := stream
 
-		kc := command.KillCursors{
-			NS:  cs.ns,
-			IDs: []int64{cs.ID()},
-		}
-
-		_, err := driver.KillCursors(ctx, kc, cs.client.topology, cs.db.writeSelector)
-		testhelpers.RequireNil(t, err, "error running killCursors cmd: %s", err)
-
-		_, err = coll.InsertOne(ctx, doc1)
+		killChangeStreamCursor(t, cs)
+		_, err := coll.InsertOne(ctx, doc1)
 		testhelpers.RequireNil(t, err, "error inserting doc: %s", err)
 
 		drainChannels()
@@ -478,13 +478,7 @@ func TestChangeStream_ReplicaSet(t *testing.T) {
 		cs := stream
 
 		// kill cursor to force a resumable error
-		kc := command.KillCursors{
-			NS:  cs.ns,
-			IDs: []int64{cs.ID()},
-		}
-
-		_, err = driver.KillCursors(ctx, kc, cs.client.topology, cs.db.writeSelector)
-		testhelpers.RequireNil(t, err, "error running killCursors cmd: %s", err)
+		killChangeStreamCursor(t, cs)
 
 		adminDb := coll.client.Database("admin")
 		modeDoc := bsonx.Doc{
@@ -527,14 +521,7 @@ func TestChangeStream_ReplicaSet(t *testing.T) {
 		cs := stream
 
 		// kill cursor to force a resumable error
-		kc := command.KillCursors{
-			NS:  cs.ns,
-			IDs: []int64{cs.ID()},
-		}
-
-		_, err = driver.KillCursors(ctx, kc, cs.client.topology, cs.db.writeSelector)
-		testhelpers.RequireNil(t, err, "error running killCursors cmd: %s", err)
-
+		killChangeStreamCursor(t, cs)
 		drainChannels()
 		stream.Next(ctx)
 
@@ -685,14 +672,7 @@ func TestChangeStream_ReplicaSet(t *testing.T) {
 
 		// kill the stream's underlying cursor to force a resumeable error
 		cs := stream
-		kc := command.KillCursors{
-			NS:  cs.ns,
-			IDs: []int64{cs.ID()},
-		}
-
-		_, err := driver.KillCursors(ctx, kc, cs.client.topology, cs.db.writeSelector)
-		testhelpers.RequireNil(t, err, "error running killCursors cmd: %s", err)
-
+		killChangeStreamCursor(t, cs)
 		ensureResumeToken(t, coll, stream)
 	})
 }

mongo/collection_internal_test.go

@@ -27,7 +27,6 @@ import (
 	"github.com/mongodb/mongo-go-driver/mongo/writeconcern"
 	"github.com/mongodb/mongo-go-driver/x/mongo/driver"
 	"github.com/mongodb/mongo-go-driver/x/network/command"
-	"github.com/mongodb/mongo-go-driver/x/network/wiremessage"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -1544,39 +1543,6 @@ func TestCollection_Find_notFound(t *testing.T) {
 	require.False(t, cursor.Next(context.Background()))
 }
 
-func killCursor(t *testing.T, c *Cursor, coll *Collection) {
-	version, err := getServerVersion(coll.db)
-	require.Nil(t, err, "error getting server version: %s", err)
-	ns := command.NewNamespace(coll.db.name, coll.name)
-
-	if compareVersions(t, version, "3.0") > 0 {
-		// not legacy
-		kc := command.KillCursors{
-			NS:  ns,
-			IDs: []int64{c.ID()},
-		}
-
-		_, err := driver.KillCursors(ctx, kc, coll.client.topology, coll.db.writeSelector)
-		require.Nil(t, err, "error killing cursor: %s", err)
-		return
-	}
-
-	// legacy
-	kc := wiremessage.KillCursors{
-		NumberOfCursorIDs: 1,
-		CursorIDs:         []int64{c.ID()},
-		CollectionName:    ns.Collection,
-		DatabaseName:      ns.DB,
-	}
-	topo := testutil.Topology(t)
-	ss, err := topo.SelectServer(ctx, coll.db.writeSelector)
-	require.Nil(t, err, "error selecting server: %s", err)
-	conn, err := ss.Connection(ctx)
-	require.Nil(t, err, "error getting connection: %s", err)
-	err = conn.WriteWireMessage(context.Background(), kc)
-	require.Nil(t, err, "error writing wire msg: %s", err)
-}
-
 func TestCollection_Find_Error(t *testing.T) {
 	t.Run("TestInvalidIdentifier", func(t *testing.T) {
 		coll := createTestCollection(t, nil, nil)
@@ -1595,7 +1561,11 @@ func TestCollection_Find_Error(t *testing.T) {
 		require.True(t, c.Next(context.Background()))
 		require.True(t, c.Next(context.Background()))
 
-		killCursor(t, c, coll)
+		_, err = driver.KillCursors(ctx, command.Namespace{
+			DB:         coll.db.name,
+			Collection: coll.name,
+		}, c.bc.Server(), c.ID())
+		require.NoError(t, err)
 		require.False(t, c.Next(context.Background()))
 		require.NotNil(t, c.Err())
 		_ = c.Close(context.Background())

mongo/sessions_test.go

@@ -91,6 +91,9 @@ func createFuncMap(t *testing.T, dbName string, collName string, monitored bool)
 	err := db.Drop(ctx)
 	testhelpers.RequireNil(t, err, "error dropping database after creation: %s", err)
 
+	// ensure database exists
+	_, _ = db.Collection("foo").InsertOne(context.Background(), doc)
+
 	coll := db.Collection(collName)
 	iv := coll.Indexes()
 

x/mongo/driver/batch_cursor.go

@@ -174,6 +174,9 @@ func (bc *BatchCursor) Next(ctx context.Context) bool {
 // DocumentSequence is only valid until the next call to Next or Close.
 func (bc *BatchCursor) Batch() *bsoncore.DocumentSequence { return bc.currentBatch }
 
+// Server returns a pointer to the cursor's server.
+func (bc *BatchCursor) Server() *topology.Server { return bc.server }
+
 // Err returns the latest error encountered.
 func (bc *BatchCursor) Err() error { return bc.err }
 

x/mongo/driver/kill_cursors.go

@@ -8,30 +8,48 @@ package driver
 
 import (
 	"context"
+	"github.com/mongodb/mongo-go-driver/x/network/connection"
+	"github.com/mongodb/mongo-go-driver/x/network/wiremessage"
 
 	"github.com/mongodb/mongo-go-driver/x/mongo/driver/topology"
 	"github.com/mongodb/mongo-go-driver/x/network/command"
-	"github.com/mongodb/mongo-go-driver/x/network/description"
 	"github.com/mongodb/mongo-go-driver/x/network/result"
 )
 
 // KillCursors handles the full cycle dispatch and execution of an aggregate command against the provided
 // topology.
 func KillCursors(
 	ctx context.Context,
-	cmd command.KillCursors,
-	topo *topology.Topology,
-	selector description.ServerSelector,
+	ns command.Namespace,
+	server *topology.Server,
+	cursorID int64,
 ) (result.KillCursors, error) {
-	ss, err := topo.SelectServer(ctx, selector)
-	if err != nil {
-		return result.KillCursors{}, err
-	}
-	desc := ss.Description()
-	conn, err := ss.Connection(ctx)
+	desc := server.SelectedDescription()
+	conn, err := server.Connection(ctx)
 	if err != nil {
 		return result.KillCursors{}, err
 	}
 	defer conn.Close()
+
+	if desc.WireVersion.Max < 4 {
+		return result.KillCursors{}, legacyKillCursors(ctx, ns, cursorID, conn)
+	}
+
+	cmd := command.KillCursors{
+		NS:  ns,
+		IDs: []int64{cursorID},
+	}
+
 	return cmd.RoundTrip(ctx, desc, conn)
 }
+
+func legacyKillCursors(ctx context.Context, ns command.Namespace, cursorID int64, conn connection.Connection) error {
+	kc := wiremessage.KillCursors{
+		NumberOfCursorIDs: 1,
+		CursorIDs:         []int64{cursorID},
+		CollectionName:    ns.Collection,
+		DatabaseName:      ns.DB,
+	}
+
+	return conn.WriteWireMessage(ctx, kc)
+}

x/mongo/driver/list_collections_batch_cursor.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 
 	"github.com/mongodb/mongo-go-driver/x/bsonx/bsoncore"
+	"github.com/mongodb/mongo-go-driver/x/mongo/driver/topology"
 )
 
 // ListCollectionsBatchCursor is a special batch cursor returned from ListCollections that properly
@@ -83,6 +84,9 @@ func (lcbc *ListCollectionsBatchCursor) Next(ctx context.Context) bool {
 // DocumentSequence is only valid until the next call to Next or Close.
 func (lcbc *ListCollectionsBatchCursor) Batch() *bsoncore.DocumentSequence { return lcbc.currentBatch }
 
+// Server returns a pointer to the cursor's server.
+func (lcbc *ListCollectionsBatchCursor) Server() *topology.Server { return lcbc.bc.server }
+
 // Err returns the latest error encountered.
 func (lcbc *ListCollectionsBatchCursor) Err() error {
 	if lcbc.err != nil {