snapshot of https://github.com/grpc/grpc-go.git (master, 2026-09-15; blobs > 131072B omitted)

balancer/autosharding/slice_map.go

@@ -0,0 +1,120 @@
+/*
+ *
+ * Copyright 2026 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package autosharding
+
+import (
+	"bytes"
+	"slices"
+)
+
+// sliceMapEntry represents an entry for a key-range in the sliceMap.
+type sliceMapEntry struct {
+	startKey  []byte // Inclusive start key of the key-range
+	endpoints []int  // Indices into list[PickerEndpoint] in the Picker
+}
+
+// sliceMap is an immutable data structure used by the Picker to map request
+// keys to their assigned key-range. Because it is immutable, it allows the
+// Picker to access it without any explicit synchronization with the LB policy.
+//
+// The sliceMap is meant to be used by the Picker in conjunction with a list of
+// pickerEndpoints such that the list can be swapped out, as long as the number
+// and order of endpoints don't change.
+type sliceMap struct {
+	slices       []sliceMapEntry // Sorted by startKey
+	fallbackPool []int           // Indices into list[PickerEndpoint] for all endpoints
+	generation   int64           // Snapshot generation number
+}
+
+// lookup returns the index of the slice covering the given key.
+//
+// Because assignments are pre-validated to have no gaps and cover the full key
+// range, and since slices is sorted by startKey, lookup boils down to a binary
+// search, looking for the insertion point.
+//
+// If the key matches the startKey of a slice, that slice index is returned.
+// Otherwise, the index of the slice immediately preceding the insertion point
+// is returned (i.e., the slice covering the range [startKey, nextStartKey)).
+//
+// Returns -1 when the sliceMap is empty, which is the case before the first
+// assignment is received.
+func (sm *sliceMap) lookup(key []byte) int {
+	if len(sm.slices) == 0 {
+		return -1
+	}
+
+	idx, found := slices.BinarySearchFunc(sm.slices, key, func(e sliceMapEntry, k []byte) int {
+		return bytes.Compare(e.startKey, k)
+	})
+
+	if found {
+		return idx
+	}
+
+	// Key falls in range [slices[idx - 1].startKey, slices[idx].startKey).
+	return idx - 1
+}
+
+// buildSliceMap is used to generate a new sliceMap from the EndpointMap and
+// Assignment when either of them change.
+func buildSliceMap(endpointMap *endpointMap, assignment *assignment) *sliceMap {
+	sm := &sliceMap{}
+
+	// Populate fallbackPool with values [0, 1, 2, ... N-1] where N is the
+	// number of endpoints returned by the name resolver. The only reason to
+	// have this slice is to allow the Picker to share code that picks a random
+	// endpoint from a slice of indices, either from the fallbackPool or from a
+	// sliceMapEntry.
+	sm.fallbackPool = make([]int, len(endpointMap.m))
+	for i := range sm.fallbackPool {
+		sm.fallbackPool[i] = i
+	}
+
+	// If no assignment has been received yet (startup case), return early with
+	// empty slices.
+	if assignment == nil {
+		return sm
+	}
+
+	sm.generation = assignment.generation
+
+	// Build sliceMapEntry for each Slice in the assignment.
+	sm.slices = make([]sliceMapEntry, 0, len(assignment.slices))
+	for _, s := range assignment.slices {
+		entry := sliceMapEntry{
+			startKey:  s.startKey,
+			endpoints: make([]int, 0, len(s.endpoints)),
+		}
+
+		// Populate the endpoints for this slice by looking up the endpoint
+		// names in the EndpointMap. If an endpoint name is not found in the
+		// EndpointMap, it is skipped. This ensures that the sliceMap only
+		// contains valid endpoints that are currently available.
+		for _, idx := range s.endpoints {
+			hostname := assignment.endpointNames[idx]
+			if es, ok := endpointMap.m[hostname]; ok {
+				entry.endpoints = append(entry.endpoints, es.index)
+			}
+		}
+
+		sm.slices = append(sm.slices, entry)
+	}
+
+	return sm
+}