Skip to content

Commit 4a16c29

Browse files
nhachichanhachicha
committed
Simplifying bookkeeping logic for Span grouping
1 parent f650d43 commit 4a16c29

File tree

34 files changed

+1195
-276
lines changed

34 files changed

+1195
-276
lines changed

‎driver-core/src/main/com/mongodb/MongoClientSettings.java

Lines changed: 13 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -30,11 +30,10 @@
3030
import com.mongodb.connection.SslSettings;
3131
import com.mongodb.connection.TransportSettings;
3232
import com.mongodb.event.CommandListener;
33-
import com.mongodb.internal.tracing.TracingManager;
33+
import com.mongodb.internal.tracing.Tracer;
3434
import com.mongodb.lang.Nullable;
3535
import com.mongodb.spi.dns.DnsClient;
3636
import com.mongodb.spi.dns.InetAddressResolver;
37-
import com.mongodb.tracing.Tracer;
3837
import org.bson.UuidRepresentation;
3938
import org.bson.codecs.BsonCodecProvider;
4039
import org.bson.codecs.BsonValueCodecProvider;
@@ -120,7 +119,7 @@ public final class MongoClientSettings {
120119
private final InetAddressResolver inetAddressResolver;
121120
@Nullable
122121
private final Long timeoutMS;
123-
private final TracingManager tracingManager;
122+
private final Tracer tracer;
124123

125124
/**
126125
* Gets the default codec registry. It includes the following providers:
@@ -241,7 +240,7 @@ public static final class Builder {
241240
private ContextProvider contextProvider;
242241
private DnsClient dnsClient;
243242
private InetAddressResolver inetAddressResolver;
244-
private TracingManager tracingManager;
243+
private Tracer tracer;
245244

246245
private Builder() {
247246
}
@@ -279,7 +278,7 @@ private Builder(final MongoClientSettings settings) {
279278
if (settings.heartbeatSocketTimeoutSetExplicitly) {
280279
heartbeatSocketTimeoutMS = settings.heartbeatSocketSettings.getReadTimeout(MILLISECONDS);
281280
}
282-
tracingManager = settings.tracingManager;
281+
tracer = settings.tracer;
283282
}
284283

285284
/**
@@ -729,16 +728,16 @@ Builder heartbeatSocketTimeoutMS(final int heartbeatSocketTimeoutMS) {
729728
}
730729

731730
/**
732-
* Sets the tracer to use for creating Spans for operations and commands.
731+
* Sets the tracer to use for creating Spans for operations, commands and transactions.
733732
*
734733
* @param tracer the tracer
735734
* @see com.mongodb.tracing.MicrometerTracer
736735
* @return this
737-
* @since 5.5
736+
* @since 5.6
738737
*/
739738
@Alpha(Reason.CLIENT)
740739
public Builder tracer(final Tracer tracer) {
741-
this.tracingManager = new TracingManager(tracer);
740+
this.tracer = tracer;
742741
return this;
743742
}
744743

@@ -1060,14 +1059,13 @@ public ContextProvider getContextProvider() {
10601059
}
10611060

10621061
/**
1063-
* Get the tracer to create Spans for operations and commands.
1062+
* Get the tracer to create Spans for operations, commands and transactions.
10641063
*
1065-
* @return this
1066-
* @since 5.5
1064+
* @return the configured Tracer
1065+
* @since 5.6
10671066
*/
1068-
@Alpha(Reason.CLIENT)
1069-
public TracingManager getTracingManager() {
1070-
return tracingManager;
1067+
public Tracer getTracer() {
1068+
return tracer;
10711069
}
10721070

10731071
@Override
@@ -1186,6 +1184,6 @@ private MongoClientSettings(final Builder builder) {
11861184
heartbeatConnectTimeoutSetExplicitly = builder.heartbeatConnectTimeoutMS != 0;
11871185
contextProvider = builder.contextProvider;
11881186
timeoutMS = builder.timeoutMS;
1189-
tracingManager = builder.tracingManager;
1187+
tracer = (builder.tracer == null) ? Tracer.NO_OP : builder.tracer;
11901188
}
11911189
}

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

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -186,10 +186,6 @@ BsonDocument getCommandDocument(final ByteBufferBsonOutput bsonOutput) {
186186
}
187187
}
188188

