Add recording point for endpoint weight not yet usable and add metrics tests

Author: zasweq

balancer/weightedroundrobin/balancer.go

@@ -415,7 +415,7 @@ func (p *picker) inc() uint32 {
 }
 
 func (p *picker) regenerateScheduler() {
-	s := p.newScheduler()
+	s := p.newScheduler(true)
 	atomic.StorePointer(&p.scheduler, unsafe.Pointer(&s))
 }
 
@@ -558,14 +558,17 @@ func (w *weightedSubConn) updateConnectivityState(cs connectivity.State) connect
 		w.SubConn.Connect()
 	case connectivity.Ready:
 		// If we transition back to READY state, reset nonEmptySince so that we
-		// apply the blackout period after we start receiving load data.  Note
-		// that we cannot guarantee that we will never receive lingering
-		// callbacks for backend metric reports from the previous connection
-		// after the new connection has been established, but they should be
-		// masked by new backend metric reports from the new connection by the
-		// time the blackout period ends.
+		// apply the blackout period after we start receiving load data. Also
+		// reset lastUpdated to trigger endpoint weight not yet usable in the
+		// case endpoint gets asked what weight it is before receiving a new
+		// load report. Note that we cannot guarantee that we will never receive
+		// lingering callbacks for backend metric reports from the previous
+		// connection after the new connection has been established, but they
+		// should be masked by new backend metric reports from the new
+		// connection by the time the blackout period ends.
 		w.mu.Lock()
 		w.nonEmptySince = time.Time{}
+		w.lastUpdated = time.Time{}
 		w.mu.Unlock()
 	case connectivity.Shutdown:
 		if w.stopORCAListener != nil {
@@ -603,6 +606,13 @@ func (w *weightedSubConn) weight(now time.Time, weightExpirationPeriod, blackout
 		}()
 	}
 
+	// The SubConn has not received a load report (i.e. just turned READY with
+	// no load report).
+	if w.lastUpdated == (time.Time{}) {
+		endpointWeightNotYetUsableMetric.Record(w.metricsRecorder, 1, w.target, w.locality)
+		return 0
+	}
+
 	// If the most recent update was longer ago than the expiration period,
 	// reset nonEmptySince so that we apply the blackout period again if we
 	// start getting data again in the future, and return 0.

balancer/weightedroundrobin/balancer_test.go

@@ -32,6 +32,7 @@ import (
 	"google.golang.org/grpc/internal/grpctest"
 	"google.golang.org/grpc/internal/stubserver"
 	"google.golang.org/grpc/internal/testutils/roundrobin"
+	"google.golang.org/grpc/internal/testutils/stats"
 	"google.golang.org/grpc/orca"
 	"google.golang.org/grpc/peer"
 	"google.golang.org/grpc/resolver"
@@ -81,6 +82,14 @@ var (
 		WeightUpdatePeriod:      stringp(".050s"),
 		ErrorUtilizationPenalty: float64p(0),
 	}
+	testMetricsConfig = iwrr.LBConfig{
+		EnableOOBLoadReport:     boolp(false),
+		OOBReportingPeriod:      stringp("0.005s"),
+		BlackoutPeriod:          stringp("0s"),
+		WeightExpirationPeriod:  stringp("60s"),
+		WeightUpdatePeriod:      stringp(".050s"),
+		ErrorUtilizationPenalty: float64p(0),
+	}
 )
 
 type testServer struct {
@@ -196,6 +205,43 @@ func (s) TestBalancer_OneAddress(t *testing.T) {
 	}
 }
 
