From 07747c59a2f21190b9fe07120d5d1a57eaaebcb0 Mon Sep 17 00:00:00 2001
From: ZHANG Dapeng <zdapeng@google.com>
Date: Thu, 2 Sep 2021 10:25:15 -0700
Subject: [PATCH] xds: Fix WeakReference bug in SharedCallCounterMap (#8466)

Fixes #8397.
#8397 is caused by mistakenly clearing up a map entry right after the entry is recreated after gc. Reproduced in regression test.
---
 .../java/io/grpc/xds/SharedCallCounterMap.java | 13 +++++++++++--
 .../io/grpc/xds/SharedCallCounterMapTest.java  | 18 ++++++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/xds/src/main/java/io/grpc/xds/SharedCallCounterMap.java b/xds/src/main/java/io/grpc/xds/SharedCallCounterMap.java
index 71cff0cedf3..7aa55c27429 100644
--- 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
index 9f2293d3c53..3051a021870 100644
--- 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();
+  }
 }