xds: Fix WeakReference bug in SharedCallCounterMap

xds/src/main/java/io/grpc/xds/SharedCallCounterMap.java

@@ -52,19 +52,28 @@ static SharedCallCounterMap getInstance() {
 
   @Override
   public synchronized AtomicLong getOrCreate(String cluster, @Nullable String edsServiceName) {
+    AtomicLong counter = null;
     Map<String, CounterReference> clusterCounters = counters.get(cluster);
+    if (clusterCounters != null) {
+      CounterReference ref = clusterCounters.get(edsServiceName);
+      if (ref != null) {
+        counter = ref.get();
+      }
+    }
+    // In the case of ref.get() == null, must call cleanQueue() prior to creating a new map entry,
+    // otherwise the new entry will be deleted by any later cleanQueue().
+    cleanQueue();
+    if (counter != null) {
+      return counter;
+    }
+    clusterCounters = counters.get(cluster);
     if (clusterCounters == null) {
       clusterCounters = new HashMap<>();
       counters.put(cluster, clusterCounters);
     }
-    CounterReference ref = clusterCounters.get(edsServiceName);
-    AtomicLong counter;
-    if (ref == null || (counter = ref.get()) == null) {
-      counter = new AtomicLong();
-      ref = new CounterReference(counter, refQueue, cluster, edsServiceName);
-      clusterCounters.put(edsServiceName, ref);
-    }
-    cleanQueue();
+    counter = new AtomicLong();
+    CounterReference ref = new CounterReference(counter, refQueue, cluster, edsServiceName);
+    clusterCounters.put(edsServiceName, ref);
     return counter;
   }
 

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();
+  }
 }