189-
BsonDocument getCommand() {
190-
return command;
191-
}
192-
193189
/**
194190
* Get the field name from a buffer positioned at the start of the document sequence identifier of an OP_MSG Section of type
195191
* `PAYLOAD_TYPE_1_DOCUMENT_SEQUENCE`.

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

Lines changed: 122 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,6 @@
5252
import com.mongodb.internal.session.SessionContext;
5353
import com.mongodb.internal.time.Timeout;
5454
import com.mongodb.internal.tracing.Span;
55-
import com.mongodb.internal.tracing.TraceContext;
5655
import com.mongodb.internal.tracing.TracingManager;
5756
import com.mongodb.lang.Nullable;
5857
import org.bson.BsonBinaryReader;
@@ -97,6 +96,18 @@
9796
import static com.mongodb.internal.connection.ProtocolHelper.isCommandOk;
9897
import static com.mongodb.internal.logging.LogMessage.Level.DEBUG;
9998
import static com.mongodb.internal.thread.InterruptionUtil.translateInterruptedException;
99+
import static com.mongodb.internal.tracing.Tags.CLIENT_CONNECTION_ID;
100+
import static com.mongodb.internal.tracing.Tags.CURSOR_ID;
101+
import static com.mongodb.internal.tracing.Tags.NAMESPACE;
102+
import static com.mongodb.internal.tracing.Tags.QUERY_SUMMARY;
103+
import static com.mongodb.internal.tracing.Tags.QUERY_TEXT;
104+
import static com.mongodb.internal.tracing.Tags.SERVER_ADDRESS;
105+
import static com.mongodb.internal.tracing.Tags.SERVER_CONNECTION_ID;
106+
import static com.mongodb.internal.tracing.Tags.SERVER_PORT;
107+
import static com.mongodb.internal.tracing.Tags.SERVER_TYPE;
108+
import static com.mongodb.internal.tracing.Tags.SESSION_ID;
109+
import static com.mongodb.internal.tracing.Tags.SYSTEM;
110+
import static com.mongodb.internal.tracing.Tags.TRANSACTION_NUMBER;
100111
import static java.util.Arrays.asList;
101112

102113
/**
@@ -377,24 +388,13 @@ public boolean isClosed() {
377388
public <T> T sendAndReceive(final CommandMessage message, final Decoder<T> decoder, final OperationContext operationContext) {
378389
Supplier<T> sendAndReceiveInternal = () -> sendAndReceiveInternal(
379390
message, decoder, operationContext);
380-
381-
Span tracingSpan = createTracingSpan(message, operationContext);
382-
383391
try {
384392
return sendAndReceiveInternal.get();
385-
} catch (MongoCommandException e) {
386-
if (tracingSpan != null) {
387-
tracingSpan.error(e);
388-
}
389-
393+
} catch (Throwable e) {
390394
if (reauthenticationIsTriggered(e)) {
391395
return reauthenticateAndRetry(sendAndReceiveInternal, operationContext);
392396
}
393397
throw e;
394-
} finally {
395-
if (tracingSpan != null) {
396-
tracingSpan.end();
397-
}
398398
}
399399
}
400400

@@ -406,9 +406,7 @@ public <T> void sendAndReceiveAsync(final CommandMessage message, final Decoder<
406406
AsyncSupplier<T> sendAndReceiveAsyncInternal = c -> sendAndReceiveAsyncInternal(
407407
message, decoder, operationContext, c);
408408

409-
beginAsync().<T>thenSupply(c -> {
410-
sendAndReceiveAsyncInternal.getAsync(c);
411-
}).onErrorIf(e -> reauthenticationIsTriggered(e), (t, c) -> {
409+
beginAsync().thenSupply(sendAndReceiveAsyncInternal::getAsync).onErrorIf(this::reauthenticationIsTriggered, (t, c) -> {
412410
reauthenticateAndRetryAsync(sendAndReceiveAsyncInternal, operationContext, c);
413411
}).finish(callback);
414412
}
@@ -447,15 +445,44 @@ public boolean reauthenticationIsTriggered(@Nullable final Throwable t) {
447445
private <T> T sendAndReceiveInternal(final CommandMessage message, final Decoder<T> decoder,
448446
final OperationContext operationContext) {
449447
CommandEventSender commandEventSender;
448+
450449
try (ByteBufferBsonOutput bsonOutput = new ByteBufferBsonOutput(this)) {
451450
message.encode(bsonOutput, operationContext);
452-
commandEventSender = createCommandEventSender(message, bsonOutput, operationContext);
453-
commandEventSender.sendStartedEvent();
451+
Span tracingSpan = createTracingSpan(message, operationContext, bsonOutput);
452+
453+
boolean isLoggingCommandNeeded = isLoggingCommandNeeded();
454+
boolean isTracingCommandPayloadNeeded = tracingSpan != null && operationContext.getTracingManager().isCommandPayloadEnabled();
455+
456+
// Only hydrate the command document if necessary
457+
BsonDocument commandDocument = null;
458+
if (isLoggingCommandNeeded || isTracingCommandPayloadNeeded) {
459+
commandDocument = message.getCommandDocument(bsonOutput);
460+
}
461+
if (isLoggingCommandNeeded) {
462+
commandEventSender = new LoggingCommandEventSender(
463+
SECURITY_SENSITIVE_COMMANDS, SECURITY_SENSITIVE_HELLO_COMMANDS, description, commandListener,
464+
operationContext, message, commandDocument,
465+
COMMAND_PROTOCOL_LOGGER, loggerSettings);
466+
commandEventSender.sendStartedEvent();
467+
} else {
468+
commandEventSender = new NoOpCommandEventSender();
469+
}
470+
if (isTracingCommandPayloadNeeded) {
471+
tracingSpan.tag(QUERY_TEXT, commandDocument.toJson());
472+
}
473+
454474
try {
455475
sendCommandMessage(message, bsonOutput, operationContext);
456476
} catch (Exception e) {
477+
if (tracingSpan != null) {
478+
tracingSpan.error(e);
479+
}
457480
commandEventSender.sendFailedEvent(e);
458481
throw e;
482+
} finally {
483+
if (tracingSpan != null) {
484+
tracingSpan.end();
485+
}
459486
}
460487
}
461488

@@ -568,7 +595,18 @@ private <T> void sendAndReceiveAsyncInternal(final CommandMessage message, final
568595

569596
try {
570597
message.encode(bsonOutput, operationContext);
571-
CommandEventSender commandEventSender = createCommandEventSender(message, bsonOutput, operationContext);
598+
599+
CommandEventSender commandEventSender;
600+
if (isLoggingCommandNeeded()) {
601+
BsonDocument commandDocument = message.getCommandDocument(bsonOutput);
602+
commandEventSender = new LoggingCommandEventSender(
603+
SECURITY_SENSITIVE_COMMANDS, SECURITY_SENSITIVE_HELLO_COMMANDS, description, commandListener,
604+
operationContext, message, commandDocument,
605+
COMMAND_PROTOCOL_LOGGER, loggerSettings);
606+
} else {
607+
commandEventSender = new NoOpCommandEventSender();
608+
}
609+
572610
commandEventSender.sendStartedEvent();
573611
Compressor localSendCompressor = sendCompressor;
574612
if (localSendCompressor == null || SECURITY_SENSITIVE_COMMANDS.contains(message.getCommandDocument(bsonOutput).getFirstKey())) {
@@ -887,42 +925,6 @@ public ByteBuf getBuffer(final int size) {
887925
return stream.getBuffer(size);
888926
}
889927

890-
@Nullable
891-
private Span createTracingSpan(final CommandMessage message, final OperationContext operationContext) {
892-
TracingManager tracingManager = operationContext.getTracingManager();
893-
Span span;
894-
if (tracingManager.isEnabled()) {
895-
BsonDocument command = message.getCommand();
896-
TraceContext parentContext = null;
897-
long cursorId = -1;
898-
if (command.containsKey("getMore")) {
899-
cursorId = command.getInt64("getMore").longValue();
900-
parentContext = tracingManager.getCursorParentContext(cursorId);
901-
} else {
902-
parentContext = tracingManager.getParentContext(operationContext.getId());
903-
}
904-
905-
span = tracingManager.addSpan("Command " + command.getFirstKey(), parentContext);
906-
span.tag("db.system", "mongodb");
907-
span.tag("db.namespace", message.getNamespace().getFullName());
908-
span.tag("db.query.summary", command.getFirstKey());
909-
span.tag("db.query.opcode", String.valueOf(message.getOpCode()));
910-
span.tag("db.query.text", command.toString());
911-
if (cursorId != -1) {
912-
span.tag("db.mongodb.cursor_id", String.valueOf(cursorId));
913-
}
914-
span.tag("server.address", serverId.getAddress().getHost());
915-
span.tag("server.port", String.valueOf(serverId.getAddress().getPort()));
916-
span.tag("server.type", message.getSettings().getServerType().name());
917-
918-
span.tag("db.mongodb.server_connection_id", this.description.getConnectionId().toString());
919-
} else {
920-
span = null;
921-
}
922-
923-
return span;
924-
}
925-
926928
private class MessageHeaderCallback implements SingleResultCallback<ByteBuf> {
927929
private final OperationContext operationContext;
928930
private final SingleResultCallback<ResponseBuffers> callback;
@@ -1003,19 +1005,75 @@ public void onResult(@Nullable final ByteBuf result, @Nullable final Throwable t
10031005

10041006
private static final StructuredLogger COMMAND_PROTOCOL_LOGGER = new StructuredLogger("protocol.command");
10051007

1006-
private CommandEventSender createCommandEventSender(final CommandMessage message, final ByteBufferBsonOutput bsonOutput,
1007-
final OperationContext operationContext) {
1008+
private boolean isLoggingCommandNeeded() {
10081009
boolean listensOrLogs = commandListener != null || COMMAND_PROTOCOL_LOGGER.isRequired(DEBUG, getClusterId());
1009-
if (!recordEverything && (isMonitoringConnection || !opened() || !authenticated.get() || !listensOrLogs)) {
1010-
return new NoOpCommandEventSender();
1011-
}
1012-
return new LoggingCommandEventSender(
1013-
SECURITY_SENSITIVE_COMMANDS, SECURITY_SENSITIVE_HELLO_COMMANDS, description, commandListener,
1014-
operationContext, message, bsonOutput,
1015-
COMMAND_PROTOCOL_LOGGER, loggerSettings);
1010+
return recordEverything || (!isMonitoringConnection && opened() && authenticated.get() && listensOrLogs);
10161011
}
10171012

10181013
private ClusterId getClusterId() {
10191014
return description.getConnectionId().getServerId().getClusterId();
10201015
}
1016+
1017+
/**
1018+
* Creates a tracing span for the given command message.
1019+
* <p>
1020+
* The span is only created if tracing is enabled and the command is not security-sensitive.
1021+
* It attaches various tags to the span, such as database system, namespace, query summary, opcode,
1022+
* server address, port, server type, client and server connection IDs, and, if applicable,
1023+
* transaction number and session ID. For cursor fetching commands, the parent context is retrieved using the cursor ID.
1024+
* If command payload tracing is enabled, the command document is also attached as a tag.
1025+
*
1026+
* @param message the command message to trace
1027+
* @param operationContext the operation context containing tracing and session information
1028+
* @param bsonOutput the BSON output used to serialize the command
1029+
* @return the created {@link Span}, or {@code null} if tracing is not enabled or the command is security-sensitive
1030+
*/
1031+
@Nullable
1032+
private Span createTracingSpan(final CommandMessage message, final OperationContext operationContext, final ByteBufferBsonOutput bsonOutput) {
1033+
1034+
TracingManager tracingManager = operationContext.getTracingManager();
1035+
BsonDocument command = message.getCommandDocument(bsonOutput);
1036+
1037+
// BsonDocument command = message.getCommand();
1038+
String commandName = command.getFirstKey();
1039+
// Span newSpan = tracingManager.addSpan("_____Command_____[ " + commandName + " ]", myparentContext);
1040+
if (!tracingManager.isEnabled()
1041+
|| SECURITY_SENSITIVE_COMMANDS.contains(commandName)
1042+
|| SECURITY_SENSITIVE_HELLO_COMMANDS.contains(commandName)) {
1043+
return null;
1044+
}
1045+
1046+
Span span = tracingManager
1047+
.addSpan("Command " + commandName, operationContext.getTracingSpanContext())
1048+
.tag(SYSTEM, "mongodb")
1049+
.tag(NAMESPACE, message.getNamespace().getDatabaseName())
1050+
.tag(QUERY_SUMMARY, commandName);
1051+
1052+
if (command.containsKey("getMore")) {
1053+
span.tag(CURSOR_ID, command.getInt64("getMore").longValue());
1054+
}
1055+
1056+
tagServerAndConnectionInfo(span, message);
1057+
tagSessionAndTransactionInfo(span, operationContext);
1058+
1059+
return span;
1060+
}
1061+
1062+
private void tagServerAndConnectionInfo(final Span span, final CommandMessage message) {
1063+
span.tag(SERVER_ADDRESS, serverId.getAddress().getHost())
1064+
.tag(SERVER_PORT, String.valueOf(serverId.getAddress().getPort()))
1065+
.tag(SERVER_TYPE, message.getSettings().getServerType().name())
1066+
.tag(CLIENT_CONNECTION_ID, this.description.getConnectionId().toString())
1067+
.tag(SERVER_CONNECTION_ID, String.valueOf(this.description.getConnectionId().getServerValue()));
1068+
}
1069+
1070+
private void tagSessionAndTransactionInfo(final Span span, final OperationContext operationContext) {
1071+
SessionContext sessionContext = operationContext.getSessionContext();
1072+
if (sessionContext.hasSession() && !sessionContext.isImplicitSession()) {
1073+
span.tag(TRANSACTION_NUMBER, String.valueOf(sessionContext.getTransactionNumber()))
1074+
.tag(SESSION_ID, String.valueOf(sessionContext.getSessionId()
1075+
.get(sessionContext.getSessionId().getFirstKey())
1076+
.asBinary().asUuid()));
1077+
}
1078+
}
10211079
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ class LoggingCommandEventSender implements CommandEventSender {
7878
@Nullable final CommandListener commandListener,
7979
final OperationContext operationContext,
8080
final CommandMessage message,
81-
final ByteBufferBsonOutput bsonOutput,
81+
final BsonDocument commandDocument,
8282
final StructuredLogger logger,
8383
final LoggerSettings loggerSettings) {
8484
this.description = description;
@@ -88,7 +88,7 @@ class LoggingCommandEventSender implements CommandEventSender {
8888
this.loggerSettings = loggerSettings;
8989
this.startTimeNanos = System.nanoTime();
9090
this.message = message;
91-
this.commandDocument = message.getCommandDocument(bsonOutput);
91+
this.commandDocument = commandDocument;
9292
this.commandName = commandDocument.getFirstKey();
9393
this.redactionRequired = securitySensitiveCommands.contains(commandName)
9494
|| (securitySensitiveHelloCommands.contains(commandName) && commandDocument.containsKey("speculativeAuthenticate"));

0 commit comments

Comments
 (0)