GODRIVER-1404 Remove topology subscription initial server selection (#221)

Author: Divjot Arora

x/mongo/driver/topology/topology.go

@@ -91,6 +91,18 @@ type Topology struct {
 var _ driver.Deployment = &Topology{}
 var _ driver.Subscriber = &Topology{}
 
+type serverSelectionState struct {
+	selector    description.ServerSelector
+	timeoutChan <-chan time.Time
+}
+
+func newServerSelectionState(selector description.ServerSelector, timeoutChan <-chan time.Time) serverSelectionState {
+	return serverSelectionState{
+		selector:    selector,
+		timeoutChan: timeoutChan,
+	}
+}
+
 // New creates a new topology.
 func New(opts ...Option) (*Topology, error) {
 	cfg, err := newConfig(opts...)
@@ -295,16 +307,39 @@ func (t *Topology) SelectServer(ctx context.Context, ss description.ServerSelect
 		defer ssTimeout.Stop()
 	}
 
-	sub, err := t.Subscribe()
-	if err != nil {
-		return nil, err
-	}
-	defer t.Unsubscribe(sub)
-
+	var doneOnce bool
+	var sub *driver.Subscription
+	selectionState := newServerSelectionState(ss, ssTimeoutCh)
 	for {
-		suitable, err := t.selectServer(ctx, sub.Updates, ss, ssTimeoutCh)
-		if err != nil {
-			return nil, err
+		var suitable []description.Server
+		var selectErr error
+
+		if !doneOnce {
+			// for the first pass, select a server from the current description.
+			// this improves selection speed for up-to-date topology descriptions.
+			suitable, selectErr = t.selectServerFromDescription(ctx, t.Description(), selectionState)
+			doneOnce = true
+		} else {
+			// if the first pass didn't select a server, the previous description did not contain a suitable server, so
+			// we subscribe to the topology and attempt to obtain a server from that subscription
+			if sub == nil {
+				var err error
+				sub, err = t.Subscribe()
+				if err != nil {
+					return nil, err
+				}
+				defer t.Unsubscribe(sub)
+			}
+
+			suitable, selectErr = t.selectServerFromSubscription(ctx, sub.Updates, selectionState)
+		}
+		if selectErr != nil {
+			return nil, selectErr
+		}
+
+		if len(suitable) == 0 {
+			// try again if there are no servers available
+			continue
 		}
 
 		selected := suitable[rand.Intn(len(suitable))]
@@ -344,8 +379,9 @@ func (t *Topology) SelectServerLegacy(ctx context.Context, ss description.Server
 	}
 	defer t.Unsubscribe(sub)
 
+	selectionState := newServerSelectionState(ss, ssTimeoutCh)
 	for {
-		suitable, err := t.selectServer(ctx, sub.Updates, ss, ssTimeoutCh)
+		suitable, err := t.selectServerFromSubscription(ctx, sub.Updates, selectionState)
 		if err != nil {
 			return nil, err
 		}
@@ -390,39 +426,60 @@ func wrapServerSelectionError(err error, t *Topology) error {
 	return fmt.Errorf("server selection error: %v, current topology: { %s }", err, t.String())
 }
 
-// selectServer is the core piece of server selection. It handles getting
-// topology descriptions and running sever selection on those descriptions.
-func (t *Topology) selectServer(ctx context.Context, subscriptionCh <-chan description.Topology, ss description.ServerSelector, timeoutCh <-chan time.Time) ([]description.Server, error) {
+// selectServerFromSubscription loops until a topology description is available for server selection. It returns
+// when the given context expires, server selection timeout is reached, or a description containing a selectable
+// server is available.
+func (t *Topology) selectServerFromSubscription(ctx context.Context, subscriptionCh <-chan description.Topology,
+	selectionState serverSelectionState) ([]description.Server, error) {
+
 	var current description.Topology
 	for {
 		select {
 		case <-ctx.Done():
 			return nil, ctx.Err()
-		case <-timeoutCh:
+		case <-selectionState.timeoutChan:
 			return nil, wrapServerSelectionError(ErrServerSelectionTimeout, t)
 		case current = <-subscriptionCh:
 		}
 
-		var allowed []description.Server
-		for _, s := range current.Servers {
-			if s.Kind != description.Unknown {
-				allowed = append(allowed, s)
-			}
-		}
-
-		suitable, err := ss.SelectServer(current, allowed)
+		suitable, err := t.selectServerFromDescription(ctx, current, selectionState)
 		if err != nil {
-			return nil, wrapServerSelectionError(err, t)
+			return nil, err
 		}
 
 		if len(suitable) > 0 {
 			return suitable, nil
 		}
-
 		t.RequestImmediateCheck()
 	}
 }
 
+// selectServerFromDescription process the given topology description and returns a slice of suitable servers.
+func (t *Topology) selectServerFromDescription(ctx context.Context, desc description.Topology,
+	selectionState serverSelectionState) ([]description.Server, error) {
+
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	case <-selectionState.timeoutChan:
+		return nil, wrapServerSelectionError(ErrServerSelectionTimeout, t)
+	default:
+	}
+
+	var allowed []description.Server
+	for _, s := range desc.Servers {
+		if s.Kind != description.Unknown {
+			allowed = append(allowed, s)
+		}
+	}
+
+	suitable, err := selectionState.selector.SelectServer(desc, allowed)
+	if err != nil {
+		return nil, wrapServerSelectionError(err, t)
+	}
+	return suitable, nil
+}
+
 func (t *Topology) pollSRVRecords() {
 	defer t.pollingwg.Done()
 

x/mongo/driver/topology/topology_test.go

@@ -13,6 +13,7 @@ import (
 	"testing"
 	"time"
 
+	"go.mongodb.org/mongo-driver/internal/testutil/assert"
 	"go.mongodb.org/mongo-driver/x/mongo/driver"
 	"go.mongodb.org/mongo-driver/x/mongo/driver/address"
 	"go.mongodb.org/mongo-driver/x/mongo/driver/connstring"
@@ -72,7 +73,9 @@ func TestServerSelection(t *testing.T) {
 		}
 		subCh := make(chan description.Topology, 1)
 		subCh <- desc
-		srvs, err := topo.selectServer(context.Background(), subCh, selectFirst, nil)
+
+		state := newServerSelectionState(selectFirst, nil)
+		srvs, err := topo.selectServerFromSubscription(context.Background(), subCh, state)
 		noerr(t, err)
 		if len(srvs) != 1 {
 			t.Errorf("Incorrect number of descriptions returned. got %d; want %d", len(srvs), 1)
@@ -90,7 +93,8 @@ func TestServerSelection(t *testing.T) {
 
 		resp := make(chan []description.Server)
 		go func() {
-			srvs, err := topo.selectServer(context.Background(), subCh, selectFirst, nil)
+			state := newServerSelectionState(selectFirst, nil)
+			srvs, err := topo.selectServerFromSubscription(context.Background(), subCh, state)
 			noerr(t, err)
 			resp <- srvs
 		}()
@@ -137,7 +141,8 @@ func TestServerSelection(t *testing.T) {
 		resp := make(chan error)
 		ctx, cancel := context.WithCancel(context.Background())
 		go func() {
-			_, err := topo.selectServer(ctx, subCh, selectNone, nil)
+			state := newServerSelectionState(selectNone, nil)
+			_, err := topo.selectServerFromSubscription(ctx, subCh, state)
 			resp <- err
 		}()
 
@@ -174,7 +179,8 @@ func TestServerSelection(t *testing.T) {
 		resp := make(chan error)
 		timeout := make(chan time.Time)
 		go func() {
-			_, err := topo.selectServer(context.Background(), subCh, selectNone, timeout)
+			state := newServerSelectionState(selectNone, timeout)
+			_, err := topo.selectServerFromSubscription(context.Background(), subCh, state)
 			resp <- err
 		}()
 
@@ -209,7 +215,8 @@ func TestServerSelection(t *testing.T) {
 		resp := make(chan error)
 		timeout := make(chan time.Time)
 		go func() {
-			_, err := topo.selectServer(context.Background(), subCh, selectError, timeout)
+			state := newServerSelectionState(selectError, timeout)
+			_, err := topo.selectServerFromSubscription(context.Background(), subCh, state)
 			resp <- err
 		}()
 
@@ -288,7 +295,8 @@ func TestServerSelection(t *testing.T) {
 
 		go func() {
 			// server selection should discover the new topology
-			srvs, err := topo.selectServer(context.Background(), subCh, description.WriteSelector(), nil)
+			state := newServerSelectionState(description.WriteSelector(), nil)
+			srvs, err := topo.selectServerFromSubscription(context.Background(), subCh, state)
 			noerr(t, err)
 			resp <- srvs
 		}()
@@ -307,6 +315,33 @@ func TestServerSelection(t *testing.T) {
 			t.Errorf("Incorrect sever selected. got %s; want %s", srvs[0].Addr, desc.Servers[1].Addr)
 		}
 	})
+	t.Run("no subscription in fast path", func(t *testing.T) {
+		// assert that a subscription is not created if there is a server available
+		topo, err := New()
+		noerr(t, err)
+		topo.cfg.cs.HeartbeatInterval = time.Minute
+		atomic.StoreInt32(&topo.connectionstate, connected)
+
+		primaryAddr := address.Address("one")
+		desc := description.Topology{
+			Servers: []description.Server{
+				{Addr: primaryAddr, Kind: description.RSPrimary},
+			},
+		}
+		topo.desc.Store(desc)
+		for _, srv := range desc.Servers {
+			s, err := ConnectServer(srv.Addr, func(desc description.Server) { topo.apply(context.Background(), desc) })
+			noerr(t, err)
+			topo.servers[srv.Addr] = s
+		}
+
+		// manually close subscriptions so calls to Subscribe will error
+		topo.subscriptionsClosed = true
+		selectedServer, err := topo.SelectServer(context.Background(), description.WriteSelector())
+		noerr(t, err)
+		selectedAddr := selectedServer.(*SelectedServer).address
+		assert.Equal(t, primaryAddr, selectedAddr, "expected address %v, got %v", primaryAddr, selectedAddr)
+	})
 }
 
 func TestSessionTimeout(t *testing.T) {