GODRIVER-2628 optimize allocations in selectServerFromDescription (#1111)

Author: isopov

x/mongo/driver/topology/topology.go

@@ -558,10 +558,19 @@ func (t *Topology) selectServerFromDescription(desc description.Topology,
 		return desc.Servers, nil
 	}
 
-	var allowed []description.Server
-	for _, s := range desc.Servers {
+	allowedIndexes := make([]int, 0, len(desc.Servers))
+	for i, s := range desc.Servers {
 		if s.Kind != description.Unknown {
-			allowed = append(allowed, s)
+			allowedIndexes = append(allowedIndexes, i)
+		}
+	}
+	var allowed []description.Server
+	if len(allowedIndexes) == len(desc.Servers) {
+		allowed = desc.Servers
+	} else {
+		allowed = make([]description.Server, len(allowedIndexes))
+		for i, idx := range allowedIndexes {
+			allowed[i] = desc.Servers[idx]
 		}
 	}
 

x/mongo/driver/topology/topology_test.go

@@ -835,3 +835,68 @@ func serverKindFromString(t *testing.T, s string) description.ServerKind {
 
 	return description.Unknown
 }
+
+func BenchmarkSelectServerFromDescription(b *testing.B) {
+	for _, bcase := range []struct {
+		name        string
+		serversHook func(servers []description.Server)
+	}{
+		{
+			name:        "AllFit",
+			serversHook: func(servers []description.Server) {},
+		},
+		{
+			name: "AllButOneFit",
+			serversHook: func(servers []description.Server) {
+				servers[0].Kind = description.Unknown
+			},
+		},
+		{
+			name: "HalfFit",
+			serversHook: func(servers []description.Server) {
+				for i := 0; i < len(servers); i += 2 {
+					servers[i].Kind = description.Unknown
+				}
+			},
+		},
+		{
+			name: "OneFit",
+			serversHook: func(servers []description.Server) {
+				for i := 1; i < len(servers); i++ {
+					servers[i].Kind = description.Unknown
+				}
+			},
+		},
+	} {
+		bcase := bcase
+
+		b.Run(bcase.name, func(b *testing.B) {
+			s := description.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:              description.Mongos,
+				WireVersion:       &description.VersionRange{Min: 0, Max: 5},
+			}
+			servers := make([]description.Server, 100)
+			for i := 0; i < len(servers); i++ {
+				servers[i] = s
+			}
+			bcase.serversHook(servers)
+			desc := description.Topology{
+				Servers: servers,
+			}
+
+			timeout := make(chan time.Time)
+			b.ResetTimer()
+			b.RunParallel(func(p *testing.PB) {
+				b.ReportAllocs()
+				for p.Next() {
+					var c Topology
+					_, _ = c.selectServerFromDescription(desc, newServerSelectionState(selectNone, timeout))
+				}
+			})
+		})
+	}
+}