Skip to content

Make SDK tracer span current around gRPC newCall #685

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
KirillKurdyukov merged 7 commits into from
Jul 23, 2026
Merged

Make SDK tracer span current around gRPC newCall #685

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
c6f0652
Make SDK tracer span current around gRPC newCall
KirillKurdyukov Jul 23, 2026
c8692fa
Make SDK tracer span current around gRPC newCall
KirillKurdyukov Jul 23, 2026
5395f3f
Test that GrpcTelemetry spans nest under SDK tracer spans
KirillKurdyukov Jul 23, 2026
4b754fa
Test that GrpcTelemetry spans nest under SDK tracer spans
KirillKurdyukov Jul 23, 2026
5ddb0f1
Test that GrpcTelemetry spans nest under SDK tracer spans
KirillKurdyukov Jul 23, 2026
8b9e7ca
rollback otel tests
KirillKurdyukov Jul 23, 2026
f7b5435
rollback otel tests
KirillKurdyukov Jul 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
158 changes: 86 additions & 72 deletions query/src/main/java/tech/ydb/query/impl/SessionImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import tech.ydb.core.grpc.GrpcRequestSettings;
import tech.ydb.core.operation.StatusExtractor;
import tech.ydb.core.settings.BaseRequestSettings;
import tech.ydb.core.tracing.Scope;
import tech.ydb.core.tracing.Span;
import tech.ydb.core.utils.URITools;
import tech.ydb.core.utils.UpdatableOptional;
Expand Down Expand Up @@ -126,10 +127,14 @@ public CompletableFuture<Result<QueryTransaction>> beginTransaction(TxMode tx, B
.setTxSettings(TxControl.txSettings(tx))
.build();

return rpc.beginTransaction(request, makeOptions(settings).build()).thenApply(result -> {
updateSessionState(result.getStatus());
return result.map(resp -> updateTransaction(new TransactionImpl(tx, resp.getTxMeta().getId())));
});
Span span = startSpan("ydb.BeginTransaction");
try (Scope ignored = span.makeCurrent()) {
return Span.endOnResult(span, rpc.beginTransaction(request, makeOptions(settings, span).build()))
.thenApply(result -> {
updateSessionState(result.getStatus());
return result.map(resp -> updateTransaction(new TransactionImpl(tx, resp.getTxMeta().getId())));
});
}
}

