Skip to content

Commit 839e51d

Browse files
committed
Remove ConnectionRegistrationState.
JAVA-5856
1 parent 4a932eb commit 839e51d

File tree

1 file changed

+23
-35
lines changed

1 file changed

+23
-35
lines changed

driver-core/src/main/com/mongodb/internal/connection/TlsChannelStreamFactoryFactory.java

+23-35
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,6 @@
5252
import java.util.concurrent.TimeUnit;
5353
import java.util.concurrent.atomic.AtomicReference;
5454

55-
import static com.mongodb.assertions.Assertions.assertFalse;
5655
import static com.mongodb.assertions.Assertions.assertTrue;
5756
import static com.mongodb.assertions.Assertions.isTrue;
5857
import static com.mongodb.internal.connection.ServerAddressHelper.getSocketAddresses;
@@ -100,35 +99,34 @@ public void close() {
10099
group.shutdown();
101100
}
102101

102+
/**
103+
* Monitors `OP_CONNECT` events for socket connections.
104+
*/
103105
private static class SelectorMonitor implements Closeable {
104106

105107
static final class SocketRegistration {
106108
private final SocketChannel socketChannel;
107-
private final Runnable attachment;
108-
private final AtomicReference<ConnectionRegistrationState> connectionRegistrationState;
109+
private final AtomicReference<Runnable> afterConnectAction;
109110

110-
enum ConnectionRegistrationState {
111-
CONNECTING,
112-
CONNECTED,
113-
TIMEOUT_OUT
111+
SocketRegistration(final SocketChannel socketChannel, final Runnable afterConnectAction) {
112+
this.socketChannel = socketChannel;
113+
this.afterConnectAction = new AtomicReference<>(afterConnectAction);
114114
}
115115

116-
private SocketRegistration(final SocketChannel socketChannel, final Runnable attachment) {
117-
this.socketChannel = socketChannel;
118-
this.attachment = attachment;
119-
this.connectionRegistrationState = new AtomicReference<>(ConnectionRegistrationState.CONNECTING);
116+
boolean tryCancelPendingConnection() {
117+
return tryTakeAction() != null;
120118
}
121119

122-
public boolean markConnectionEstablishmentTimedOut() {
123-
return connectionRegistrationState.compareAndSet(
124-
ConnectionRegistrationState.CONNECTING,
125-
ConnectionRegistrationState.TIMEOUT_OUT);
120+
void runAfterConnectActionIfNotCanceled() {
121+
Runnable afterConnectActionToExecute = tryTakeAction();
122+
if (afterConnectActionToExecute != null) {
123+
afterConnectActionToExecute.run();
124+
}
126125
}
127126

128-
public boolean markConnectionEstablishmentCompleted() {
129-
return connectionRegistrationState.compareAndSet(
130-
ConnectionRegistrationState.CONNECTING,
131-
ConnectionRegistrationState.CONNECTED);
127+
@Nullable
128+
private Runnable tryTakeAction() {
129+
return afterConnectAction.getAndSet(null);
132130
}
133131
}
134132

@@ -144,7 +142,6 @@ public boolean markConnectionEstablishmentCompleted() {
144142
}
145143
}
146144

147-
// Monitors OP_CONNECT events.
148145
void start() {
149146
Thread selectorThread = new Thread(() -> {
150147
try {
@@ -153,13 +150,7 @@ void start() {
153150
selector.select();
154151
for (SelectionKey selectionKey : selector.selectedKeys()) {
155152
selectionKey.cancel();
156-
SocketRegistration socketRegistration = (SocketRegistration) selectionKey.attachment();
157-
158-
boolean markedCompleted = socketRegistration.markConnectionEstablishmentCompleted();
159-
if (markedCompleted) {
160-
Runnable runnable = socketRegistration.attachment;
161-
runnable.run();
162-
}
153+
((SocketRegistration) selectionKey.attachment()).runAfterConnectActionIfNotCanceled();
163154
}
164155

165156
for (Iterator<SocketRegistration> iter = pendingRegistrations.iterator(); iter.hasNext();) {
@@ -228,15 +219,14 @@ public void openAsync(final OperationContext operationContext, final AsyncComple
228219
if (getSettings().getSendBufferSize() > 0) {
229220
socketChannel.setOption(StandardSocketOptions.SO_SNDBUF, getSettings().getSendBufferSize());
230221
}
231-
222+
//getConnectTimeoutMs MUST be called before connection attempt, as it might throw MongoOperationTimeout exception.
223+
int connectTimeoutMs = operationContext.getTimeoutContext().getConnectTimeoutMs();
232224
socketChannel.connect(getSocketAddresses(getServerAddress(), inetAddressResolver).get(0));
233-
234225
SelectorMonitor.SocketRegistration socketRegistration = new SelectorMonitor.SocketRegistration(
235226
socketChannel, () -> initializeTslChannel(handler, socketChannel));
236227

237-
int connectTimeoutMs = getSettings().getConnectTimeout(TimeUnit.MILLISECONDS);
238228
if (connectTimeoutMs > 0) {
239-
scheduleTimeoutInterruption(handler, socketRegistration, socketChannel, connectTimeoutMs);
229+
scheduleTimeoutInterruption(handler, socketRegistration, connectTimeoutMs);
240230
}
241231
selectorMonitor.register(socketRegistration);
242232
} catch (IOException e) {
@@ -248,12 +238,10 @@ public void openAsync(final OperationContext operationContext, final AsyncComple
248238

249239
private void scheduleTimeoutInterruption(final AsyncCompletionHandler<Void> handler,
250240
final SelectorMonitor.SocketRegistration socketRegistration,
251-
final SocketChannel socketChannel,
252241
final int connectTimeoutMs) {
253242
group.getTimeoutExecutor().schedule(() -> {
254-
boolean markedTimedOut = socketRegistration.markConnectionEstablishmentTimedOut();
255-
if (markedTimedOut) {
256-
closeAndTimeout(handler, socketChannel);
243+
if (socketRegistration.tryCancelPendingConnection()) {
244+
closeAndTimeout(handler, socketRegistration.socketChannel);
257245
}
258246
}, connectTimeoutMs, TimeUnit.MILLISECONDS);
259247
}

0 commit comments

Comments
 (0)