3535import io .netty .handler .logging .LoggingHandler ;
3636import io .netty .handler .ssl .SslHandler ;
3737import io .netty .handler .stream .ChunkedWriteHandler ;
38+ import io .netty .util .AttributeKey ;
3839import io .netty .util .Timer ;
3940import io .netty .util .concurrent .DefaultThreadFactory ;
4041import io .netty .util .concurrent .GlobalEventExecutor ;
@@ -87,6 +88,8 @@ public class ChannelManager {
8788public static final String AHC_WS_HANDLER ="ahc-ws" ;
8889public static final String LOGGING_HANDLER ="logging" ;
8990
91+ private static final AttributeKey <Object >partitionKeyAttr =AttributeKey .valueOf ("partitionKey" );
92+
9093private final AsyncHttpClientConfig config ;
9194private final SslEngineFactory sslEngineFactory ;
9295private final EventLoopGroup eventLoopGroup ;
@@ -99,7 +102,6 @@ public class ChannelManager {
99102
100103private final ChannelPool channelPool ;
101104private final ChannelGroup openChannels ;
102- private final ConcurrentHashMap <Channel ,Object >channelId2PartitionKey =new ConcurrentHashMap <>();
103105private final boolean maxTotalConnectionsEnabled ;
104106private final Semaphore freeChannels ;
105107private final boolean maxConnectionsPerHostEnabled ;
@@ -141,7 +143,7 @@ public boolean remove(Object o) {
141143if (maxTotalConnectionsEnabled )
142144freeChannels .release ();
143145if (maxConnectionsPerHostEnabled ) {
144- Object partitionKey =channelId2PartitionKey . remove ( Channel .class .cast (o ));
146+ Object partitionKey =Channel .class .cast (o ). attr ( partitionKeyAttr ). getAndSet ( null );
145147if (partitionKey !=null ) {
146148Semaphore hostFreeChannels =freeChannelsPerHost .get (partitionKey );
147149if (hostFreeChannels !=null )
@@ -308,7 +310,7 @@ public final void tryToOfferChannelToPool(Channel channel, AsyncHandler<?> async
308310AsyncHandlerExtensions .class .cast (asyncHandler ).onConnectionOffer (channel );
309311if (channelPool .offer (channel ,partitionKey )) {
310312if (maxConnectionsPerHostEnabled )
311- channelId2PartitionKey . putIfAbsent ( channel , partitionKey );
313+ channel . attr ( partitionKeyAttr ). setIfAbsent ( partitionKey );
312314 }else {
313315// rejected by pool
314316closeChannel (channel );
@@ -383,7 +385,7 @@ public void releaseChannelLock(Object partitionKey) {
383385public void registerOpenChannel (Channel channel ,Object partitionKey ) {
384386openChannels .add (channel );
385387if (maxConnectionsPerHostEnabled ) {
386- channelId2PartitionKey . put ( channel , partitionKey );
388+ channel . attr ( partitionKeyAttr ). set ( partitionKey );
387389 }
388390 }
389391