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

internal/xds/balancer/clusterimpl/picker.go

@@ -0,0 +1,239 @@
+/*
+ *
+ * Copyright 2020 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 clusterimpl
+
+import (
+	"context"
+
+	v3orcapb "github.com/cncf/xds/go/xds/data/orca/v3"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/internal/envconfig"
+	"google.golang.org/grpc/internal/stats"
+	"google.golang.org/grpc/internal/wrr"
+	xdsinternal "google.golang.org/grpc/internal/xds"
+	"google.golang.org/grpc/internal/xds/clients"
+	"google.golang.org/grpc/internal/xds/xdsclient"
+	"google.golang.org/grpc/internal/xds/xdsclient/xdsresource"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+)
+
+// NewRandomWRR is used when calculating drops. It's exported so that tests can
+// override it.
+var NewRandomWRR = wrr.NewRandom
+
+const million = 1000000
+
+type dropper struct {
+	category string
+	w        wrr.WRR
+}
+
+// greatest common divisor (GCD) via Euclidean algorithm
+func gcd(a, b uint32) uint32 {
+	for b != 0 {
+		t := b
+		b = a % b
+		a = t
+	}
+	return a
+}
+
+func newDropper(c DropConfig) *dropper {
+	w := NewRandomWRR()
+	gcdv := gcd(c.RequestsPerMillion, million)
+	// Return true for RequestPerMillion, false for the rest.
+	w.Add(true, int64(c.RequestsPerMillion/gcdv))
+	w.Add(false, int64((million-c.RequestsPerMillion)/gcdv))
+
+	return &dropper{
+		category: c.Category,
+		w:        w,
+	}
+}
+
+func (d *dropper) drop() (ret bool) {
+	return d.w.Next().(bool)
+}
+
+// loadReporter wraps the methods from the loadStore that are used here.
+type loadReporter interface {
+	CallStarted(locality clients.Locality)
+	CallFinished(locality clients.Locality, err error)
+	CallServerLoad(locality clients.Locality, name string, val float64)
+	CallDropped(category string)
+}
+
+// Picker implements RPC drop, circuit breaking drop and load reporting.
+type picker struct {
+	drops           []*dropper
+	s               balancer.State
+	loadStore       loadReporter
+	counter         *xdsclient.ClusterRequestsCounter
+	countMax        uint32
+	telemetryLabels map[string]string
+	clusterName     string
+	metrics         *xdsresource.LRSReportEndpointMetricsConfig
+}
+
+func (d *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
+	// Unconditionally update labels if present, even dropped or queued RPC's can
+	// use these labels.
+	stats.UpdateLabels(info.Ctx, d.telemetryLabels)
+
+	// Don't drop unless the inner picker is READY. Similar to
+	// https://github.com/grpc/grpc-go/issues/2622.
+	if d.s.ConnectivityState == connectivity.Ready {
+		// Check if this RPC should be dropped by category.
+		for _, dp := range d.drops {
+			if dp.drop() {
+				if d.loadStore != nil {
+					d.loadStore.CallDropped(dp.category)
+				}
+				return balancer.PickResult{}, status.Errorf(codes.Unavailable, "RPC is dropped")
+			}
+		}
+	}
+
+	// Check if this RPC should be dropped by circuit breaking.
+	if d.counter != nil {
+		if err := d.counter.StartRequest(d.countMax); err != nil {
+			// Drops by circuit breaking are reported with empty category. They
+			// will be reported only in total drops, but not in per category.
+			if d.loadStore != nil {
+				d.loadStore.CallDropped("")
+			}
+			return balancer.PickResult{}, status.Error(codes.Unavailable, err.Error())
+		}
+	}
+
+	var lID clients.Locality
+	pr, err := d.s.Picker.Pick(info)
+	if scw, ok := pr.SubConn.(*scWrapper); ok {
+		// This OK check also covers the case err!=nil, because SubConn will be
+		// nil.
+		pr.SubConn = scw.SubConn
+		// If locality ID isn't found in the wrapper, an empty locality ID will
+		// be used.
+		lID = scw.localityID
+
+		if scw.hostname != "" && autoHostRewriteEnabled(info.Ctx) {
+			if pr.Metadata == nil {
+				pr.Metadata = metadata.Pairs(":authority", scw.hostname)
+			} else {
+				pr.Metadata.Set(":authority", scw.hostname)
+			}
+		}
+	}
+
+	if err != nil {
+		if d.counter != nil {
+			// Release one request count if this pick fails.
+			d.counter.EndRequest()
+		}
+		return pr, err
+	}
+
+	stats.UpdateLabels(
+		info.Ctx,
+		map[string]string{
+			"grpc.lb.locality":        xdsinternal.LocalityString(lID),
+			"grpc.lb.backend_service": d.clusterName,
+		},
+	)
+
+	if d.loadStore != nil {
+		locality := clients.Locality{Region: lID.Region, Zone: lID.Zone, SubZone: lID.SubZone}
+		d.loadStore.CallStarted(locality)
+		oldDone := pr.Done
+		pr.Done = func(info balancer.DoneInfo) {
+			if oldDone != nil {
+				oldDone(info)
+			}
+			d.loadStore.CallFinished(locality, info.Err)
+
+			load, ok := info.ServerLoad.(*v3orcapb.OrcaLoadReport)
+			if !ok || load == nil {
+				return
+			}
+
+			if envconfig.XDSORCAToLRSPropEnabled {
+				if d.metrics != nil {
+					if d.metrics.CPUUtilization {
+						d.loadStore.CallServerLoad(locality, "cpu_utilization", load.CpuUtilization)
+					}
+					if d.metrics.MemUtilization {
+						d.loadStore.CallServerLoad(locality, "mem_utilization", load.MemUtilization)
+					}
+					if d.metrics.ApplicationUtilization {
+						d.loadStore.CallServerLoad(locality, "application_utilization", load.ApplicationUtilization)
+					}
+					for n, c := range load.NamedMetrics {
+						_, ok := d.metrics.NamedMetrics[n]
+						if d.metrics.NamedMetricsAll || ok {
+							d.loadStore.CallServerLoad(locality, "named_metrics."+n, c)
+						}
+					}
+				}
+			} else {
+				for n, c := range load.NamedMetrics {
+					d.loadStore.CallServerLoad(locality, n, c)
+				}
+			}
+		}
+	}
+
+	if d.counter != nil {
+		// Update Done() so that when the RPC finishes, the request count will
+		// be released.
+		oldDone := pr.Done
+		pr.Done = func(doneInfo balancer.DoneInfo) {
+			d.counter.EndRequest()
+			if oldDone != nil {
+				oldDone(doneInfo)
+			}
+		}
+	}
+
+	return pr, err
+}
+
+// autoHostRewriteKey is the context key used to store the value of
+// route's autoHostRewrite in the RPC context.
+type autoHostRewriteKey struct{}
+
+// autoHostRewriteEnabled retrieves the autoHostRewrite value from the provided context.
+func autoHostRewriteEnabled(ctx context.Context) bool {
+	v, _ := ctx.Value(autoHostRewriteKey{}).(bool)
+	return v
+}
+
+// AutoHostRewriteEnabledForTesting returns the value of autoHostRewrite field;
+// to be used for testing only.
+func AutoHostRewriteEnabledForTesting(ctx context.Context) bool {
+	return autoHostRewriteEnabled(ctx)
+}
+
+// EnableAutoHostRewrite adds the autoHostRewrite value to the context for
+// the xds_cluster_impl LB policy to pick.
+func EnableAutoHostRewrite(ctx context.Context) context.Context {
+	return context.WithValue(ctx, autoHostRewriteKey{}, true)
+}