core,xds: Metrics recording in WRR LB

Adds the recording of the four metrics documented in:

https://github.com/grpc/proposal/blob/master/A78-grpc-metrics-wrr-pf-xds.md#weighted-round-robin-lb-policy

api/src/main/java/io/grpc/MetricInstrumentRegistry.java

@@ -16,6 +16,7 @@
 
 package io.grpc;
 
+import com.google.common.annotations.VisibleForTesting;
 import java.util.Collections;
 import java.util.List;
 import java.util.Set;
@@ -46,6 +47,16 @@ public static synchronized MetricInstrumentRegistry getDefaultRegistry() {
     return instance;
   }
 
+  /**
+   * Allows the registry to be reset from unit tests.
+   */
+  @VisibleForTesting
+  public static synchronized void reset() {
+    if (instance != null) {
+      instance = new MetricInstrumentRegistry();
+    }
+  }
+
   /**
    * Returns a list of registered metric instruments.
    */

api/src/testFixtures/java/io/grpc/FakeMetricRecorder.java

@@ -0,0 +1,101 @@
+/*
+ * Copyright 2024 The 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 io.grpc;
+
+import com.google.common.base.MoreObjects;
+import com.google.common.collect.Maps;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * A fake implementation of the {@link MetricRecorder} that collects all metrics in memory to allow
+ * tests to make assertions against the collected metrics.
+ */
+public class FakeMetricRecorder implements MetricRecorder {
+  public Map<String, MetricEntry<Double>> doubleCounterEntries = Maps.newHashMap();
+  public Map<String, MetricEntry<Long>> longCounterEntries = Maps.newHashMap();
+  public Map<String, MetricEntry<Double>> doubleHistogramCounterEntries = Maps.newHashMap();
+  public Map<String, MetricEntry<Long>> longHistogramCounterEntries = Maps.newHashMap();
+
+
+  @Override
+  public void recordDoubleCounter(DoubleCounterMetricInstrument metricInstrument, double value,
+      List<String> requiredLabelValues, List<String> optionalLabelValues) {
+    doubleCounterEntries.put(metricInstrument.getName(),
+        new MetricEntry<>(value, requiredLabelValues, optionalLabelValues));
+  }
+
+  @Override
+  public void recordLongCounter(LongCounterMetricInstrument metricInstrument, long value,
+      List<String> requiredLabelValues, List<String> optionalLabelValues) {
+    longCounterEntries.put(metricInstrument.getName(),
+        new MetricEntry<>(value, requiredLabelValues, optionalLabelValues));
+  }
+
+  @Override
+  public void recordDoubleHistogram(DoubleHistogramMetricInstrument metricInstrument,
+      double value, List<String> requiredLabelValues, List<String> optionalLabelValues) {
+    doubleHistogramCounterEntries.put(metricInstrument.getName(),
+        new MetricEntry<>(value, requiredLabelValues, optionalLabelValues));
+  }
+
+  @Override
+  public void recordLongHistogram(LongHistogramMetricInstrument metricInstrument, long value,
+      List<String> requiredLabelValues, List<String> optionalLabelValues) {
+    longHistogramCounterEntries.put(metricInstrument.getName(),
+        new MetricEntry<>(value, requiredLabelValues, optionalLabelValues));
+  }
+
+  public Long getLongCounterValue(String metricName) {
+    return longCounterEntries.get(metricName).value;
+  }
+
+  // Returns the last recorded double histogram value.
+  public Double getDoubleHistogramValue(String metricName) {
+    return doubleHistogramCounterEntries.get(metricName).value;
+  }
+
+  public boolean hasLongCounterValue(String metricName) {
+    return longCounterEntries.containsKey(metricName);
+  }
+
+  public void clear() {
+    doubleCounterEntries.clear();
+    longCounterEntries.clear();
+    doubleHistogramCounterEntries.clear();
+    longHistogramCounterEntries.clear();
+  }
+
+  public static class MetricEntry<T> {
+    public T value;
+    public List<String> requiredLabelValues;
+    public List<String> optionalLabelValues;
+
+    public MetricEntry(T value, List<String> requiredValues, List<String> optionalValues) {
+      this.value = value;
+      this.requiredLabelValues = requiredValues;
+      this.optionalLabelValues = optionalValues;
+    }
+
+    @Override
+    public String toString() {
+      return MoreObjects.toStringHelper(this).add("value", value)
+          .add("requiredLabelValues", requiredLabelValues)
+          .add("optionalLabelValues", optionalLabelValues).toString();
+    }
+  }
+}

