Callbacks for Connection initialization

apache · Apr 8, 2024 · 333757f · 333757f
1 parent ccff406
commit 333757f
Show file tree

Hide file tree

Showing 6 changed files with 83 additions and 16 deletions.
diff --git a/...src/main/java/org/apache/hc/client5/http/impl/io/DefaultHttpClientConnectionOperator.java b/...src/main/java/org/apache/hc/client5/http/impl/io/DefaultHttpClientConnectionOperator.java
@@ -192,6 +192,7 @@ public void connect(
             final InetAddress address = remoteAddresses[i];
             final boolean last = i == remoteAddresses.length - 1;
             final InetSocketAddress remoteAddress = new InetSocketAddress(address, port);
+            onBeforeSocketConnect(context, endpointHost);
             if (LOG.isDebugEnabled()) {
                 LOG.debug("{} connecting {}->{} ({})", endpointHost, localAddress, remoteAddress, connectTimeout);
             }
@@ -221,6 +222,7 @@ public void connect(
                 }
                 socket.connect(remoteAddress, TimeValue.isPositive(connectTimeout) ? connectTimeout.toMillisecondsIntBound() : 0);
                 conn.bind(socket);
+                onAfterSocketConnect(context, endpointHost);
                 if (LOG.isDebugEnabled()) {
                     LOG.debug("{} {} connected {}->{}", ConnPoolSupport.getId(conn), endpointHost,
                             conn.getLocalAddress(), conn.getRemoteAddress());
@@ -229,11 +231,16 @@ public void connect(
                 final TlsSocketStrategy tlsSocketStrategy = tlsSocketStrategyLookup != null ? tlsSocketStrategyLookup.lookup(endpointHost.getSchemeName()) : null;
                 if (tlsSocketStrategy != null) {
                     final NamedEndpoint tlsName = endpointName != null ? endpointName : endpointHost;
+                    onBeforeTlsHandshake(context, endpointHost);
                     if (LOG.isDebugEnabled()) {
                         LOG.debug("{} {} upgrading to TLS", ConnPoolSupport.getId(conn), tlsName);
                     }
                     final Socket upgradedSocket = tlsSocketStrategy.upgrade(socket, tlsName.getHostName(), tlsName.getPort(), attachment, context);
                     conn.bind(upgradedSocket);
+                    onAfterTlsHandshake(context, endpointHost);
+                    if (LOG.isDebugEnabled()) {
+                        LOG.debug("{} {} upgraded to TLS", ConnPoolSupport.getId(conn), tlsName);
+                    }
                 }
                 return;
             } catch (final RuntimeException ex) {
@@ -278,14 +285,31 @@ public void upgrade(
         final TlsSocketStrategy tlsSocketStrategy = tlsSocketStrategyLookup != null ? tlsSocketStrategyLookup.lookup(newProtocol) : null;
         if (tlsSocketStrategy != null) {
             final NamedEndpoint tlsName = endpointName != null ? endpointName : endpointHost;
+            onBeforeTlsHandshake(context, endpointHost);
             if (LOG.isDebugEnabled()) {
                 LOG.debug("{} upgrading to TLS {}:{}", ConnPoolSupport.getId(conn), tlsName.getHostName(), tlsName.getPort());
             }
             final SSLSocket upgradedSocket = tlsSocketStrategy.upgrade(socket, tlsName.getHostName(), tlsName.getPort(), attachment, context);
             conn.bind(upgradedSocket);
+            onAfterTlsHandshake(context, endpointHost);
+            if (LOG.isDebugEnabled()) {
+                LOG.debug("{} upgraded to TLS {}:{}", ConnPoolSupport.getId(conn), tlsName.getHostName(), tlsName.getPort());
+            }
         } else {
             throw new UnsupportedSchemeException(newProtocol + " protocol is not supported");
         }
     }
 
+    protected void onBeforeSocketConnect(final HttpContext httpContext, final HttpHost endpointHost) {
+    }
+
+    protected void onAfterSocketConnect(final HttpContext httpContext, final HttpHost endpointHost) {
+    }
+
+    protected void onBeforeTlsHandshake(final HttpContext httpContext, final HttpHost endpointHost) {
+    }
+
+    protected void onAfterTlsHandshake(final HttpContext httpContext, final HttpHost endpointHost) {
+    }
+
 }
diff --git a/.../src/main/java/org/apache/hc/client5/http/impl/io/PoolingHttpClientConnectionManager.java b/.../src/main/java/org/apache/hc/client5/http/impl/io/PoolingHttpClientConnectionManager.java
@@ -206,7 +206,7 @@ public PoolingHttpClientConnectionManager(
     }
 
     @Internal
-    protected PoolingHttpClientConnectionManager(
+    public PoolingHttpClientConnectionManager(
             final HttpClientConnectionOperator httpClientConnectionOperator,
             final PoolConcurrencyPolicy poolConcurrencyPolicy,
             final PoolReusePolicy poolReusePolicy,

diff --git a/...in/java/org/apache/hc/client5/http/impl/io/PoolingHttpClientConnectionManagerBuilder.java b/...in/java/org/apache/hc/client5/http/impl/io/PoolingHttpClientConnectionManagerBuilder.java
@@ -34,9 +34,11 @@
 import org.apache.hc.client5.http.SchemePortResolver;
 import org.apache.hc.client5.http.config.ConnectionConfig;
 import org.apache.hc.client5.http.config.TlsConfig;
+import org.apache.hc.client5.http.io.HttpClientConnectionOperator;
 import org.apache.hc.client5.http.io.ManagedHttpClientConnection;
 import org.apache.hc.client5.http.ssl.DefaultClientTlsStrategy;
 import org.apache.hc.client5.http.ssl.TlsSocketStrategy;
+import org.apache.hc.core5.annotation.Internal;
 import org.apache.hc.core5.function.Resolver;
 import org.apache.hc.core5.http.HttpHost;
 import org.apache.hc.core5.http.URIScheme;
@@ -273,15 +275,24 @@ public final PoolingHttpClientConnectionManagerBuilder useSystemProperties() {
         return this;
     }
 
+    @Internal
+    protected HttpClientConnectionOperator createConnectionOperator(
+            final SchemePortResolver schemePortResolver,
+            final DnsResolver dnsResolver,
+            final TlsSocketStrategy tlsSocketStrategy,
+            final boolean systemProperties) {
+        return new DefaultHttpClientConnectionOperator(schemePortResolver, dnsResolver,
+                RegistryBuilder.<TlsSocketStrategy>create()
+                        .register(URIScheme.HTTPS.id, tlsSocketStrategy != null ? tlsSocketStrategy :
+                                (systemProperties ?
+                                        DefaultClientTlsStrategy.createSystemDefault() :
+                                        DefaultClientTlsStrategy.createDefault()))
+                        .build());
+    }
+
     public PoolingHttpClientConnectionManager build() {
         final PoolingHttpClientConnectionManager poolingmgr = new PoolingHttpClientConnectionManager(
-                new DefaultHttpClientConnectionOperator(schemePortResolver, dnsResolver,
-                        RegistryBuilder.<TlsSocketStrategy>create()
-                                .register(URIScheme.HTTPS.id, tlsSocketStrategy != null ? tlsSocketStrategy :
-                                        (systemProperties ?
-                                                DefaultClientTlsStrategy.createSystemDefault() :
-                                                DefaultClientTlsStrategy.createDefault()))
-                                .build()),
+                createConnectionOperator(schemePortResolver, dnsResolver, tlsSocketStrategy, systemProperties),
                 poolConcurrencyPolicy,
                 poolReusePolicy,
                 null,

diff --git a/...c/main/java/org/apache/hc/client5/http/impl/nio/DefaultAsyncClientConnectionOperator.java b/...c/main/java/org/apache/hc/client5/http/impl/nio/DefaultAsyncClientConnectionOperator.java
@@ -41,6 +41,7 @@
 import org.apache.hc.client5.http.nio.AsyncClientConnectionOperator;
 import org.apache.hc.client5.http.nio.ManagedAsyncClientConnection;
 import org.apache.hc.client5.http.routing.RoutingSupport;
+import org.apache.hc.core5.annotation.Internal;
 import org.apache.hc.core5.concurrent.CallbackContribution;
 import org.apache.hc.core5.concurrent.ComplexFuture;
 import org.apache.hc.core5.concurrent.FutureCallback;
@@ -59,7 +60,8 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-final class DefaultAsyncClientConnectionOperator implements AsyncClientConnectionOperator {
+@Internal
+public class DefaultAsyncClientConnectionOperator implements AsyncClientConnectionOperator {
 
     private static final Logger LOG = LoggerFactory.getLogger(DefaultAsyncClientConnectionOperator.class);
 
@@ -105,6 +107,7 @@ public Future<ManagedAsyncClientConnection> connect(
         final InetAddress remoteAddress = endpointHost.getAddress();
         final TlsConfig tlsConfig = attachment instanceof TlsConfig ? (TlsConfig) attachment : TlsConfig.DEFAULT;
 
+        onBeforeSocketConnect(context, endpointHost);
         if (LOG.isDebugEnabled()) {
             LOG.debug("{} connecting {}->{} ({})", endpointHost, localAddress, remoteAddress, connectTimeout);
         }
@@ -121,6 +124,7 @@ public Future<ManagedAsyncClientConnection> connect(
                     @Override
                     public void completed(final IOSession session) {
                         final DefaultManagedAsyncClientConnection connection = new DefaultManagedAsyncClientConnection(session);
+                        onAfterSocketConnect(context, endpointHost);
                         if (LOG.isDebugEnabled()) {
                             LOG.debug("{} {} connected {}->{}", ConnPoolSupport.getId(connection), endpointHost,
                                     connection.getLocalAddress(), connection.getRemoteAddress());
@@ -131,6 +135,7 @@ public void completed(final IOSession session) {
                                 final Timeout socketTimeout = connection.getSocketTimeout();
                                 final Timeout handshakeTimeout = tlsConfig.getHandshakeTimeout();
                                 final NamedEndpoint tlsName = endpointName != null ? endpointName : endpointHost;
+                                onBeforeTlsHandshake(context, endpointHost);
                                 if (LOG.isDebugEnabled()) {
                                     LOG.debug("{} {} upgrading to TLS", ConnPoolSupport.getId(connection), tlsName);
                                 }
@@ -145,6 +150,10 @@ public void completed(final IOSession session) {
                                             public void completed(final TransportSecurityLayer transportSecurityLayer) {
                                                 connection.setSocketTimeout(socketTimeout);
                                                 future.completed(connection);
+                                                onAfterTlsHandshake(context, endpointHost);
+                                                if (LOG.isDebugEnabled()) {
+                                                    LOG.debug("{} {} upgraded to TLS", ConnPoolSupport.getId(connection), tlsName);
+                                                }
                                             }
 
                                         });
@@ -214,4 +223,16 @@ public void completed(final TransportSecurityLayer transportSecurityLayer) {
         }
     }
 
+    protected void onBeforeSocketConnect(final HttpContext httpContext, final HttpHost endpointHost) {
+    }
+
+    protected void onAfterSocketConnect(final HttpContext httpContext, final HttpHost endpointHost) {
+    }
+
+    protected void onBeforeTlsHandshake(final HttpContext httpContext, final HttpHost endpointHost) {
+    }
+
+    protected void onAfterTlsHandshake(final HttpContext httpContext, final HttpHost endpointHost) {
+    }
+
 }
diff --git a/...rc/main/java/org/apache/hc/client5/http/impl/nio/PoolingAsyncClientConnectionManager.java b/...rc/main/java/org/apache/hc/client5/http/impl/nio/PoolingAsyncClientConnectionManager.java
@@ -164,7 +164,7 @@ public PoolingAsyncClientConnectionManager(
     }
 
     @Internal
-    protected PoolingAsyncClientConnectionManager(
+    public PoolingAsyncClientConnectionManager(
             final AsyncClientConnectionOperator connectionOperator,
             final PoolConcurrencyPolicy poolConcurrencyPolicy,
             final PoolReusePolicy poolReusePolicy,

diff --git a/.../java/org/apache/hc/client5/http/impl/nio/PoolingAsyncClientConnectionManagerBuilder.java b/.../java/org/apache/hc/client5/http/impl/nio/PoolingAsyncClientConnectionManagerBuilder.java
@@ -32,8 +32,10 @@
 import org.apache.hc.client5.http.SchemePortResolver;
 import org.apache.hc.client5.http.config.ConnectionConfig;
 import org.apache.hc.client5.http.config.TlsConfig;
+import org.apache.hc.client5.http.nio.AsyncClientConnectionOperator;
 import org.apache.hc.client5.http.ssl.ConscryptClientTlsStrategy;
 import org.apache.hc.client5.http.ssl.DefaultClientTlsStrategy;
+import org.apache.hc.core5.annotation.Internal;
 import org.apache.hc.core5.function.Resolver;
 import org.apache.hc.core5.http.HttpHost;
 import org.apache.hc.core5.http.URIScheme;
@@ -229,6 +231,19 @@ public final PoolingAsyncClientConnectionManagerBuilder useSystemProperties() {
         return this;
     }
 
+    @Internal
+    protected AsyncClientConnectionOperator createConnectionOperator(
+            final TlsStrategy tlsStrategy,
+            final SchemePortResolver schemePortResolver,
+            final DnsResolver dnsResolver) {
+        return new DefaultAsyncClientConnectionOperator(
+                RegistryBuilder.<TlsStrategy>create()
+                        .register(URIScheme.HTTPS.getId(), tlsStrategy)
+                        .build(),
+                schemePortResolver,
+                dnsResolver);
+    }
+
     public PoolingAsyncClientConnectionManager build() {
         final TlsStrategy tlsStrategyCopy;
         if (tlsStrategy != null) {
@@ -249,14 +264,10 @@ public PoolingAsyncClientConnectionManager build() {
             }
         }
         final PoolingAsyncClientConnectionManager poolingmgr = new PoolingAsyncClientConnectionManager(
-                RegistryBuilder.<TlsStrategy>create()
-                        .register(URIScheme.HTTPS.getId(), tlsStrategyCopy)
-                        .build(),
+                createConnectionOperator(tlsStrategyCopy, schemePortResolver, dnsResolver),
                 poolConcurrencyPolicy,
                 poolReusePolicy,
-                null,
-                schemePortResolver,
-                dnsResolver);
+                null);
         poolingmgr.setConnectionConfigResolver(connectionConfigResolver);
         poolingmgr.setTlsConfigResolver(tlsConfigResolver);
         if (maxConnTotal > 0) {