nastra commented on code in PR #14161: URL: https://github.com/apache/iceberg/pull/14161#discussion_r2526891181
########## aws/src/main/java/org/apache/iceberg/aws/HttpClientCache.java: ########## @@ -0,0 +1,240 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.iceberg.aws; + +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Supplier; +import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting; +import org.apache.iceberg.relocated.com.google.common.collect.Maps; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import software.amazon.awssdk.http.ExecutableHttpRequest; +import software.amazon.awssdk.http.HttpExecuteRequest; +import software.amazon.awssdk.http.SdkHttpClient; + +/** + * A cache that manages the lifecycle of shared HTTP clients for AWS SDK v2 using reference + * counting. Package-private - only accessed via {@link BaseHttpClientConfigurations}. + */ +final class HttpClientCache { + private static final Logger LOG = LoggerFactory.getLogger(HttpClientCache.class); + + private final ConcurrentMap<String, ManagedHttpClient> clientMap; + + private static volatile HttpClientCache instance; + + static HttpClientCache getInstance() { + if (instance == null) { + synchronized (HttpClientCache.class) { + if (instance == null) { + instance = new HttpClientCache(); + } + } + } + return instance; + } + + private HttpClientCache() { + this.clientMap = Maps.newConcurrentMap(); + } + + /** + * Get or create a managed HTTP client for the given configuration. Each call increments the + * reference count for the client and returns a ref-counted wrapper. + * + * @param clientKey unique key identifying the client configuration + * @param clientFactory factory to create the HTTP client if not cached + * @return a ref-counted HTTP client wrapper + */ + SdkHttpClient getOrCreateClient(String clientKey, Supplier<SdkHttpClient> clientFactory) { + ManagedHttpClient managedClient = + clientMap.computeIfAbsent( + clientKey, + k -> { + LOG.debug("Creating new managed HTTP client for key: {}", k); + SdkHttpClient httpClient = clientFactory.get(); + return new ManagedHttpClient(httpClient, k); + }); + // Return the cached ref-counted wrapper + return managedClient.acquire(); + } + + /** + * Release a reference to the HTTP client. When the reference count reaches zero, the client is + * closed and removed from the cache. + * + * @param clientKey the key identifying the client to release + */ + void releaseClient(String clientKey) { + ManagedHttpClient managedClient = clientMap.get(clientKey); + if (managedClient != null) { + if (managedClient.release()) { + // Client was closed, remove from map + clientMap.remove(clientKey, managedClient); + } + } + } + + @VisibleForTesting + ConcurrentMap<String, ManagedHttpClient> clientMap() { + return clientMap; + } + + @VisibleForTesting + void shutdown() { + clientMap.values().forEach(ManagedHttpClient::close); + clientMap.clear(); + } + + /** + * Managed HTTP client wrapper that provides reference counting for lifecycle management. The HTTP + * client is closed when the reference count reaches zero. + */ + static class ManagedHttpClient { + private final SdkHttpClient httpClient; + private final String clientKey; + private final AtomicInteger refCount = new AtomicInteger(0); + private final AtomicBoolean closed = new AtomicBoolean(false); + private final WrappedSdkHttpClient wrapper; + + ManagedHttpClient(SdkHttpClient httpClient, String clientKey) { + this.httpClient = httpClient; + this.clientKey = clientKey; + this.wrapper = new WrappedSdkHttpClient(httpClient, clientKey); + LOG.debug("Created managed HTTP client: key={}", clientKey); + } + + /** + * Acquire a reference to the HTTP client, incrementing the reference count. + * + * @return the ref-counted wrapper client + * @throws IllegalStateException if the client has already been closed + */ + WrappedSdkHttpClient acquire() { + + int count = refCount.incrementAndGet(); + if (closed.get()) { + refCount.decrementAndGet(); + throw new IllegalStateException("Cannot acquire closed HTTP client: " + clientKey); + } + LOG.debug("Acquired HTTP client: key={}, refCount={}", clientKey, count); + return wrapper; + } + + /** + * Release a reference to the HTTP client, decrementing the reference count. If the count + * reaches zero, the client is closed. + * + * @return true if the client was closed, false otherwise + */ + boolean release() { + if (closed.get()) { + LOG.warn("Attempted to release already closed HTTP client: key={}", clientKey); + return false; + } + + int count = refCount.decrementAndGet(); + LOG.debug("Released HTTP client: key={}, refCount={}", clientKey, count); + if (count == 0) { + return close(); + } else if (count < 0) { + LOG.warn("HTTP client reference count went negative key={}, refCount={}", clientKey, count); + refCount.set(0); // Reset to prevent further corruption + } + return false; + } + + /** + * Close the HTTP client if not already closed. + * + * @return true if the client was closed by this call, false if already closed or if an error + * occurred + */ + @VisibleForTesting + boolean close() { + if (closed.compareAndSet(false, true)) { + LOG.debug("Closing HTTP client: key={}", clientKey); + try { + httpClient.close(); + return true; + } catch (Exception e) { + LOG.error("Failed to close HTTP client: key={}", clientKey, e); + return false; + } + } + return false; + } + + @VisibleForTesting + int refCount() { + return refCount.get(); + } + + @VisibleForTesting + boolean isClosed() { + return closed.get(); + } + } + + /** + * A delegating wrapper around {@link SdkHttpClient} that handles reference counting for lifecycle + * management. + * + * <p><strong>Lifecycle Contract:</strong> + * + * <ul> + * <li>Calling {@link #close()} decrements the reference count in the registry + * <li>The underlying HTTP client is only closed when the reference count reaches zero + * <li>Multiple wrappers for the same client configuration share the same lifecycle + * <li>After the underlying client is closed, operations on this wrapper may fail + * </ul> + */ + static class WrappedSdkHttpClient implements SdkHttpClient { Review Comment: I believe we can remove this additional layer of wrapping: ``` diff --git a/aws/src/main/java/org/apache/iceberg/aws/HttpClientCache.java b/aws/src/main/java/org/apache/iceberg/aws/HttpClientCache.java index 9a560611d5..1150ce394b 100644 --- a/aws/src/main/java/org/apache/iceberg/aws/HttpClientCache.java +++ b/aws/src/main/java/org/apache/iceberg/aws/HttpClientCache.java @@ -108,17 +108,15 @@ final class HttpClientCache { * Managed HTTP client wrapper that provides reference counting for lifecycle management. The HTTP * client is closed when the reference count reaches zero. */ - static class ManagedHttpClient { + static class ManagedHttpClient implements SdkHttpClient { private final SdkHttpClient httpClient; private final String clientKey; private final AtomicInteger refCount = new AtomicInteger(0); private final AtomicBoolean closed = new AtomicBoolean(false); - private final WrappedSdkHttpClient wrapper; ManagedHttpClient(SdkHttpClient httpClient, String clientKey) { this.httpClient = httpClient; this.clientKey = clientKey; - this.wrapper = new WrappedSdkHttpClient(httpClient, clientKey); LOG.debug("Created managed HTTP client: key={}", clientKey); } @@ -128,15 +126,14 @@ final class HttpClientCache { * @return the ref-counted wrapper client * @throws IllegalStateException if the client has already been closed */ - WrappedSdkHttpClient acquire() { - + ManagedHttpClient acquire() { int count = refCount.incrementAndGet(); if (closed.get()) { refCount.decrementAndGet(); throw new IllegalStateException("Cannot acquire closed HTTP client: " + clientKey); } LOG.debug("Acquired HTTP client: key={}, refCount={}", clientKey, count); - return wrapper; + return this; } /** @@ -154,7 +151,7 @@ final class HttpClientCache { int count = refCount.decrementAndGet(); LOG.debug("Released HTTP client: key={}, refCount={}", clientKey, count); if (count == 0) { - return close(); + return closeHttpClient(); } else if (count < 0) { LOG.warn("HTTP client reference count went negative key={}, refCount={}", clientKey, count); refCount.set(0); // Reset to prevent further corruption @@ -162,14 +159,18 @@ final class HttpClientCache { return false; } + @VisibleForTesting + SdkHttpClient httpClient() { + return httpClient; + } + /** * Close the HTTP client if not already closed. * * @return true if the client was closed by this call, false if already closed or if an error * occurred */ - @VisibleForTesting - boolean close() { + private boolean closeHttpClient() { if (closed.compareAndSet(false, true)) { LOG.debug("Closing HTTP client: key={}", clientKey); try { @@ -192,49 +193,20 @@ final class HttpClientCache { boolean isClosed() { return closed.get(); } - } - - /** - * A delegating wrapper around {@link SdkHttpClient} that handles reference counting for lifecycle - * management. - * - * <p><strong>Lifecycle Contract:</strong> - * - * <ul> - * <li>Calling {@link #close()} decrements the reference count in the registry - * <li>The underlying HTTP client is only closed when the reference count reaches zero - * <li>Multiple wrappers for the same client configuration share the same lifecycle - * <li>After the underlying client is closed, operations on this wrapper may fail - * </ul> - */ - static class WrappedSdkHttpClient implements SdkHttpClient { - private final SdkHttpClient delegate; - private final String clientKey; - - WrappedSdkHttpClient(SdkHttpClient delegate, String clientKey) { - this.delegate = delegate; - this.clientKey = clientKey; - } @Override public ExecutableHttpRequest prepareRequest(HttpExecuteRequest request) { - return delegate.prepareRequest(request); - } - - @Override - public void close() { - // Delegate close to the cache which manages ref counting - HttpClientCache.getInstance().releaseClient(clientKey); + return httpClient.prepareRequest(request); } @Override public String clientName() { - return delegate.clientName(); + return httpClient.clientName(); } - @VisibleForTesting - SdkHttpClient delegate() { - return delegate; + @Override + public void close() { + HttpClientCache.getInstance().releaseClient(clientKey); } } } diff --git a/aws/src/test/java/org/apache/iceberg/aws/TestHttpClientCache.java b/aws/src/test/java/org/apache/iceberg/aws/TestHttpClientCache.java index 00a13fd968..6c7164031e 100644 --- a/aws/src/test/java/org/apache/iceberg/aws/TestHttpClientCache.java +++ b/aws/src/test/java/org/apache/iceberg/aws/TestHttpClientCache.java @@ -27,7 +27,7 @@ import static org.mockito.Mockito.when; import java.util.concurrent.ConcurrentMap; import java.util.function.Supplier; -import org.apache.iceberg.aws.HttpClientCache.WrappedSdkHttpClient; +import org.apache.iceberg.aws.HttpClientCache.ManagedHttpClient; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.Mock; @@ -77,7 +77,7 @@ public class TestHttpClientCache { assertThat(client1).isSameAs(client2); // Verify reference count is 2 - HttpClientCache.ManagedHttpClient managedClient = cache.clientMap().get(cacheKey); + ManagedHttpClient managedClient = cache.clientMap().get(cacheKey); assertThat(managedClient.refCount()).isEqualTo(2); } @@ -100,12 +100,11 @@ public class TestHttpClientCache { SdkHttpClient mockClient = mock(SdkHttpClient.class); final String cacheKey = "test-key"; - HttpClientCache.ManagedHttpClient managedClient = - new HttpClientCache.ManagedHttpClient(mockClient, cacheKey); + ManagedHttpClient managedClient = new ManagedHttpClient(mockClient, cacheKey); // Acquire twice - WrappedSdkHttpClient client1 = managedClient.acquire(); - WrappedSdkHttpClient client2 = managedClient.acquire(); + ManagedHttpClient client1 = managedClient.acquire(); + ManagedHttpClient client2 = managedClient.acquire(); assertThat(client1).isSameAs(client2); assertThat(managedClient.refCount()).isEqualTo(2); @@ -128,8 +127,7 @@ public class TestHttpClientCache { SdkHttpClient mockClient = mock(SdkHttpClient.class); final String cacheKey = "test-key"; - HttpClientCache.ManagedHttpClient managedClient = - new HttpClientCache.ManagedHttpClient(mockClient, cacheKey); + ManagedHttpClient managedClient = new ManagedHttpClient(mockClient, cacheKey); // Acquire and release to close managedClient.acquire(); @@ -151,7 +149,7 @@ public class TestHttpClientCache { SdkHttpClient client1 = cache.getOrCreateClient(cacheKey, mockFactory1); assertThat(client1).isNotNull(); - ConcurrentMap<String, HttpClientCache.ManagedHttpClient> clientMap = cache.clientMap(); + ConcurrentMap<String, ManagedHttpClient> clientMap = cache.clientMap(); assertThat(clientMap).containsKey(cacheKey); // Verify ref count is 1 @@ -202,13 +200,13 @@ public class TestHttpClientCache { } // Verify reference count equals number of threads - HttpClientCache.ManagedHttpClient managedClient = cache.clientMap().get(cacheKey); + ManagedHttpClient managedClient = cache.clientMap().get(cacheKey); assertThat(managedClient.refCount()).isEqualTo(threadCount); } @Test public void testRegistryShutdown() { - ConcurrentMap<String, HttpClientCache.ManagedHttpClient> clientMap = cache.clientMap(); + ConcurrentMap<String, ManagedHttpClient> clientMap = cache.clientMap(); // Create some clients cache.getOrCreateClient("key1", mockFactory1); @@ -233,11 +231,10 @@ public class TestHttpClientCache { SdkHttpClient mockClient = mock(SdkHttpClient.class); final String cacheKey = "test-key"; - HttpClientCache.ManagedHttpClient managedClient = - new HttpClientCache.ManagedHttpClient(mockClient, cacheKey); + ManagedHttpClient managedClient = new ManagedHttpClient(mockClient, cacheKey); // Acquire once - WrappedSdkHttpClient client = managedClient.acquire(); + ManagedHttpClient client = managedClient.acquire(); assertThat(managedClient.refCount()).isEqualTo(1); // First release should close the client (refCount goes to 0) @@ -261,8 +258,7 @@ public class TestHttpClientCache { SdkHttpClient mockClient = mock(SdkHttpClient.class); final String cacheKey = "test-key"; - HttpClientCache.ManagedHttpClient managedClient = - new HttpClientCache.ManagedHttpClient(mockClient, cacheKey); + ManagedHttpClient managedClient = new ManagedHttpClient(mockClient, cacheKey); // Acquire once managedClient.acquire(); diff --git a/aws/src/test/java/org/apache/iceberg/aws/TestHttpClientProperties.java b/aws/src/test/java/org/apache/iceberg/aws/TestHttpClientProperties.java index 10377412b6..378e5e6ca9 100644 --- a/aws/src/test/java/org/apache/iceberg/aws/TestHttpClientProperties.java +++ b/aws/src/test/java/org/apache/iceberg/aws/TestHttpClientProperties.java @@ -24,7 +24,7 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.verify; import java.util.Map; -import org.apache.iceberg.aws.HttpClientCache.WrappedSdkHttpClient; +import org.apache.iceberg.aws.HttpClientCache.ManagedHttpClient; import org.apache.iceberg.relocated.com.google.common.collect.Maps; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; @@ -50,12 +50,12 @@ public class TestHttpClientProperties { SdkHttpClient capturedHttpClient = httpClientCaptor.getValue(); assertThat(capturedHttpClient) - .as("Should use wrapped SDK http client") - .isInstanceOf(WrappedSdkHttpClient.class); + .as("Should use managed SDK http client") + .isInstanceOf(ManagedHttpClient.class); // Verify the underlying delegate is UrlConnectionHttpClient - WrappedSdkHttpClient wrappedClient = (WrappedSdkHttpClient) capturedHttpClient; - assertThat(wrappedClient.delegate()) + ManagedHttpClient managedClient = (ManagedHttpClient) capturedHttpClient; + assertThat(managedClient.httpClient()) .as("Underlying client should be UrlConnectionHttpClient") .isInstanceOf(UrlConnectionHttpClient.class); } @@ -73,12 +73,12 @@ public class TestHttpClientProperties { SdkHttpClient capturedHttpClient = httpClientCaptor.getValue(); assertThat(capturedHttpClient) - .as("Should use wrapped SDK http client") - .isInstanceOf(WrappedSdkHttpClient.class); + .as("Should use managed SDK http client") + .isInstanceOf(ManagedHttpClient.class); // Verify the underlying delegate is ApacheHttpClient - WrappedSdkHttpClient wrappedClient = (WrappedSdkHttpClient) capturedHttpClient; - assertThat(wrappedClient.delegate()) + ManagedHttpClient managedClient = (ManagedHttpClient) capturedHttpClient; + assertThat(managedClient.httpClient()) .as("Underlying client should be ApacheHttpClient") .isInstanceOf(ApacheHttpClient.class); } @@ -103,8 +103,8 @@ public class TestHttpClientProperties { apacheConfig.configureHttpClientBuilder(mockS3ClientBuilder); - // Verify that httpClient() is called with a wrapped client (as a shared resource) - verify(mockS3ClientBuilder).httpClient(any(WrappedSdkHttpClient.class)); + // Verify that httpClient() is called with a managed client (as a shared resource) + verify(mockS3ClientBuilder).httpClient(any(ManagedHttpClient.class)); } @Test @@ -116,7 +116,7 @@ public class TestHttpClientProperties { urlConfig.configureHttpClientBuilder(mockS3ClientBuilder); - // Verify that httpClient() is called with a wrapped client (as a shared resource) - verify(mockS3ClientBuilder).httpClient(any(WrappedSdkHttpClient.class)); + // Verify that httpClient() is called with a managed client (as a shared resource) + verify(mockS3ClientBuilder).httpClient(any(ManagedHttpClient.class)); } } ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