core/src/main/java/io/grpc/internal/MetricLongCounter.java

@@ -0,0 +1,64 @@
+/*
+ * Copyright 2024 The 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 io.grpc.internal;
+
+import io.grpc.LongCounterMetricInstrument;
+import io.grpc.MetricRecorder;
+import java.util.List;
+
+/**
+ * Maintains a counter while also recording new values with a {@link MetricRecorder}.
+ */
+public class MetricLongCounter {
+
+  private final MetricRecorder recorder;
+  private final LongCounterMetricInstrument instrument;
+  private final LongCounter longCounter;
+
+  /**
+   * Creates a new {@code MetricLongCounter} for a given recorder and instrument.
+   */
+  public MetricLongCounter(MetricRecorder recorder, LongCounterMetricInstrument instrument) {
+    this.recorder = recorder;
+    this.instrument = instrument;
+    this.longCounter = LongCounterFactory.create();
+  }
+
+  /**
+   * Updates the counter with the given delta value and records the change with the
+   * {@link MetricRecorder}.
+   */
+  public void add(long delta, List<String> requiredLabelValues, List<String> optionalLabelValues) {
+    longCounter.add(delta);
+    recorder.recordLongCounter(instrument, longCounter.value(), requiredLabelValues,
+        optionalLabelValues);
+  }
+
+  /**
+   * Increments the counter by one and records the change with the {@link MetricRecorder}.
+   */
+  public void increment(List<String> requiredLabelValues, List<String> optionalLabelValues) {
+    add(1, requiredLabelValues, optionalLabelValues);
+  }
+
+  /**
+   * Returns the current value of the counter.
+   */
+  public long value() {
+    return longCounter.value();
+  }
+}

core/src/test/java/io/grpc/internal/MetricLongCounterTest.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2024 The 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 io.grpc.internal;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.verify;
+
+import com.google.common.collect.Lists;
+import io.grpc.LongCounterMetricInstrument;
+import io.grpc.MetricInstrumentRegistry;
+import io.grpc.MetricRecorder;
+import java.util.List;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+
+@RunWith(JUnit4.class)
+public class MetricLongCounterTest {
+
+  @Rule
+  public MockitoRule mockitoRule = MockitoJUnit.rule();
+
+  @Mock
+  MetricRecorder mockMetricRecorder;
+
+  @Test
+  public void incrementAndAdd() {
+    LongCounterMetricInstrument instrument = MetricInstrumentRegistry.getDefaultRegistry()
+        .registerLongCounter("long",
+            "description", "unit", Lists.newArrayList(), Lists.newArrayList(), true);
+    List<String> requiredLabelValues = Lists.newArrayList();
+    List<String> optionalLabelValues = Lists.newArrayList();
+    MetricLongCounter counter = new MetricLongCounter(mockMetricRecorder, instrument);
+
+    counter.increment(requiredLabelValues, optionalLabelValues);
+    verify(mockMetricRecorder).recordLongCounter(instrument, 1, requiredLabelValues,
+        optionalLabelValues);
+
+    counter.increment(requiredLabelValues, optionalLabelValues);
+    verify(mockMetricRecorder).recordLongCounter(instrument, 2, requiredLabelValues,
+        optionalLabelValues);
+
+    counter.add(2L, requiredLabelValues, optionalLabelValues);
+    verify(mockMetricRecorder).recordLongCounter(instrument, 4, requiredLabelValues,
+        optionalLabelValues);
+
+    counter.add(3L, requiredLabelValues, optionalLabelValues);
+    verify(mockMetricRecorder).recordLongCounter(instrument, 7, requiredLabelValues,
+        optionalLabelValues);
+
+    counter.add(5L, requiredLabelValues, optionalLabelValues);
+    verify(mockMetricRecorder).recordLongCounter(instrument, 12, requiredLabelValues,
+        optionalLabelValues);
+
+    assertThat(counter.value()).isEqualTo(12);
+  }
+}