From 96819f090d4edc748a9aacdf3517406f6aa47e56 Mon Sep 17 00:00:00 2001
From: emeroad <emeroad@gmail.com>
Date: Wed, 16 Oct 2024 16:01:37 +0900
Subject: [PATCH] [#11278] Fix atomicity of transportTerminated

---
 .../ConnectionCountServerTransportFilter.java | 22 +++++++++-
 ...nectionCountServerTransportFilterTest.java | 41 +++++++++++++++++++
 2 files changed, 61 insertions(+), 2 deletions(-)
 create mode 100644 grpc/src/test/java/com/navercorp/pinpoint/grpc/server/ConnectionCountServerTransportFilterTest.java

diff --git a/grpc/src/main/java/com/navercorp/pinpoint/grpc/server/ConnectionCountServerTransportFilter.java b/grpc/src/main/java/com/navercorp/pinpoint/grpc/server/ConnectionCountServerTransportFilter.java
index ca869c57f57a..20a8f8a66ccd 100644
--- a/grpc/src/main/java/com/navercorp/pinpoint/grpc/server/ConnectionCountServerTransportFilter.java
+++ b/grpc/src/main/java/com/navercorp/pinpoint/grpc/server/ConnectionCountServerTransportFilter.java
@@ -18,11 +18,17 @@
 
 import io.grpc.Attributes;
 import io.grpc.ServerTransportFilter;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 
 public class ConnectionCountServerTransportFilter extends ServerTransportFilter {
 
+    private static final Attributes.Key<AtomicInteger> TERMINATED = Attributes.Key.create("TransportTerminated");
+
+    private final Logger logger = LogManager.getLogger(this.getClass());
     private final AtomicLong currentConnection = new AtomicLong();
 
     public ConnectionCountServerTransportFilter() {
@@ -31,13 +37,25 @@ public ConnectionCountServerTransportFilter() {
 
     @Override
     public Attributes transportReady(Attributes transportAttrs) {
+        Attributes.Builder builder = transportAttrs.toBuilder();
+        builder.set(TERMINATED, new AtomicInteger());
+        Attributes attributes = builder.build();
+
         currentConnection.incrementAndGet();
-        return transportAttrs;
+        return attributes;
     }
 
     @Override
     public void transportTerminated(Attributes transportAttrs) {
-        currentConnection.decrementAndGet();
+        final AtomicInteger terminated = transportAttrs.get(TERMINATED);
+        if (terminated == null) {
+            return;
+        }
+        if (terminated.getAndIncrement() == 0) {
+            currentConnection.decrementAndGet();
+        } else {
+            logger.info("transportTerminated() already terminated attribute:{}", transportAttrs);
+        }
     }
 
     public long getCurrentConnection() {
diff --git a/grpc/src/test/java/com/navercorp/pinpoint/grpc/server/ConnectionCountServerTransportFilterTest.java b/grpc/src/test/java/com/navercorp/pinpoint/grpc/server/ConnectionCountServerTransportFilterTest.java
new file mode 100644
index 000000000000..1d53d0c595f2
--- /dev/null
+++ b/grpc/src/test/java/com/navercorp/pinpoint/grpc/server/ConnectionCountServerTransportFilterTest.java
@@ -0,0 +1,41 @@
+package com.navercorp.pinpoint.grpc.server;
+
+import io.grpc.Attributes;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+class ConnectionCountServerTransportFilterTest {
+
+    @Test
+    void transportReady() {
+        ConnectionCountServerTransportFilter filter = new ConnectionCountServerTransportFilter();
+        Attributes attributes = Attributes.newBuilder().build();
+        filter.transportReady(attributes);
+
+        Assertions.assertEquals(1, filter.getCurrentConnection());
+    }
+
+    @Test
+    void transportReadyAndTerminated() {
+        ConnectionCountServerTransportFilter filter = new ConnectionCountServerTransportFilter();
+        Attributes attributes = Attributes.newBuilder().build();
+        attributes = filter.transportReady(attributes);
+        filter.transportTerminated(attributes);
+
+        Assertions.assertEquals(0, filter.getCurrentConnection());
+    }
+
+    @Test
+    void transportTerminated_duplicate() {
+        ConnectionCountServerTransportFilter filter = new ConnectionCountServerTransportFilter();
+        Attributes attributes = Attributes.newBuilder().build();
+        attributes = filter.transportReady(attributes);
+        filter.transportTerminated(attributes);
+        filter.transportTerminated(attributes);
+        filter.transportTerminated(attributes);
+
+        Assertions.assertEquals(0, filter.getCurrentConnection());
+
+
+    }
+}
\ No newline at end of file