From 9c0fc765bf440d3646aa1037988ddb25e3921063 Mon Sep 17 00:00:00 2001
From: Santiago Pericasgeertsen <santiago.pericasgeertsen@oracle.com>
Date: Mon, 6 Nov 2023 11:16:18 -0500
Subject: [PATCH] Properly removes a channel from the cache. If a channel set
 is empty, it must be removed from the map. See issue #7821.

---
 .../WebClientRequestBuilderImpl.java          | 19 ++++++---
 .../WebClientRequestBuilderImplTest.java      | 39 +++++++++++++++++++
 2 files changed, 53 insertions(+), 5 deletions(-)
 create mode 100644 webclient/webclient/src/test/java/io/helidon/webclient/WebClientRequestBuilderImplTest.java

diff --git a/webclient/webclient/src/main/java/io/helidon/webclient/WebClientRequestBuilderImpl.java b/webclient/webclient/src/main/java/io/helidon/webclient/WebClientRequestBuilderImpl.java
index 5571b19e559..4d433b949b6 100644
--- a/webclient/webclient/src/main/java/io/helidon/webclient/WebClientRequestBuilderImpl.java
+++ b/webclient/webclient/src/main/java/io/helidon/webclient/WebClientRequestBuilderImpl.java
@@ -89,10 +89,10 @@ class WebClientRequestBuilderImpl implements WebClientRequestBuilder {
 
     private static final Logger LOGGER = Logger.getLogger(WebClientRequestBuilderImpl.class.getName());
 
-    private static final Map<ConnectionIdent, Set<ChannelRecord>> CHANNEL_CACHE = new ConcurrentHashMap<>();
     private static final List<DataPropagationProvider> PROPAGATION_PROVIDERS = HelidonServiceLoader
             .builder(ServiceLoader.load(DataPropagationProvider.class)).build().asList();
 
+    static final Map<ConnectionIdent, Set<ChannelRecord>> CHANNEL_CACHE = new ConcurrentHashMap<>();
     static final AttributeKey<WebClientRequestImpl> REQUEST = AttributeKey.valueOf("request");
     static final AttributeKey<CompletableFuture<WebClientServiceResponse>> RECEIVED = AttributeKey.valueOf("received");
     static final AttributeKey<CompletableFuture<WebClientServiceResponse>> COMPLETED = AttributeKey.valueOf("completed");
@@ -242,7 +242,16 @@ static void removeChannelFromCache(ConnectionIdent key, Channel channel) {
             LOGGER.finest(() -> "Removing from channel cache. Connection ident ->  " + key
                     + ", channel -> " + channel.hashCode());
         }
-        CHANNEL_CACHE.get(key).remove(new ChannelRecord(channel));
+        Set<ChannelRecord> channelSet = CHANNEL_CACHE.get(key);
+        if (channelSet != null) {
+            // remove entry from set
+            channelSet.remove(new ChannelRecord(channel));
+
+            // remove set from map if empty
+            if (channelSet.isEmpty()) {
+                CHANNEL_CACHE.remove(key);
+            }
+        }
     }
 
     @Override
@@ -958,17 +967,17 @@ HttpRequest.Path createSubpath(String path, String rawPath,
         }
     }
 
-    private static class ChannelRecord {
+    static class ChannelRecord {
 
         private final ChannelFuture channelFuture;
         private final Channel channel;
 
-        private ChannelRecord(ChannelFuture channelFuture) {
+        ChannelRecord(ChannelFuture channelFuture) {
             this.channelFuture = channelFuture;
             this.channel = channelFuture.channel();
         }
 
-        private ChannelRecord(Channel channel) {
+        ChannelRecord(Channel channel) {
             this.channelFuture = null;
             this.channel = channel;
         }
diff --git a/webclient/webclient/src/test/java/io/helidon/webclient/WebClientRequestBuilderImplTest.java b/webclient/webclient/src/test/java/io/helidon/webclient/WebClientRequestBuilderImplTest.java
new file mode 100644
index 00000000000..071c899874e
--- /dev/null
+++ b/webclient/webclient/src/test/java/io/helidon/webclient/WebClientRequestBuilderImplTest.java
@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) 2023 Oracle and/or its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.helidon.webclient;
+
+import java.util.HashSet;
+
+import io.netty.channel.Channel;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.nullValue;
+
+class WebClientRequestBuilderImplTest {
+
+    @Test
+    void testChannelRemovalFromCache() {
+        var conn = Mockito.mock(WebClientRequestBuilderImpl.ConnectionIdent.class);
+        var channel = Mockito.mock(Channel.class);
+        HashSet<WebClientRequestBuilderImpl.ChannelRecord> set = new HashSet<>();
+        set.add(new WebClientRequestBuilderImpl.ChannelRecord(channel));
+        WebClientRequestBuilderImpl.CHANNEL_CACHE.put(conn, set);
+        WebClientRequestBuilderImpl.removeChannelFromCache(conn, channel);
+        assertThat(WebClientRequestBuilderImpl.CHANNEL_CACHE.get(conn), nullValue());
+    }
+}