private QueryTransaction updateTransaction(TransactionImpl newTx) {
Expand Down Expand Up @@ -331,14 +336,16 @@ GrpcReadStream<YdbQuery.ExecuteQueryResponsePart> createGrpcStream(
public QueryStream createQuery(String query, TxMode tx, Params prms, ExecuteQuerySettings settings) {
YdbQuery.TransactionControl tc = TxControl.txModeCtrl(tx, true);
Span span = startSpan("ydb.ExecuteQuery");
return new StreamImpl(createGrpcStream(query, tc, prms, settings, span), span) {
@Override
void handleTxMeta(String txID) {
if (txID != null && !txID.isEmpty()) {
logger.warn("{} got unexpected transaction id {}", SessionImpl.this, txID);
try (Scope ignored = span.makeCurrent()) {
return new StreamImpl(createGrpcStream(query, tc, prms, settings, span), span) {
@Override
void handleTxMeta(String txID) {
if (txID != null && !txID.isEmpty()) {
logger.warn("{} got unexpected transaction id {}", SessionImpl.this, txID);
}
}
}
};
};
}
}

public CompletableFuture<Result<YdbQuery.DeleteSessionResponse>> delete(DeleteSessionSettings settings) {
Expand Down Expand Up @@ -478,44 +485,46 @@ public QueryStream createQuery(String query, boolean commitAtEnd, Params prms, E
: TxControl.txModeCtrl(txMode, commitAtEnd);

Span span = startSpan("ydb.ExecuteQuery");
return new StreamImpl(createGrpcStream(query, tc, prms, settings, span), span) {
@Override
void handleTxMeta(String txID) {
String newId = txID == null || txID.isEmpty() ? null : txID;
if (!txId.compareAndSet(currentId, newId)) {
logger.warn("{} lost transaction meta id {}", SessionImpl.this, newId);
try (Scope ignored = span.makeCurrent()) {
return new StreamImpl(createGrpcStream(query, tc, prms, settings, span), span) {
@Override
void handleTxMeta(String txID) {
String newId = txID == null || txID.isEmpty() ? null : txID;
if (!txId.compareAndSet(currentId, newId)) {
logger.warn("{} lost transaction meta id {}", SessionImpl.this, newId);
}
}
}

@Override
void handleCompletion(Status status, Throwable th) {
if (th != null) {
currentStatusFuture.completeExceptionally(
new RuntimeException("Query on transaction failed with exception ", th));
}
if (status.isSuccess()) {
if (commitAtEnd) {
currentStatusFuture.complete(Status.SUCCESS);
@Override
void handleCompletion(Status status, Throwable th) {
if (th != null) {
currentStatusFuture.completeExceptionally(
new RuntimeException("Query on transaction failed with exception ", th));
}
} else {
if (txId.compareAndSet(currentId, null)) {
logger.warn("{} transaction with id {} was failed", SessionImpl.this, currentId);
if (status.isSuccess()) {
if (commitAtEnd) {
currentStatusFuture.complete(Status.SUCCESS);
}
} else {
if (txId.compareAndSet(currentId, null)) {
logger.warn("{} transaction with id {} was failed", SessionImpl.this, currentId);
}
currentStatusFuture.complete(Status
.of(StatusCode.ABORTED)
.withIssues(Issue.of("Query on transaction failed with status "
+ status, Issue.Severity.ERROR)));
}
currentStatusFuture.complete(Status
.of(StatusCode.ABORTED)
.withIssues(Issue.of("Query on transaction failed with status "
+ status, Issue.Severity.ERROR)));
}
}

@Override
public void cancel() {
super.cancel();
if (txId.compareAndSet(currentId, null)) {
logger.warn("{} transaction with id {} was cancelled", SessionImpl.this, currentId);
@Override
public void cancel() {
super.cancel();
if (txId.compareAndSet(currentId, null)) {
logger.warn("{} transaction with id {} was cancelled", SessionImpl.this, currentId);
}
}
}
};
};
}
}

@Override
Expand All @@ -534,22 +543,25 @@ public CompletableFuture<Result<QueryInfo>> commit(CommitTransactionSettings set
.setTxId(transactionId)
.build();

return Span.endOnResult(span, rpc.commitTransaction(request, makeOptions(settings, span).build()))
.thenApply(res -> {
Status status = res.getStatus();
currentStatusFuture.complete(status);
updateSessionState(status);
if (!txId.compareAndSet(transactionId, null)) {
logger.warn("{} lost commit response for transaction {}", SessionImpl.this, transactionId);
}
// TODO: CommitTransactionResponse must contain exec_stats
return res.map(resp -> new QueryInfo(null));
}).whenComplete(((status, th) -> {
if (th != null) {
currentStatusFuture.completeExceptionally(
new RuntimeException("Transaction commit failed with exception", th));
}
}));
try (Scope ignored = span.makeCurrent()) {
return Span.endOnResult(span, rpc.commitTransaction(request, makeOptions(settings, span).build()))
.thenApply(res -> {
Status status = res.getStatus();
currentStatusFuture.complete(status);
updateSessionState(status);
if (!txId.compareAndSet(transactionId, null)) {
logger.warn("{} lost commit response for transaction {}", SessionImpl.this,
transactionId);
}
// TODO: CommitTransactionResponse must contain exec_stats
return res.map(resp -> new QueryInfo(null));
}).whenComplete(((status, th) -> {
if (th != null) {
currentStatusFuture.completeExceptionally(
new RuntimeException("Transaction commit failed with exception", th));
}
}));
}
}

@Override
Expand All @@ -568,20 +580,22 @@ public CompletableFuture<Status> rollback(RollbackTransactionSettings settings)
.setSessionId(sessionId)
.setTxId(transactionId)
.build();
return Span.endOnResult(span, rpc.rollbackTransaction(request, makeOptions(settings, span).build()))
.thenApply(result -> {
updateSessionState(result.getStatus());
if (!txId.compareAndSet(transactionId, null)) {
logger.warn("{} lost rollback response for transaction {}", SessionImpl.this,
transactionId);
}
return result.getStatus();
})
.whenComplete((status, th) -> {
currentStatusFuture.complete(Status
.of(StatusCode.ABORTED)
.withIssues(Issue.of("Transaction was rolled back", Issue.Severity.ERROR)));
});
try (Scope ignored = span.makeCurrent()) {
return Span.endOnResult(span, rpc.rollbackTransaction(request, makeOptions(settings, span).build()))
.thenApply(result -> {
updateSessionState(result.getStatus());
if (!txId.compareAndSet(transactionId, null)) {
logger.warn("{} lost rollback response for transaction {}", SessionImpl.this,
transactionId);
}
return result.getStatus();
})
.whenComplete((status, th) -> {
currentStatusFuture.complete(Status
.of(StatusCode.ABORTED)
.withIssues(Issue.of("Transaction was rolled back", Issue.Severity.ERROR)));
});
}
}
}
}
34 changes: 20 additions & 14 deletions query/src/main/java/tech/ydb/query/impl/SessionPool.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import tech.ydb.core.UnexpectedResultException;
import tech.ydb.core.grpc.GrpcReadStream;
import tech.ydb.core.metrics.Meter;
import tech.ydb.core.tracing.Scope;
import tech.ydb.core.tracing.Span;
import tech.ydb.core.utils.FutureTools;
import tech.ydb.proto.query.YdbQuery;
Expand Down Expand Up @@ -286,20 +287,25 @@ public CompletableFuture<PooledQuerySession> create() {
long startNanos = System.nanoTime();
stats.requested.increment();
metrics.onSessionRequested();
return Span.endOnResult(createSpan, SessionImpl.createSession(rpc, CREATE_SETTINGS, true, createSpan))
.thenCompose(r -> {
metrics.onCreateTime(System.nanoTime() - startNanos);

if (!r.isSuccess()) {
stats.failed.increment();
metrics.onSessionFailed(r.getStatus());
throw new UnexpectedResultException("create session problem", r.getStatus());
}
metrics.onSessionCreated();
PooledQuerySession session = new PooledQuerySession(rpc, r.getValue());
return session.start();
})
.thenApply(Result::getValue);
try (Scope ignored = createSpan.makeCurrent()) {
return Span.endOnResult(
createSpan,
SessionImpl.createSession(rpc, CREATE_SETTINGS, true, createSpan)
)
.thenCompose(r -> {
metrics.onCreateTime(System.nanoTime() - startNanos);

if (!r.isSuccess()) {
stats.failed.increment();
metrics.onSessionFailed(r.getStatus());
throw new UnexpectedResultException("create session problem", r.getStatus());
}
metrics.onSessionCreated();
PooledQuerySession session = new PooledQuerySession(rpc, r.getValue());
return session.start();
})
.thenApply(Result::getValue);
}
} finally {
ctx.detach(previous);
}
Expand Down
90 changes: 50 additions & 40 deletions query/src/main/java/tech/ydb/query/impl/TableClientImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
import tech.ydb.core.StatusCode;
import tech.ydb.core.UnexpectedResultException;
import tech.ydb.core.grpc.GrpcTransport;
import tech.ydb.core.tracing.Scope;
import tech.ydb.core.tracing.Span;
import tech.ydb.core.tracing.Tracer;
import tech.ydb.proto.ValueProtos;
Expand Down Expand Up @@ -136,49 +137,54 @@ public CompletableFuture<Result<DataQueryResult>> executeDataQueryInternal(
final List<ValueProtos.ResultSet> results = new ArrayList<>();
Span span = querySession.startSpan("ydb.ExecuteQuery");

QueryStream stream = querySession.new StreamImpl(querySession.createGrpcStream(query, tc, prms, qs, span),
span) {
@Override
void handleTxMeta(String txID) {
txRef.set(txID);
}
};

CompletableFuture<Result<QueryInfo>> future = stream.execute(new QueryStream.PartsHandler() {
@Override
public void onIssues(Issue[] issueArr) {
issues.addAll(Arrays.asList(issueArr));
}

@Override
public void onNextPart(QueryResultPart part) {
} // not used
try (Scope ignored = span.makeCurrent()) {
QueryStream stream = querySession.new StreamImpl(
querySession.createGrpcStream(query, tc, prms, qs, span),
span
) {
@Override
void handleTxMeta(String txID) {
txRef.set(txID);
}
};

@Override
public void onNextRawPart(long index, ValueProtos.ResultSet rs) {
int idx = (int) index;
while (results.size() <= idx) {
results.add(null);
CompletableFuture<Result<QueryInfo>> future = stream.execute(new QueryStream.PartsHandler() {
@Override
public void onIssues(Issue[] issueArr) {
issues.addAll(Arrays.asList(issueArr));
}
if (results.get(idx) == null) {
results.set(idx, rs);
} else {
results.set(idx, results.get(idx).toBuilder().addAllRows(rs.getRowsList()).build());

@Override
public void onNextPart(QueryResultPart part) {
} // not used

@Override
public void onNextRawPart(long index, ValueProtos.ResultSet rs) {
int idx = (int) index;
while (results.size() <= idx) {
results.add(null);
}
if (results.get(idx) == null) {
results.set(idx, rs);
} else {
results.set(idx, results.get(idx).toBuilder().addAllRows(rs.getRowsList()).build());
}
}
}
});
});


return future.thenApply(res -> {
if (!res.isSuccess()) {
return res.map(v -> null);
}
QueryStats stats = res.getValue().getStats();
String txId = txRef.get();
Status status = res.getStatus().withIssues(issues.toArray(new Issue[0]));
DataQueryResult value = new DataQueryResult(txId, results, stats != null ? stats.toProtobuf() : null);
return Result.success(value, status);
});
return future.thenApply(res -> {
if (!res.isSuccess()) {
return res.map(v -> null);
}
QueryStats stats = res.getValue().getStats();
String txId = txRef.get();
Status status = res.getStatus().withIssues(issues.toArray(new Issue[0]));
DataQueryResult value = new DataQueryResult(
txId, results, stats != null ? stats.toProtobuf() : null);
return Result.success(value, status);
});
}
}

@Override
Expand Down Expand Up @@ -214,7 +220,9 @@ protected CompletableFuture<Status> commitTransactionInternal(String txId, Commi
.withTraceId(settings.getTraceId())
.withRequestTimeout(settings.getTimeoutDuration())
.build();
return Span.endOnStatus(span, querySession.commitById(txId, querySettings, span));
try (Scope ignored = span.makeCurrent()) {
return Span.endOnStatus(span, querySession.commitById(txId, querySettings, span));
}
}

@Override
Expand All @@ -224,7 +232,9 @@ protected CompletableFuture<Status> rollbackTransactionInternal(String txId, Rol
.withTraceId(settings.getTraceId())
.withRequestTimeout(settings.getTimeoutDuration())
.build();
return Span.endOnStatus(span, querySession.rollbackById(txId, querySettings, span));
try (Scope ignored = span.makeCurrent()) {
return Span.endOnStatus(span, querySession.rollbackById(txId, querySettings, span));
}
}

private final class TracedTableTransaction implements TableTransaction {
Expand Down
Loading