GODRIVER-2626 further optimize selectByKind for special case of all fitting servers (#1107)

Author: isopov

mongo/description/selector_test.go

@@ -300,32 +300,69 @@ func BenchmarkLatencySelector(b *testing.B) {
 }
 
 func BenchmarkSelector_Sharded(b *testing.B) {
-	subject := readpref.Primary()
+	for _, bcase := range []struct {
+		name        string
+		serversHook func(servers []Server)
+	}{
+		{
+			name:        "AllFit",
+			serversHook: func(servers []Server) {},
+		},
+		{
+			name: "AllButOneFit",
+			serversHook: func(servers []Server) {
+				servers[0].Kind = LoadBalancer
+			},
+		},
+		{
+			name: "HalfFit",
+			serversHook: func(servers []Server) {
+				for i := 0; i < len(servers); i += 2 {
+					servers[i].Kind = LoadBalancer
+				}
+			},
+		},
+		{
+			name: "OneFit",
+			serversHook: func(servers []Server) {
+				for i := 1; i < len(servers); i++ {
+					servers[i].Kind = LoadBalancer
+				}
+			},
+		},
+	} {
+		bcase := bcase
 
-	s := Server{
-		Addr:              address.Address("localhost:27017"),
-		HeartbeatInterval: time.Duration(10) * time.Second,
-		LastWriteTime:     time.Date(2017, 2, 11, 14, 0, 0, 0, time.UTC),
-		LastUpdateTime:    time.Date(2017, 2, 11, 14, 0, 2, 0, time.UTC),
-		Kind:              Mongos,
-		WireVersion:       &VersionRange{Min: 0, Max: 5},
-	}
-	servers := make([]Server, 100)
-	for i := 0; i < len(servers); i++ {
-		servers[i] = s
-	}
-	servers[0].Kind = LoadBalancer
-	c := Topology{
-		Kind:    Sharded,
-		Servers: servers,
-	}
+		b.Run(bcase.name, func(b *testing.B) {
+			subject := readpref.Primary()
 
-	b.ResetTimer()
-	b.RunParallel(func(p *testing.PB) {
-		for p.Next() {
-			_, _ = ReadPrefSelector(subject).SelectServer(c, c.Servers)
-		}
-	})
+			s := Server{
+				Addr:              address.Address("localhost:27017"),
+				HeartbeatInterval: time.Duration(10) * time.Second,
+				LastWriteTime:     time.Date(2017, 2, 11, 14, 0, 0, 0, time.UTC),
+				LastUpdateTime:    time.Date(2017, 2, 11, 14, 0, 2, 0, time.UTC),
+				Kind:              Mongos,
+				WireVersion:       &VersionRange{Min: 0, Max: 5},
+			}
+			servers := make([]Server, 100)
+			for i := 0; i < len(servers); i++ {
+				servers[i] = s
+			}
+			bcase.serversHook(servers)
+			c := Topology{
+				Kind:    Sharded,
+				Servers: servers,
+			}
+
+			b.ResetTimer()
+			b.RunParallel(func(p *testing.PB) {
+				b.ReportAllocs()
+				for p.Next() {
+					_, _ = ReadPrefSelector(subject).SelectServer(c, c.Servers)
+				}
+			})
+		})
+	}
 }
 
 func TestSelector_Single(t *testing.T) {

mongo/description/server_selector.go

@@ -310,6 +310,9 @@ func selectByKind(candidates []Server, kind ServerKind) []Server {
 			viableIndexes = append(viableIndexes, i)
 		}
 	}
+	if len(viableIndexes) == len(candidates) {
+		return candidates
+	}
 	result := make([]Server, len(viableIndexes))
 	for i, idx := range viableIndexes {
 		result[i] = candidates[idx]