Skip to content

Commit a1f833f

Browse files
rozzarozza
committed
Added operation context to authentication
JAVA-5211
1 parent 1a9785d commit a1f833f

22 files changed

+348
-154
lines changed

‎driver-core/src/main/com/mongodb/internal/async/function/RetryState.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,9 @@ private void doAdvanceOrThrow(final Throwable attemptException,
182182
final boolean onlyRuntimeExceptions) throws Throwable {
183183
assertTrue(attempt() < attempts);
184184
assertNotNull(attemptException);
185+
if (attemptException instanceof MongoOperationTimeoutException) {
186+
throw attemptException;
187+
}
185188
if (onlyRuntimeExceptions) {
186189
assertTrue(isRuntime(attemptException));
187190
}

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -89,8 +89,9 @@ <T> T getNonNullMechanismProperty(final String key, @Nullable final T defaultVal
8989

9090
}
9191

92-
abstract void authenticate(InternalConnection connection, ConnectionDescription connectionDescription);
92+
abstract void authenticate(InternalConnection connection, ConnectionDescription connectionDescription,
93+
OperationContext operationContext);
9394

9495
abstract void authenticateAsync(InternalConnection connection, ConnectionDescription connectionDescription,
95-
SingleResultCallback<Void> callback);
96+
OperationContext operationContext, SingleResultCallback<Void> callback);
9697
}

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

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -47,24 +47,23 @@ public final class CommandHelper {
4747
static final String LEGACY_HELLO_LOWER = LEGACY_HELLO.toLowerCase(Locale.ROOT);
4848

4949
static BsonDocument executeCommand(final String database, final BsonDocument command, final ClusterConnectionMode clusterConnectionMode,
50-
@Nullable final ServerApi serverApi, final InternalConnection internalConnection) {
51-
return sendAndReceive(database, command, clusterConnectionMode, serverApi, internalConnection);
50+
@Nullable final ServerApi serverApi, final InternalConnection internalConnection, final OperationContext operationContext) {
51+
return sendAndReceive(database, command, clusterConnectionMode, serverApi, internalConnection, operationContext);
5252
}
5353

5454
static BsonDocument executeCommandWithoutCheckingForFailure(final String database, final BsonDocument command,
55-
final ClusterConnectionMode clusterConnectionMode,
56-
@Nullable final ServerApi serverApi,
57-
final InternalConnection internalConnection) {
55+
final ClusterConnectionMode clusterConnectionMode, @Nullable final ServerApi serverApi,
56+
final InternalConnection internalConnection, final OperationContext operationContext) {
5857
try {
59-
return sendAndReceive(database, command, clusterConnectionMode, serverApi, internalConnection);
58+
return executeCommand(database, command, clusterConnectionMode, serverApi, internalConnection, operationContext);
6059
} catch (MongoServerException e) {
6160
return new BsonDocument();
6261
}
6362
}
6463

6564
static void executeCommandAsync(final String database, final BsonDocument command, final ClusterConnectionMode clusterConnectionMode,
66-
@Nullable final ServerApi serverApi, final InternalConnection internalConnection,
67-
final SingleResultCallback<BsonDocument> callback) {
65+
@Nullable final ServerApi serverApi, final InternalConnection internalConnection, final OperationContext operationContext,
66+
final SingleResultCallback<BsonDocument> callback) {
6867
internalConnection.sendAndReceiveAsync(getCommandMessage(database, command, internalConnection, clusterConnectionMode, serverApi),
6968
new BsonDocumentCodec(), createOperationContext(NoOpSessionContext.INSTANCE, serverApi),
7069
(result, t) -> {
@@ -98,11 +97,12 @@ static OperationContext createOperationContext(final SessionContext sessionConte
9897
private static BsonDocument sendAndReceive(final String database, final BsonDocument command,
9998
final ClusterConnectionMode clusterConnectionMode,
10099
@Nullable final ServerApi serverApi,
101-
final InternalConnection internalConnection) {
100+
final InternalConnection internalConnection,
101+
final OperationContext operationContext) {
102102
return assertNotNull(
103103
internalConnection.sendAndReceive(
104104
getCommandMessage(database, command, internalConnection, clusterConnectionMode, serverApi),
105-
new BsonDocumentCodec(), createOperationContext(NoOpSessionContext.INSTANCE, serverApi))
105+
new BsonDocumentCodec(), operationContext)
106106
);
107107
}
108108

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

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -47,14 +47,15 @@ class DefaultAuthenticator extends Authenticator implements SpeculativeAuthentic
4747
}
4848

4949
@Override
50-
void authenticate(final InternalConnection connection, final ConnectionDescription connectionDescription) {
50+
void authenticate(final InternalConnection connection, final ConnectionDescription connectionDescription,
51+
final OperationContext operationContext) {
5152
if (serverIsLessThanVersionFourDotZero(connectionDescription)) {
5253
new ScramShaAuthenticator(getMongoCredentialWithCache().withMechanism(SCRAM_SHA_1), getClusterConnectionMode(), getServerApi())
53-
.authenticate(connection, connectionDescription);
54+
.authenticate(connection, connectionDescription, operationContext);
5455
} else {
5556
try {
5657
setDelegate(connectionDescription);
57-
delegate.authenticate(connection, connectionDescription);
58+
delegate.authenticate(connection, connectionDescription, operationContext);
5859
} catch (Exception e) {
5960
throw wrapException(e);
6061
}
@@ -63,13 +64,13 @@ void authenticate(final InternalConnection connection, final ConnectionDescripti
6364

6465
@Override
6566
void authenticateAsync(final InternalConnection connection, final ConnectionDescription connectionDescription,
66-
final SingleResultCallback<Void> callback) {
67+
final OperationContext operationContext, final SingleResultCallback<Void> callback) {
6768
if (serverIsLessThanVersionFourDotZero(connectionDescription)) {
6869
new ScramShaAuthenticator(getMongoCredentialWithCache().withMechanism(SCRAM_SHA_1), getClusterConnectionMode(), getServerApi())
69-
.authenticateAsync(connection, connectionDescription, callback);
70+
.authenticateAsync(connection, connectionDescription, operationContext, callback);
7071
} else {
7172
setDelegate(connectionDescription);
72-
delegate.authenticateAsync(connection, connectionDescription, callback);
73+
delegate.authenticateAsync(connection, connectionDescription, operationContext, callback);
7374
}
7475
}
7576

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -437,9 +437,10 @@ private void initialize() {
437437

438438
private void pingServer(final InternalConnection connection) {
439439
long start = System.nanoTime();
440+
OperationContext operationContext = operationContextFactory.create();
440441
executeCommand("admin",
441442
new BsonDocument(getHandshakeCommandName(connection.getInitialServerDescription()), new BsonInt32(1)),
442-
clusterConnectionMode, serverApi, connection);
443+
clusterConnectionMode, serverApi, connection, operationContext);
443444
long elapsedTimeNanos = System.nanoTime() - start;
444445
roundTripTimeSampler.addSample(elapsedTimeNanos);
445446
}

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

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,14 +20,19 @@
2020

2121
interface InternalConnectionInitializer {
2222

23-
InternalConnectionInitializationDescription startHandshake(InternalConnection internalConnection);
23+
InternalConnectionInitializationDescription startHandshake(InternalConnection internalConnection,
24+
OperationContext operationContext);
2425

2526
InternalConnectionInitializationDescription finishHandshake(InternalConnection internalConnection,
26-
InternalConnectionInitializationDescription description);
27+
InternalConnectionInitializationDescription description,
28+
OperationContext operationContext);
2729

2830
void startHandshakeAsync(InternalConnection internalConnection,
31+
OperationContext operationContext,
2932
SingleResultCallback<InternalConnectionInitializationDescription> callback);
3033

31-
void finishHandshakeAsync(InternalConnection internalConnection, InternalConnectionInitializationDescription description,
34+
void finishHandshakeAsync(InternalConnection internalConnection,
35+
InternalConnectionInitializationDescription description,
36+
OperationContext operationContext,
3237
SingleResultCallback<InternalConnectionInitializationDescription> callback);
3338
}

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -202,10 +202,10 @@ public void open(final OperationContext operationContext) {
202202
try {
203203
stream.open(operationContext);
204204

205-
InternalConnectionInitializationDescription initializationDescription = connectionInitializer.startHandshake(this);
205+
InternalConnectionInitializationDescription initializationDescription = connectionInitializer.startHandshake(this, operationContext);
206206
initAfterHandshakeStart(initializationDescription);
207207

208-
initializationDescription = connectionInitializer.finishHandshake(this, initializationDescription);
208+
initializationDescription = connectionInitializer.finishHandshake(this, initializationDescription, operationContext);
209209
initAfterHandshakeFinish(initializationDescription);
210210
} catch (Throwable t) {
211211
close();
@@ -226,7 +226,7 @@ public void openAsync(final OperationContext operationContext, final SingleResul
226226

227227
@Override
228228
public void completed(@Nullable final Void aVoid) {
229-
connectionInitializer.startHandshakeAsync(InternalStreamConnection.this,
229+
connectionInitializer.startHandshakeAsync(InternalStreamConnection.this, operationContext,
230230
(initialResult, initialException) -> {
231231
if (initialException != null) {
232232
close();
@@ -235,7 +235,7 @@ public void completed(@Nullable final Void aVoid) {
235235
assertNotNull(initialResult);
236236
initAfterHandshakeStart(initialResult);
237237
connectionInitializer.finishHandshakeAsync(InternalStreamConnection.this,
238-
initialResult, (completedResult, completedException) -> {
238+
initialResult, operationContext, (completedResult, completedException) -> {
239239
if (completedException != null) {
240240
close();
241241
callback.onResult(null, completedException);

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

Lines changed: 26 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -72,27 +72,29 @@ public InternalStreamConnectionInitializer(final ClusterConnectionMode clusterCo
7272
}
7373

7474
@Override
75-
public InternalConnectionInitializationDescription startHandshake(final InternalConnection internalConnection) {
75+
public InternalConnectionInitializationDescription startHandshake(final InternalConnection internalConnection,
76+
final OperationContext operationContext) {
7677
notNull("internalConnection", internalConnection);
7778

78-
return initializeConnectionDescription(internalConnection);
79+
return initializeConnectionDescription(internalConnection, operationContext);
7980
}
8081

8182
public InternalConnectionInitializationDescription finishHandshake(final InternalConnection internalConnection,
82-
final InternalConnectionInitializationDescription description) {
83+
final InternalConnectionInitializationDescription description,
84+
final OperationContext operationContext) {
8385
notNull("internalConnection", internalConnection);
8486
notNull("description", description);
8587

86-
authenticate(internalConnection, description.getConnectionDescription());
87-
return completeConnectionDescriptionInitialization(internalConnection, description);
88+
authenticate(internalConnection, description.getConnectionDescription(), operationContext);
89+
return completeConnectionDescriptionInitialization(internalConnection, description, operationContext);
8890
}
8991

9092
@Override
91-
public void startHandshakeAsync(final InternalConnection internalConnection,
93+
public void startHandshakeAsync(final InternalConnection internalConnection, final OperationContext operationContext,
9294
final SingleResultCallback<InternalConnectionInitializationDescription> callback) {
9395
long startTime = System.nanoTime();
9496
executeCommandAsync("admin", createHelloCommand(authenticator, internalConnection), clusterConnectionMode, serverApi,
95-
internalConnection, (helloResult, t) -> {
97+
internalConnection, operationContext, (helloResult, t) -> {
9698
if (t != null) {
9799
callback.onResult(null, t instanceof MongoException ? mapHelloException((MongoException) t) : t);
98100
} else {
@@ -105,31 +107,35 @@ public void startHandshakeAsync(final InternalConnection internalConnection,
105107
@Override
106108
public void finishHandshakeAsync(final InternalConnection internalConnection,
107109
final InternalConnectionInitializationDescription description,
110+
final OperationContext operationContext,
108111
final SingleResultCallback<InternalConnectionInitializationDescription> callback) {
109112
if (authenticator == null || description.getConnectionDescription().getServerType()
110113
== ServerType.REPLICA_SET_ARBITER) {
111-
completeConnectionDescriptionInitializationAsync(internalConnection, description, callback);
114+
completeConnectionDescriptionInitializationAsync(internalConnection, description, operationContext, callback);
112115
} else {
113-
authenticator.authenticateAsync(internalConnection, description.getConnectionDescription(),
116+
authenticator.authenticateAsync(internalConnection, description.getConnectionDescription(), operationContext,
114117
(result1, t1) -> {
115118
if (t1 != null) {
116119
callback.onResult(null, t1);
117120
} else {
118-
completeConnectionDescriptionInitializationAsync(internalConnection, description, callback);
121+
completeConnectionDescriptionInitializationAsync(internalConnection, description, operationContext, callback);
119122
}
120123
});
121124
}
122125
}
123126

124-
private InternalConnectionInitializationDescription initializeConnectionDescription(final InternalConnection internalConnection) {
127+
private InternalConnectionInitializationDescription initializeConnectionDescription(final InternalConnection internalConnection,
128+
final OperationContext operationContext) {
125129
BsonDocument helloResult;
126130
BsonDocument helloCommandDocument = createHelloCommand(authenticator, internalConnection);
127131

128132
long start = System.nanoTime();
129133
try {
130-
helloResult = executeCommand("admin", helloCommandDocument, clusterConnectionMode, serverApi, internalConnection);
134+
helloResult = executeCommand("admin", helloCommandDocument, clusterConnectionMode, serverApi, internalConnection, operationContext);
131135
} catch (MongoException e) {
132136
throw mapHelloException(e);
137+
} finally {
138+
operationContext.getTimeoutContext().resetMaintenanceTimeout();
133139
}
134140
setSpeculativeAuthenticateResponse(helloResult);
135141
return createInitializationDescription(helloResult, internalConnection, start);
@@ -189,21 +195,23 @@ private BsonDocument createHelloCommand(final Authenticator authenticator, final
189195

190196
private InternalConnectionInitializationDescription completeConnectionDescriptionInitialization(
191197
final InternalConnection internalConnection,
192-
final InternalConnectionInitializationDescription description) {
198+
final InternalConnectionInitializationDescription description,
199+
final OperationContext operationContext) {
193200

194201
if (description.getConnectionDescription().getConnectionId().getServerValue() != null) {
195202
return description;
196203
}
197204

198205
return applyGetLastErrorResult(executeCommandWithoutCheckingForFailure("admin",
199206
new BsonDocument("getlasterror", new BsonInt32(1)), clusterConnectionMode, serverApi,
200-
internalConnection),
207+
internalConnection, operationContext),
201208
description);
202209
}
203210

204-
private void authenticate(final InternalConnection internalConnection, final ConnectionDescription connectionDescription) {
211+
private void authenticate(final InternalConnection internalConnection, final ConnectionDescription connectionDescription,
212+
final OperationContext operationContext) {
205213
if (authenticator != null && connectionDescription.getServerType() != ServerType.REPLICA_SET_ARBITER) {
206-
authenticator.authenticate(internalConnection, connectionDescription);
214+
authenticator.authenticate(internalConnection, connectionDescription, operationContext);
207215
}
208216
}
209217

@@ -217,6 +225,7 @@ private void setSpeculativeAuthenticateResponse(final BsonDocument helloResult)
217225
private void completeConnectionDescriptionInitializationAsync(
218226
final InternalConnection internalConnection,
219227
final InternalConnectionInitializationDescription description,
228+
final OperationContext operationContext,
220229
final SingleResultCallback<InternalConnectionInitializationDescription> callback) {
221230

222231
if (description.getConnectionDescription().getConnectionId().getServerValue() != null) {
@@ -225,7 +234,7 @@ private void completeConnectionDescriptionInitializationAsync(
225234
}
226235

227236
executeCommandAsync("admin", new BsonDocument("getlasterror", new BsonInt32(1)), clusterConnectionMode, serverApi,
228-
internalConnection,
237+
internalConnection, operationContext,
229238
(result, t) -> {
230239
if (t != null) {
231240
callback.onResult(description, null);

0 commit comments

Comments
 (0)