Ensure PooledConnectionProvider check only for event loop in the thread local

reactor-netty-core/src/main/java/reactor/netty/resources/ColocatedEventLoopGroup.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2011-2021 VMware, Inc. or its affiliates, All Rights Reserved.
+ * Copyright (c) 2011-2022 VMware, Inc. or its affiliates, All Rights Reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -33,6 +33,7 @@
 import io.netty.util.concurrent.EventExecutor;
 import io.netty.util.concurrent.FastThreadLocal;
 import io.netty.util.concurrent.ScheduledFuture;
+import reactor.util.annotation.Nullable;
 
 /**
  * Reuse local event loop if already working inside one.
@@ -124,10 +125,8 @@ public Iterator<EventExecutor> iterator() {
 
 	@Override
 	public EventLoop next() {
-		if (localLoop.isSet()) {
-			return localLoop.get();
-		}
-		return eventLoopGroup.next();
+		EventLoop loop = nextInternal();
+		return loop != null ? loop : eventLoopGroup.next();
 	}
 
 	@Override
@@ -218,4 +217,9 @@ void clean() {
 			ex.execute(() -> localLoop.set(null));
 		}
 	}
+
+	@Nullable
+	EventLoop nextInternal() {
+		return localLoop.isSet() ? localLoop.get() : null;
+	}
 }

reactor-netty-core/src/main/java/reactor/netty/resources/DefaultPooledConnectionProvider.java

@@ -499,9 +499,15 @@ static final class PooledConnectionAllocator {
 		Publisher<PooledConnection> connectChannel() {
 			return Mono.create(sink -> {
 				PooledConnectionInitializer initializer = new PooledConnectionInitializer(sink);
-				EventLoop callerEventLoop = sink.contextView().get(CONTEXT_CALLER_EVENTLOOP);
-				TransportConnector.connect(config, remoteAddress, resolver, initializer, callerEventLoop)
-				                  .subscribe(initializer);
+				EventLoop callerEventLoop = sink.contextView().hasKey(CONTEXT_CALLER_EVENTLOOP) ?
+						sink.contextView().get(CONTEXT_CALLER_EVENTLOOP) : null;
+				if (callerEventLoop != null) {
+					TransportConnector.connect(config, remoteAddress, resolver, initializer, callerEventLoop)
+							.subscribe(initializer);
+				}
+				else {
+					TransportConnector.connect(config, remoteAddress, resolver, initializer).subscribe(initializer);
+				}
 			});
 		}
 

reactor-netty-core/src/main/java/reactor/netty/resources/PooledConnectionProvider.java

@@ -17,6 +17,7 @@
 
 import io.netty.channel.Channel;
 import io.netty.channel.EventLoop;
+import io.netty.channel.EventLoopGroup;
 import io.netty.resolver.AddressResolverGroup;
 import org.reactivestreams.Publisher;
 import reactor.core.CoreSubscriber;
@@ -146,12 +147,26 @@ else if (Metrics.isInstrumentationAvailable()) {
 				return newPool;
 			});
 
-			EventLoop eventLoop = sink.contextView().getOrDefault(CONTEXT_CALLER_EVENTLOOP,
-					config.loopResources().onClient(config.isPreferNative()).next());
-			pool.acquire(Duration.ofMillis(poolFactory.pendingAcquireTimeout))
-			    .contextWrite(ctx -> ctx.put(CONTEXT_CALLER_EVENTLOOP, eventLoop))
-			    .subscribe(createDisposableAcquire(config, connectionObserver,
-			            poolFactory.pendingAcquireTimeout, pool, sink));
+			EventLoop eventLoop;
+			if (sink.contextView().hasKey(CONTEXT_CALLER_EVENTLOOP)) {
+				eventLoop = sink.contextView().get(CONTEXT_CALLER_EVENTLOOP);
+			}
+			else {
+				EventLoopGroup group = config.loopResources().onClient(config.isPreferNative());
+				if (group instanceof ColocatedEventLoopGroup) {
+					eventLoop = ((ColocatedEventLoopGroup) group).nextInternal();
+				}
+				else {
+					eventLoop = null;
+				}
+			}
+
+			Mono<PooledRef<T>> mono = pool.acquire(Duration.ofMillis(poolFactory.pendingAcquireTimeout));
+			if (eventLoop != null) {
+				mono = mono.contextWrite(ctx -> ctx.put(CONTEXT_CALLER_EVENTLOOP, eventLoop));
+			}
+			mono.subscribe(createDisposableAcquire(config, connectionObserver,
+					poolFactory.pendingAcquireTimeout, pool, sink));
 		});
 	}