xds: Fix WeakReference bug in SharedCallCounterMap (grpc#8466)

dapengzhang0 · web-flow · commit 07747c59a2f2 · 2021-09-02T10:25:15.000-07:00
Fixes grpc#8397. grpc#8397 is caused by mistakenly clearing up a map entry right after the entry is recreated after gc. Reproduced in regression test.
diff --git a/xds/src/main/java/io/grpc/xds/SharedCallCounterMap.java b/xds/src/main/java/io/grpc/xds/SharedCallCounterMap.java
@@ -58,8 +58,14 @@ public synchronized AtomicLong getOrCreate(String cluster, @Nullable String edsS
       counters.put(cluster, clusterCounters);
     }
     CounterReference ref = clusterCounters.get(edsServiceName);
-    AtomicLong counter;
-    if (ref == null || (counter = ref.get()) == null) {
+    AtomicLong counter = null;
+    if (ref != null) {
+      counter = ref.get();
+      if (counter == null) {
+        ref.enqueue();
+      }
+    }
+    if (counter == null) {
       counter = new AtomicLong();
       ref = new CounterReference(counter, refQueue, cluster, edsServiceName);
       clusterCounters.put(edsServiceName, ref);
@@ -73,6 +79,9 @@ void cleanQueue() {
     CounterReference ref;
     while ((ref = (CounterReference) refQueue.poll()) != null) {
       Map<String, CounterReference> clusterCounter = counters.get(ref.cluster);
+      if (clusterCounter.get(ref.edsServiceName) != ref) {
+        continue;
+      }
       clusterCounter.remove(ref.edsServiceName);
       if (clusterCounter.isEmpty()) {
         counters.remove(ref.cluster);
diff --git a/xds/src/test/java/io/grpc/xds/SharedCallCounterMapTest.java b/xds/src/test/java/io/grpc/xds/SharedCallCounterMapTest.java
@@ -62,4 +62,22 @@ public boolean isDone() {
     map.cleanQueue();
     assertThat(counters).isEmpty();
   }
+
+  @Test
+  public void gcAndRecreate() {
+    @SuppressWarnings("UnusedVariable") // assign to null for GC only
+    AtomicLong counter = map.getOrCreate(CLUSTER, EDS_SERVICE_NAME);
+    final CounterReference ref = counters.get(CLUSTER).get(EDS_SERVICE_NAME);
+    assertThat(counter.get()).isEqualTo(0);
+    counter = null;
+    GcFinalization.awaitDone(new FinalizationPredicate() {
+      @Override
+      public boolean isDone() {
+        return ref.isEnqueued();
+      }
+    });
+    map.getOrCreate(CLUSTER, EDS_SERVICE_NAME);
+    assertThat(counters.get(CLUSTER)).isNotNull();
+    assertThat(counters.get(CLUSTER).get(EDS_SERVICE_NAME)).isNotNull();
+  }
 }
