xds: Copy data in least request to avoid picker data race

ejona86 · ejona86 · commit f4cc166f18fb · 2024-02-16T13:44:25.000-08:00
In 0d39bf5 the ReadyPicker was changed holding List<Subchannel> to List<ChildLbState>, but ChildLbState mutates over time and is not synchronized. We want the picker to have a snapshot of the data, so copy the data from ChildLbState instead of using it directly. Unfortunately the tests depended on the ChildLbState a bit, so we need to save the EAG only to use it in tests. That's okay for now, but in the future we'll probably want to remove that unnecessary memory usage.
diff --git a/xds/src/main/java/io/grpc/xds/LeastRequestLoadBalancer.java b/xds/src/main/java/io/grpc/xds/LeastRequestLoadBalancer.java
@@ -32,6 +32,7 @@
 import io.grpc.ClientStreamTracer;
 import io.grpc.ClientStreamTracer.StreamInfo;
 import io.grpc.ConnectivityState;
+import io.grpc.EquivalentAddressGroup;
 import io.grpc.LoadBalancer;
 import io.grpc.LoadBalancerProvider;
 import io.grpc.Metadata;
@@ -153,28 +154,37 @@ private static AtomicInteger getInFlights(ChildLbState childLbState) {
 
   @VisibleForTesting
   static final class ReadyPicker extends SubchannelPicker {
-    private final List<ChildLbState> childLbStates; // non-empty
+    private final List<SubchannelPicker> childPickers; // non-empty
+    private final List<AtomicInteger> childInFlights; // 1:1 with childPickers
+    private final List<EquivalentAddressGroup> childEags; // 1:1 with childPickers
     private final int choiceCount;
     private final ThreadSafeRandom random;
     private final int hashCode;
 
     ReadyPicker(List<ChildLbState> childLbStates, int choiceCount, ThreadSafeRandom random) {
       checkArgument(!childLbStates.isEmpty(), "empty list");
-      this.childLbStates = childLbStates;
+      this.childPickers = new ArrayList<>(childLbStates.size());
+      this.childInFlights = new ArrayList<>(childLbStates.size());
+      this.childEags = new ArrayList<>(childLbStates.size());
+      for (ChildLbState state : childLbStates) {
+        childPickers.add(state.getCurrentPicker());
+        childInFlights.add(getInFlights(state));
+        childEags.add(state.getEag());
+      }
       this.choiceCount = choiceCount;
       this.random = checkNotNull(random, "random");
 
       int sum = 0;
-      for (ChildLbState child : childLbStates) {
+      for (SubchannelPicker child : childPickers) {
         sum += child.hashCode();
       }
       this.hashCode = sum ^ choiceCount;
     }
 
     @Override
     public PickResult pickSubchannel(PickSubchannelArgs args) {
-      final ChildLbState childLbState = nextChildToUse();
-      PickResult childResult = childLbState.getCurrentPicker().pickSubchannel(args);
+      int child = nextChildToUse();
+      PickResult childResult = childPickers.get(child).pickSubchannel(args);
 
       if (!childResult.getStatus().isOk() || childResult.getSubchannel() == null) {
         return childResult;
@@ -186,33 +196,38 @@ public PickResult pickSubchannel(PickSubchannelArgs args) {
       } else {
         // Wrap the subchannel
         OutstandingRequestsTracingFactory factory =
-            new OutstandingRequestsTracingFactory(getInFlights(childLbState));
+            new OutstandingRequestsTracingFactory(childInFlights.get(child));
         return PickResult.withSubchannel(childResult.getSubchannel(), factory);
       }
     }
 
     @Override
     public String toString() {
       return MoreObjects.toStringHelper(ReadyPicker.class)
-                        .add("list", childLbStates)
+                        .add("list", childPickers)
                         .add("choiceCount", choiceCount)
                         .toString();
     }
 
-    private ChildLbState nextChildToUse() {
-      ChildLbState candidate = childLbStates.get(random.nextInt(childLbStates.size()));
+    private int nextChildToUse() {
+      int candidate = random.nextInt(childPickers.size());
       for (int i = 0; i < choiceCount - 1; ++i) {
-        ChildLbState sampled = childLbStates.get(random.nextInt(childLbStates.size()));
-        if (getInFlights(sampled).get() < getInFlights(candidate).get()) {
+        int sampled = random.nextInt(childPickers.size());
+        if (childInFlights.get(sampled).get() < childInFlights.get(candidate).get()) {
           candidate = sampled;
         }
       }
       return candidate;
     }
 
     @VisibleForTesting
-    List<ChildLbState> getChildLbStates() {
-      return childLbStates;
+    List<SubchannelPicker> getChildPickers() {
+      return childPickers;
+    }
+
+    @VisibleForTesting
+    List<EquivalentAddressGroup> getChildEags() {
+      return childEags;
     }
 
     @Override
@@ -232,8 +247,8 @@ public boolean equals(Object o) {
       // the lists cannot contain duplicate children
       return hashCode == other.hashCode
           && choiceCount == other.choiceCount
-          && childLbStates.size() == other.childLbStates.size()
-          && new HashSet<>(childLbStates).containsAll(other.childLbStates);
+          && childPickers.size() == other.childPickers.size()
+          && new HashSet<>(childPickers).containsAll(other.childPickers);
     }
   }
 
diff --git a/xds/src/test/java/io/grpc/xds/LeastRequestLoadBalancerTest.java b/xds/src/test/java/io/grpc/xds/LeastRequestLoadBalancerTest.java
@@ -384,12 +384,12 @@ private String getStatusString(SubchannelPicker picker) {
     }
 
     if (picker instanceof ReadyPicker) {
-      List<ChildLbState> childLbStates = ((ReadyPicker)picker).getChildLbStates();
-      if (childLbStates == null || childLbStates.isEmpty()) {
+      List<SubchannelPicker> childPickers = ((ReadyPicker)picker).getChildPickers();
+      if (childPickers == null || childPickers.isEmpty()) {
         return "";
       }
 
-      picker = childLbStates.get(0).getCurrentPicker();
+      picker = childPickers.get(0);
     }
 
     Status status = picker.pickSubchannel(mockArgs).getStatus();
@@ -460,7 +460,8 @@ public void pickerLeastRequest() throws Exception {
 
     ReadyPicker picker = (ReadyPicker) pickerCaptor.getValue();
 
-    assertThat(picker.getChildLbStates()).containsExactlyElementsIn(childLbStates);
+    assertThat(picker.getChildEags())
+        .containsExactlyElementsIn(childLbStates.stream().map(ChildLbState::getEag).toArray());
 
     // Make random return 0, then 2 for the sample indexes.
     when(mockRandom.nextInt(childLbStates.size())).thenReturn(0, 2);
@@ -647,7 +648,7 @@ public void emptyAddresses() {
 
   private List<Subchannel> getList(SubchannelPicker picker) {
     if (picker instanceof ReadyPicker) {
-      return ((ReadyPicker) picker).getChildLbStates().stream()
+      return ((ReadyPicker) picker).getChildEags().stream()
           .map(this::getSubchannel)
           .collect(Collectors.toList());
     } else {
