Add new ServerSelector for PinnedServers

GODRIVER-1130

Change-Id: If6ec2e5137eff02ee8ee53b4fe7cdd676e4a26d2

Author: eppingere

mongo/client.go

@@ -454,13 +454,10 @@ func (c *Client) ListDatabases(ctx context.Context, filter interface{}, opts ...
 		return ListDatabasesResult{}, err
 	}
 
-	selector := description.CompositeSelector([]description.ServerSelector{
+	selector := makePinnedSelector(sess, description.CompositeSelector([]description.ServerSelector{
 		description.ReadPrefSelector(readpref.Primary()),
 		description.LatencySelector(c.localThreshold),
-	})
-	if sess != nil && sess.PinnedServer != nil {
-		selector = sess.PinnedServer
-	}
+	}))
 
 	ldo := options.MergeListDatabasesOptions(opts...)
 	op := operation.NewListDatabases(filterDoc).

mongo/collection.go

@@ -267,10 +267,7 @@ func (coll *Collection) insert(ctx context.Context, documents []interface{},
 		sess = nil
 	}
 
-	selector := coll.writeSelector
-	if sess != nil && sess.PinnedServer != nil {
-		selector = sess.PinnedServer
-	}
+	selector := makePinnedSelector(sess, coll.writeSelector)
 
 	op := operation.NewInsert(docs...).
 		Session(sess).WriteConcern(wc).CommandMonitor(coll.client.monitor).
@@ -386,10 +383,7 @@ func (coll *Collection) delete(ctx context.Context, filter interface{}, deleteOn
 		sess = nil
 	}
 
-	selector := coll.writeSelector
-	if sess != nil && sess.PinnedServer != nil {
-		selector = sess.PinnedServer
-	}
+	selector := makePinnedSelector(sess, coll.writeSelector)
 
 	var limit int32
 	if deleteOne {
@@ -490,10 +484,7 @@ func (coll *Collection) updateOrReplace(ctx context.Context, filter, update bson
 		sess = nil
 	}
 
-	selector := coll.writeSelector
-	if sess != nil && sess.PinnedServer != nil {
-		selector = sess.PinnedServer
-	}
+	selector := makePinnedSelector(sess, coll.writeSelector)
 
 	op := operation.NewUpdate(updateDoc).
 		Session(sess).WriteConcern(wc).CommandMonitor(coll.client.monitor).
@@ -672,13 +663,11 @@ func aggregate(a aggregateParams) (*Cursor, error) {
 		sess = nil
 	}
 
-	selector := a.readSelector
+	defaultSelector := a.readSelector
 	if hasDollarOut {
-		selector = a.writeSelector
-	}
-	if sess != nil && sess.PinnedServer != nil {
-		selector = sess.PinnedServer
+		defaultSelector = a.writeSelector
 	}
+	selector := makePinnedSelector(sess, defaultSelector)
 
 	ao := options.MergeAggregateOptions(a.opts...)
 	cursorOpts := driver.CursorOptions{
@@ -870,10 +859,7 @@ func (coll *Collection) Distinct(ctx context.Context, fieldName string, filter i
 		rc = nil
 	}
 
-	selector := coll.readSelector
-	if sess != nil && sess.PinnedServer != nil {
-		selector = sess.PinnedServer
-	}
+	selector := makePinnedSelector(sess, coll.readSelector)
 
 	option := options.MergeDistinctOptions(opts...)
 
@@ -951,10 +937,7 @@ func (coll *Collection) Find(ctx context.Context, filter interface{},
 		rc = nil
 	}
 
-	selector := coll.writeSelector
-	if sess != nil && sess.PinnedServer != nil {
-		selector = sess.PinnedServer
-	}
+	selector := makePinnedSelector(sess, coll.writeSelector)
 
 	op := operation.NewFind(f).
 		Session(sess).ReadConcern(rc).ReadPreference(coll.readPreference).
@@ -1142,10 +1125,7 @@ func (coll *Collection) findAndModify(ctx context.Context, op *operation.FindAnd
 		sess = nil
 	}
 
-	selector := coll.writeSelector
-	if sess != nil && sess.PinnedServer != nil {
-		selector = sess.PinnedServer
-	}
+	selector := makePinnedSelector(sess, coll.writeSelector)
 
 	retry := driver.RetryNone
 	if coll.client.retryWrites {
@@ -1378,10 +1358,7 @@ func (coll *Collection) Drop(ctx context.Context) error {
 		sess = nil
 	}
 
-	selector := coll.writeSelector
-	if sess != nil && sess.PinnedServer != nil {
-		selector = sess.PinnedServer
-	}
+	selector := makePinnedSelector(sess, coll.writeSelector)
 
 	op := operation.NewDropCollection().
 		Session(sess).WriteConcern(wc).CommandMonitor(coll.client.monitor).
@@ -1397,3 +1374,15 @@ func (coll *Collection) Drop(ctx context.Context) error {
 	}
 	return nil
 }
+
+// makePinnedSelector makes a selector for a pinned session with a pinned server. Will attempt to do server selection on
+// the pinned server but if that fails it will go through a list of default selectors
+func makePinnedSelector(sess *session.Client, defaultSelector description.ServerSelector) description.ServerSelectorFunc {
+	return func(t description.Topology, svrs []description.Server) ([]description.Server, error) {
+		if sess != nil && sess.PinnedServer != nil {
+			return sess.PinnedServer.SelectServer(t, svrs)
+		}
+
+		return defaultSelector.SelectServer(t, svrs)
+	}
+}

mongo/database.go

@@ -231,10 +231,7 @@ func (db *Database) Drop(ctx context.Context) error {
 		sess = nil
 	}
 
-	selector := db.writeSelector
-	if sess != nil && sess.PinnedServer != nil {
-		selector = sess.PinnedServer
-	}
+	selector := makePinnedSelector(sess, db.writeSelector)
 
 	op := operation.NewDropDatabase().
 		Session(sess).WriteConcern(wc).CommandMonitor(db.client.monitor).
@@ -275,10 +272,7 @@ func (db *Database) ListCollections(ctx context.Context, filter interface{}, opt
 		return nil, err
 	}
 
-	selector := db.readSelector
-	if sess != nil && sess.PinnedServer != nil {
-		selector = sess.PinnedServer
-	}
+	selector := makePinnedSelector(sess, db.readSelector)
 
 	lco := options.MergeListCollectionsOptions(opts...)
 	op := operation.NewListCollections(filterDoc).

mongo/session.go

@@ -217,12 +217,7 @@ func (s *sessionImpl) CommitTransaction(ctx context.Context) error {
 		s.clientSession.RetryingCommit = true
 	}
 
-	var selector description.ServerSelectorFunc = func(t description.Topology, svrs []description.Server) ([]description.Server, error) {
-		if s.clientSession.PinnedServer != nil {
-			return s.clientSession.PinnedServer.SelectServer(t, svrs)
-		}
-		return description.WriteSelector().SelectServer(t, svrs)
-	}
+	selector := makePinnedSelector(s.clientSession, description.WriteSelector())
 
 	s.clientSession.Committing = true
 	err = operation.NewCommitTransaction().