+// TestWRRMetricsBasic tests metrics emitted from the WRR balancer. It
+// configures a weighted round robin balancer as the top level balancer of a
+// ClientConn, and configures a fake stats handler on the ClientConn to receive
+// metrics. It verifies stats emitted from the Weighted Round Robin Balancer on
+// balancer startup case which triggers the first picker and scheduler update
+// before any load reports are received.
+//
+// Note that this test and others, metrics emission asssertions are a snapshot
+// of the most recently emitted metrics. This is due to the nondeterminism of
+// scheduler updates with respect to test bodies, so the assertions made are
+// from the most recently synced state of the system (picker/scheduler) from the
+// test body.
+func (s) TestWRRMetricsBasic(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	srv := startServer(t, reportCall)
+	sc := svcConfig(t, testMetricsConfig)
+
+	mr := stats.NewTestMetricsRecorder(t, []string{"grpc.lb.wrr.rr_fallback", "grpc.lb.wrr.endpoint_weight_not_yet_usable", "grpc.lb.wrr.endpoint_weight_stale", "grpc.lb.wrr.endpoint_weights"})
+	if err := srv.StartClient(grpc.WithDefaultServiceConfig(sc), grpc.WithStatsHandler(mr)); err != nil {
+		t.Fatalf("Error starting client: %v", err)
+	}
+	srv.callMetrics.SetQPS(float64(1))
+
+	if _, err := srv.Client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+		t.Fatalf("Error from EmptyCall: %v", err)
+	}
+
+	mr.AssertDataForMetric("grpc.lb.wrr.rr_fallback", 1)           // Falls back because only one SubConn.
+	mr.AssertDataForMetric("grpc.lb.wrr.endpoint_weight_stale", 0) // The endpoint weight has not expired so this is 0 (never emitted).
+	mr.AssertDataForMetric("grpc.lb.wrr.endpoint_weight_not_yet_usable", 1)
+	// Unusable, so no endpoint weight. Due to only one SubConn, this will never
+	// update the weight. Thus, this will stay 0.
+	mr.AssertDataForMetric("grpc.lb.wrr.endpoint_weights", 0)
+}
+
 // Tests two addresses with ORCA reporting disabled (should fall back to pure
 // RR).
 func (s) TestBalancer_TwoAddresses_ReportingDisabled(t *testing.T) {
@@ -509,7 +555,8 @@ func (s) TestBalancer_TwoAddresses_BlackoutPeriod(t *testing.T) {
 		cfg := oobConfig
 		cfg.BlackoutPeriod = tc.blackoutPeriodCfg
 		sc := svcConfig(t, cfg)
-		if err := srv1.StartClient(grpc.WithDefaultServiceConfig(sc)); err != nil {
+		mr := stats.NewTestMetricsRecorder(t, []string{"grpc.lb.wrr.rr_fallback", "grpc.lb.wrr.endpoint_weight_not_yet_usable", "grpc.lb.wrr.endpoint_weight_stale", "grpc.lb.wrr.endpoint_weights"})
+		if err := srv1.StartClient(grpc.WithDefaultServiceConfig(sc), grpc.WithStatsHandler(mr)); err != nil {
 			t.Fatalf("Error starting client: %v", err)
 		}
 		addrs := []resolver.Address{{Addr: srv1.Address}, {Addr: srv2.Address}}
@@ -536,12 +583,20 @@ func (s) TestBalancer_TwoAddresses_BlackoutPeriod(t *testing.T) {
 		// Wait for the weight update period to allow the new weights to be processed.
 		time.Sleep(weightUpdatePeriod)
 		checkWeights(ctx, t, srvWeight{srv1, 1}, srvWeight{srv2, 10})
+
+		mr.AssertDataForMetric("grpc.lb.wrr.rr_fallback", 1)
+		mr.AssertDataForMetric("grpc.lb.wrr.endpoint_weight_stale", 0)
+		mr.AssertDataForMetric("grpc.lb.wrr.endpoint_weight_not_yet_usable", 1)
+		// Either 10 or 100, dependent on whatever ordering SubConns are
+		// processed, which is nondeterministic.
+		mr.AssertEitherDataForMetric("grpc.lb.wrr.endpoint_weights", 10, 100)
 	}
 }
 
 // Tests that the weight expiration period causes backends to use 0 as their
 // weight (meaning to use the average weight) once the expiration period
-// elapses.
+// elapses. After the weight expires, the expected metrics to be emitted from
+// WRR are also configured.
 func (s) TestBalancer_TwoAddresses_WeightExpiration(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
 	defer cancel()
@@ -577,7 +632,8 @@ func (s) TestBalancer_TwoAddresses_WeightExpiration(t *testing.T) {
 	cfg := oobConfig
 	cfg.OOBReportingPeriod = stringp("60s")
 	sc := svcConfig(t, cfg)
-	if err := srv1.StartClient(grpc.WithDefaultServiceConfig(sc)); err != nil {
+	mr := stats.NewTestMetricsRecorder(t, []string{"grpc.lb.wrr.rr_fallback", "grpc.lb.wrr.endpoint_weight_not_yet_usable", "grpc.lb.wrr.endpoint_weight_stale", "grpc.lb.wrr.endpoint_weights"})
+	if err := srv1.StartClient(grpc.WithDefaultServiceConfig(sc), grpc.WithStatsHandler(mr)); err != nil {
 		t.Fatalf("Error starting client: %v", err)
 	}
 	addrs := []resolver.Address{{Addr: srv1.Address}, {Addr: srv2.Address}}
@@ -605,6 +661,11 @@ func (s) TestBalancer_TwoAddresses_WeightExpiration(t *testing.T) {
 	// Wait for the weight expiration period so the weights have expired.
 	time.Sleep(weightUpdatePeriod)
 	checkWeights(ctx, t, srvWeight{srv1, 1}, srvWeight{srv2, 1})
+
+	mr.AssertDataForMetric("grpc.lb.wrr.rr_fallback", 1)
+	mr.AssertDataForMetric("grpc.lb.wrr.endpoint_weight_stale", 1)
+	mr.AssertDataForMetric("grpc.lb.wrr.endpoint_weight_not_yet_usable", 1)
+	mr.AssertDataForMetric("grpc.lb.wrr.endpoint_weights", 0)
 }
 
 // Tests logic surrounding subchannel management.

balancer/weightedroundrobin/scheduler.go

@@ -31,14 +31,16 @@ type scheduler interface {
 // len(scWeights)-1 are zero or there is only a single subconn, otherwise it
 // will return an Earliest Deadline First (EDF) scheduler implementation that
 // selects the subchannels according to their weights.
-func (p *picker) newScheduler() scheduler {
-	scWeights := p.scWeights(true)
+func (p *picker) newScheduler(recordMetrics bool) scheduler {
+	scWeights := p.scWeights(recordMetrics)
 	n := len(scWeights)
 	if n == 0 {
 		return nil
 	}
 	if n == 1 {
-		rrFallbackMetric.Record(p.metricsRecorder, 1, p.target, p.locality)
+		if recordMetrics {
+			rrFallbackMetric.Record(p.metricsRecorder, 1, p.target, p.locality)
+		}
 		return &rrScheduler{numSCs: 1, inc: p.inc}
 	}
 	sum := float64(0)
@@ -55,7 +57,9 @@ func (p *picker) newScheduler() scheduler {
 	}
 
 	if numZero >= n-1 {
-		rrFallbackMetric.Record(p.metricsRecorder, 1, p.target, p.locality)
+		if recordMetrics {
+			rrFallbackMetric.Record(p.metricsRecorder, 1, p.target, p.locality)
+		}
 		return &rrScheduler{numSCs: uint32(n), inc: p.inc}
 	}
 	unscaledMean := sum / float64(n-numZero)

internal/testutils/stats/test_metrics_recorder.go

@@ -21,7 +21,9 @@ package stats
 
 import (
 	"context"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/google/go-cmp/cmp"
 	estats "google.golang.org/grpc/experimental/stats"
@@ -41,18 +43,65 @@ type TestMetricsRecorder struct {
 	intHistoCh   *testutils.Channel
 	floatHistoCh *testutils.Channel
 	intGaugeCh   *testutils.Channel
+
+	// The most recent update for each metric name.
+	mu   sync.Mutex
+	data map[estats.Metric]float64
 }
 
 func NewTestMetricsRecorder(t *testing.T, metrics []string) *TestMetricsRecorder {
-	return &TestMetricsRecorder{
+	tmr := &TestMetricsRecorder{
 		t: t,
 
-		intCountCh:   testutils.NewChannelWithSize(10),
-		floatCountCh: testutils.NewChannelWithSize(10),
-		intHistoCh:   testutils.NewChannelWithSize(10),
-		floatHistoCh: testutils.NewChannelWithSize(10),
-		intGaugeCh:   testutils.NewChannelWithSize(10),
+		intCountCh:   testutils.NewChannelWithSize(1000),
+		floatCountCh: testutils.NewChannelWithSize(1000),
+		intHistoCh:   testutils.NewChannelWithSize(1000),
+		floatHistoCh: testutils.NewChannelWithSize(1000),
+		intGaugeCh:   testutils.NewChannelWithSize(1000),
+
+		data: make(map[estats.Metric]float64),
+	}
+
+	for _, metric := range metrics {
+		tmr.data[estats.Metric(metric)] = 0
+	}
+
+	return tmr
+}
+
+// AssertDataForMetric asserts data is present for metric. The zero value in the
+// check is equivalent to unset.
+func (r *TestMetricsRecorder) AssertDataForMetric(metricName string, wantVal float64) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if r.data[estats.Metric(metricName)] != wantVal {
+		r.t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", metricName, r.data[estats.Metric(metricName)], wantVal)
+	}
+}
+
+// AssertEitherDataForMetric asserts either data point is present for metric.
+// The zero value in the check is equivalent to unset.
+
+func (r *TestMetricsRecorder) AssertEitherDataForMetric(metricName string, wantVal1 float64, wantVal2 float64) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if r.data[estats.Metric(metricName)] != wantVal1 && r.data[estats.Metric(metricName)] != wantVal2 {
+		r.t.Fatalf("Unexpected data for metric %v, got: %v, want: %v or %v", metricName, r.data[estats.Metric(metricName)], wantVal1, wantVal2)
+	}
+}
+
+// PollForDataForMetric polls the metric data for the want. Fails if context
+// provided expires before data for metric is found.
+func (r *TestMetricsRecorder) PollForDataForMetric(ctx context.Context, metricName string, wantVal float64) {
+	for ; ctx.Err() == nil; <-time.After(time.Millisecond) {
+		r.mu.Lock()
+		if r.data[estats.Metric(metricName)] == wantVal {
+			r.mu.Unlock()
+			break
+		}
+		r.mu.Unlock()
 	}
+	r.t.Fatalf("Timeout waiting for data %v for metric %v", wantVal, metricName)
 }
 
 type MetricsData struct {
@@ -85,6 +134,10 @@ func (r *TestMetricsRecorder) RecordInt64Count(handle *estats.Int64CountHandle,
 		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
 		LabelVals: labels,
 	})
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.data[handle.Name] = float64(incr)
 }
 
 func (r *TestMetricsRecorder) WaitForFloat64Count(ctx context.Context, metricsDataWant MetricsData) {
@@ -105,6 +158,10 @@ func (r *TestMetricsRecorder) RecordFloat64Count(handle *estats.Float64CountHand
 		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
 		LabelVals: labels,
 	})
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.data[handle.Name] = incr
 }
 
 func (r *TestMetricsRecorder) WaitForInt64Histo(ctx context.Context, metricsDataWant MetricsData) {
@@ -125,6 +182,10 @@ func (r *TestMetricsRecorder) RecordInt64Histo(handle *estats.Int64HistoHandle,
 		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
 		LabelVals: labels,
 	})
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.data[handle.Name] = float64(incr)
 }
 
 func (r *TestMetricsRecorder) WaitForFloat64Histo(ctx context.Context, metricsDataWant MetricsData) {
@@ -145,6 +206,10 @@ func (r *TestMetricsRecorder) RecordFloat64Histo(handle *estats.Float64HistoHand
 		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
 		LabelVals: labels,
 	})
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.data[handle.Name] = incr
 }
 
 func (r *TestMetricsRecorder) WaitForInt64Gauge(ctx context.Context, metricsDataWant MetricsData) {
@@ -165,6 +230,10 @@ func (r *TestMetricsRecorder) RecordInt64Gauge(handle *estats.Int64GaugeHandle,
 		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
 		LabelVals: labels,
 	})
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.data[handle.Name] = float64(incr)
 }
 
 // To implement a stats.Handler, which allows it to be set as a dial option:

internal/testutils/xds/e2e/setup/setup.go

@@ -0,0 +1,61 @@
+/*
+ *
+ * Copyright 2024 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 setup
+
+import (
+	"testing"
+
+	"github.com/google/uuid"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/testutils/xds/e2e"
+	"google.golang.org/grpc/resolver"
+	_ "google.golang.org/grpc/xds" // Register the xds_resolver.
+)
+
+// ManagementServerAndResolver sets up an xDS management server, creates
+// bootstrap configuration pointing to that server and creates an xDS resolver
+// using that configuration.
+//
+// Registers a cleanup function on t to stop the management server.
+//
+// Returns the following:
+// - the xDS management server
+// - the node ID to use when talking to this management server
+// - bootstrap configuration to use (if creating an xDS-enabled gRPC server)
+// - xDS resolver builder (if creating an xDS-enabled gRPC client)
+func ManagementServerAndResolver(t *testing.T) (*e2e.ManagementServer, string, []byte, resolver.Builder) {
+	// Start an xDS management server.
+	xdsServer := e2e.StartManagementServer(t, e2e.ManagementServerOptions{AllowResourceSubset: true})
+
+	// Create bootstrap configuration pointing to the above management server.
+	nodeID := uuid.New().String()
+	bc := e2e.DefaultBootstrapContents(t, nodeID, xdsServer.Address)
+
+	// Create an xDS resolver with the above bootstrap configuration.
+	var r resolver.Builder
+	var err error
+	if newResolver := internal.NewXDSResolverWithConfigForTesting; newResolver != nil {
+		r, err = newResolver.(func([]byte) (resolver.Builder, error))(bc)
+		if err != nil {
+			t.Fatalf("Failed to create xDS resolver for testing: %v", err)
+		}
+	}
+
+	return xdsServer, nodeID, bc, r
+}