Skip to content

Commit 5b09213

Browse files
KirillKurdyukovKirillKurdyukov
authored
Make SDK tracer span current around gRPC newCall (#685)
1 parent d0d6ac1 commit 5b09213

3 files changed

Lines changed: 156 additions & 126 deletions

File tree

‎query/src/main/java/tech/ydb/query/impl/SessionImpl.java‎

Lines changed: 86 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import tech.ydb.core.grpc.GrpcRequestSettings;
2727
import tech.ydb.core.operation.StatusExtractor;
2828
import tech.ydb.core.settings.BaseRequestSettings;
29+
import tech.ydb.core.tracing.Scope;
2930
import tech.ydb.core.tracing.Span;
3031
import tech.ydb.core.utils.URITools;
3132
import tech.ydb.core.utils.UpdatableOptional;
@@ -126,10 +127,14 @@ public CompletableFuture<Result<QueryTransaction>> beginTransaction(TxMode tx, B
126127
.setTxSettings(TxControl.txSettings(tx))
127128
.build();
128129

129-
return rpc.beginTransaction(request, makeOptions(settings).build()).thenApply(result -> {
130-
updateSessionState(result.getStatus());
131-
return result.map(resp -> updateTransaction(new TransactionImpl(tx, resp.getTxMeta().getId())));
132-
});
130+
Span span = startSpan("ydb.BeginTransaction");
131+
try (Scope ignored = span.makeCurrent()) {
132+
return Span.endOnResult(span, rpc.beginTransaction(request, makeOptions(settings, span).build()))
133+
.thenApply(result -> {
134+
updateSessionState(result.getStatus());
135+
return result.map(resp -> updateTransaction(new TransactionImpl(tx, resp.getTxMeta().getId())));
136+
});
137+
}
133138
}
134139

135140
private QueryTransaction updateTransaction(TransactionImpl newTx) {
@@ -331,14 +336,16 @@ GrpcReadStream<YdbQuery.ExecuteQueryResponsePart> createGrpcStream(
331336
public QueryStream createQuery(String query, TxMode tx, Params prms, ExecuteQuerySettings settings) {
332337
YdbQuery.TransactionControl tc = TxControl.txModeCtrl(tx, true);
333338
Span span = startSpan("ydb.ExecuteQuery");
334-
return new StreamImpl(createGrpcStream(query, tc, prms, settings, span), span) {
335-
@Override
336-
void handleTxMeta(String txID) {
337-
if (txID != null && !txID.isEmpty()) {
338-
logger.warn("{} got unexpected transaction id {}", SessionImpl.this, txID);
339+
try (Scope ignored = span.makeCurrent()) {
340+
return new StreamImpl(createGrpcStream(query, tc, prms, settings, span), span) {
341+
@Override
342+
void handleTxMeta(String txID) {
343+
if (txID != null && !txID.isEmpty()) {
344+
logger.warn("{} got unexpected transaction id {}", SessionImpl.this, txID);
345+
}
339346
}
340-
}
341-
};
347+
};
348+
}
342349
}
343350

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

480487
Span span = startSpan("ydb.ExecuteQuery");
481-
return new StreamImpl(createGrpcStream(query, tc, prms, settings, span), span) {
482-
@Override
483-
void handleTxMeta(String txID) {
484-
String newId = txID == null || txID.isEmpty() ? null : txID;
485-
if (!txId.compareAndSet(currentId, newId)) {
486-
logger.warn("{} lost transaction meta id {}", SessionImpl.this, newId);
488+
try (Scope ignored = span.makeCurrent()) {
489+
return new StreamImpl(createGrpcStream(query, tc, prms, settings, span), span) {
490+
@Override
491+
void handleTxMeta(String txID) {
492+
String newId = txID == null || txID.isEmpty() ? null : txID;
493+
if (!txId.compareAndSet(currentId, newId)) {
494+
logger.warn("{} lost transaction meta id {}", SessionImpl.this, newId);
495+
}
487496
}
488-
}
489497

490-
@Override
491-
void handleCompletion(Status status, Throwable th) {
492-
if (th != null) {
493-
currentStatusFuture.completeExceptionally(
494-
new RuntimeException("Query on transaction failed with exception ", th));
495-
}
496-
if (status.isSuccess()) {
497-
if (commitAtEnd) {
498-
currentStatusFuture.complete(Status.SUCCESS);
498+
@Override
499+
void handleCompletion(Status status, Throwable th) {
500+
if (th != null) {
501+
currentStatusFuture.completeExceptionally(
502+
new RuntimeException("Query on transaction failed with exception ", th));
499503
}
500-
} else {
501-
if (txId.compareAndSet(currentId, null)) {
502-
logger.warn("{} transaction with id {} was failed", SessionImpl.this, currentId);
504+
if (status.isSuccess()) {
505+
if (commitAtEnd) {
506+
currentStatusFuture.complete(Status.SUCCESS);
507+
}
508+
} else {
509+
if (txId.compareAndSet(currentId, null)) {
510+
logger.warn("{} transaction with id {} was failed", SessionImpl.this, currentId);
511+
}
512+
currentStatusFuture.complete(Status
513+
.of(StatusCode.ABORTED)
514+
.withIssues(Issue.of("Query on transaction failed with status "
515+
+ status, Issue.Severity.ERROR)));
503516
}
504-
currentStatusFuture.complete(Status
505-
.of(StatusCode.ABORTED)
506-
.withIssues(Issue.of("Query on transaction failed with status "
507-
+ status, Issue.Severity.ERROR)));
508517
}
509-
}
510518

511-
@Override
512-
public void cancel() {
513-
super.cancel();
514-
if (txId.compareAndSet(currentId, null)) {
515-
logger.warn("{} transaction with id {} was cancelled", SessionImpl.this, currentId);
519+
@Override
520+
public void cancel() {
521+
super.cancel();
522+
if (txId.compareAndSet(currentId, null)) {
523+
logger.warn("{} transaction with id {} was cancelled", SessionImpl.this, currentId);
524+
}
516525
}
517-
}
518-
};
526+
};
527+
}
519528
}
520529

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

537-
return Span.endOnResult(span, rpc.commitTransaction(request, makeOptions(settings, span).build()))
538-
.thenApply(res -> {
539-
Status status = res.getStatus();
540-
currentStatusFuture.complete(status);
541-
updateSessionState(status);
542-
if (!txId.compareAndSet(transactionId, null)) {
543-
logger.warn("{} lost commit response for transaction {}", SessionImpl.this, transactionId);
544-
}
545-
// TODO: CommitTransactionResponse must contain exec_stats
546-
return res.map(resp -> new QueryInfo(null));
547-
}).whenComplete(((status, th) -> {
548-
if (th != null) {
549-
currentStatusFuture.completeExceptionally(
550-
new RuntimeException("Transaction commit failed with exception", th));
551-
}
552-
}));
546+
try (Scope ignored = span.makeCurrent()) {
547+
return Span.endOnResult(span, rpc.commitTransaction(request, makeOptions(settings, span).build()))
548+
.thenApply(res -> {
549+
Status status = res.getStatus();
550+
currentStatusFuture.complete(status);
551+
updateSessionState(status);
552+
if (!txId.compareAndSet(transactionId, null)) {
553+
logger.warn("{} lost commit response for transaction {}", SessionImpl.this,
554+
transactionId);
555+
}
556+
// TODO: CommitTransactionResponse must contain exec_stats
557+
return res.map(resp -> new QueryInfo(null));
558+
}).whenComplete(((status, th) -> {
559+
if (th != null) {
560+
currentStatusFuture.completeExceptionally(
561+
new RuntimeException("Transaction commit failed with exception", th));
562+
}
563+
}));
564+
}
553565
}
554566

555567
@Override
@@ -568,20 +580,22 @@ public CompletableFuture<Status> rollback(RollbackTransactionSettings settings)
568580
.setSessionId(sessionId)
569581
.setTxId(transactionId)
570582
.build();
571-
return Span.endOnResult(span, rpc.rollbackTransaction(request, makeOptions(settings, span).build()))
572-
.thenApply(result -> {
573-
updateSessionState(result.getStatus());
574-
if (!txId.compareAndSet(transactionId, null)) {
575-
logger.warn("{} lost rollback response for transaction {}", SessionImpl.this,
576-
transactionId);
577-
}
578-
return result.getStatus();
579-
})
580-
.whenComplete((status, th) -> {
581-
currentStatusFuture.complete(Status
582-
.of(StatusCode.ABORTED)
583-
.withIssues(Issue.of("Transaction was rolled back", Issue.Severity.ERROR)));
584-
});
583+
try (Scope ignored = span.makeCurrent()) {
584+
return Span.endOnResult(span, rpc.rollbackTransaction(request, makeOptions(settings, span).build()))
585+
.thenApply(result -> {
586+
updateSessionState(result.getStatus());
587+
if (!txId.compareAndSet(transactionId, null)) {
588+
logger.warn("{} lost rollback response for transaction {}", SessionImpl.this,
589+
transactionId);
590+
}
591+
return result.getStatus();
592+
})
593+
.whenComplete((status, th) -> {
594+
currentStatusFuture.complete(Status
595+
.of(StatusCode.ABORTED)
596+
.withIssues(Issue.of("Transaction was rolled back", Issue.Severity.ERROR)));
597+
});
598+
}
585599
}
586600
}
587601
}

‎query/src/main/java/tech/ydb/query/impl/SessionPool.java‎

Lines changed: 20 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import tech.ydb.core.grpc.GrpcReadStream;
2525
import tech.ydb.core.impl.Observability;
2626
import tech.ydb.core.metrics.Meter;
27+
import tech.ydb.core.tracing.Scope;
2728
import tech.ydb.core.tracing.Span;
2829
import tech.ydb.core.utils.FutureTools;
2930
import tech.ydb.proto.query.YdbQuery;
@@ -288,20 +289,25 @@ public CompletableFuture<PooledQuerySession> create() {
288289
long startNanos = System.nanoTime();
289290
stats.requested.increment();
290291
metrics.onSessionRequested();
291-
return Span.endOnResult(createSpan, SessionImpl.createSession(rpc, CREATE_SETTINGS, true, createSpan))
292-
.thenCompose(r -> {
293-
metrics.onCreateTime(System.nanoTime() - startNanos);
294-
295-
if (!r.isSuccess()) {
296-
stats.failed.increment();
297-
metrics.onSessionFailed(r.getStatus());
298-
throw new UnexpectedResultException("create session problem", r.getStatus());
299-
}
300-
metrics.onSessionCreated();
301-
PooledQuerySession session = new PooledQuerySession(rpc, r.getValue());
302-
return session.start();
303-
})
304-
.thenApply(Result::getValue);
292+
try (Scope ignored = createSpan.makeCurrent()) {
293+
return Span.endOnResult(
294+
createSpan,
295+
SessionImpl.createSession(rpc, CREATE_SETTINGS, true, createSpan)
296+
)
297+
.thenCompose(r -> {
298+
metrics.onCreateTime(System.nanoTime() - startNanos);
299+
300+
if (!r.isSuccess()) {
301+
stats.failed.increment();
302+
metrics.onSessionFailed(r.getStatus());
303+
throw new UnexpectedResultException("create session problem", r.getStatus());
304+
}
305+
metrics.onSessionCreated();
306+
PooledQuerySession session = new PooledQuerySession(rpc, r.getValue());
307+
return session.start();
308+
})
309+
.thenApply(Result::getValue);
310+
}
305311
} finally {
306312
ctx.detach(previous);
307313
}

‎query/src/main/java/tech/ydb/query/impl/TableClientImpl.java‎

Lines changed: 50 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import tech.ydb.core.StatusCode;
1616
import tech.ydb.core.UnexpectedResultException;
1717
import tech.ydb.core.grpc.GrpcTransport;
18+
import tech.ydb.core.tracing.Scope;
1819
import tech.ydb.core.tracing.Span;
1920
import tech.ydb.core.tracing.Tracer;
2021
import tech.ydb.proto.ValueProtos;
@@ -136,49 +137,54 @@ public CompletableFuture<Result<DataQueryResult>> executeDataQueryInternal(
136137
final List<ValueProtos.ResultSet> results = new ArrayList<>();
137138
Span span = querySession.startSpan("ydb.ExecuteQuery");
138139

139-
QueryStream stream = querySession.new StreamImpl(querySession.createGrpcStream(query, tc, prms, qs, span),
140-
span) {
141-
@Override
142-
void handleTxMeta(String txID) {
143-
txRef.set(txID);
144-
}
145-
};
146-
147-
CompletableFuture<Result<QueryInfo>> future = stream.execute(new QueryStream.PartsHandler() {
148-
@Override
149-
public void onIssues(Issue[] issueArr) {
150-
issues.addAll(Arrays.asList(issueArr));
151-
}
152-
153-
@Override
154-
public void onNextPart(QueryResultPart part) {
155-
} // not used
140+
try (Scope ignored = span.makeCurrent()) {
141+
QueryStream stream = querySession.new StreamImpl(
142+
querySession.createGrpcStream(query, tc, prms, qs, span),
143+
span
144+
) {
145+
@Override
146+
void handleTxMeta(String txID) {
147+
txRef.set(txID);
148+
}
149+
};
156150

157-
@Override
158-
public void onNextRawPart(long index, ValueProtos.ResultSet rs) {
159-
int idx = (int) index;
160-
while (results.size() <= idx) {
161-
results.add(null);
151+
CompletableFuture<Result<QueryInfo>> future = stream.execute(new QueryStream.PartsHandler() {
152+
@Override
153+
public void onIssues(Issue[] issueArr) {
154+
issues.addAll(Arrays.asList(issueArr));
162155
}
163-
if (results.get(idx) == null) {
164-
results.set(idx, rs);
165-
} else {
166-
results.set(idx, results.get(idx).toBuilder().addAllRows(rs.getRowsList()).build());
156+
157+
@Override
158+
public void onNextPart(QueryResultPart part) {
159+
} // not used
160+
161+
@Override
162+
public void onNextRawPart(long index, ValueProtos.ResultSet rs) {
163+
int idx = (int) index;
164+
while (results.size() <= idx) {
165+
results.add(null);
166+
}
167+
if (results.get(idx) == null) {
168+
results.set(idx, rs);
169+
} else {
170+
results.set(idx, results.get(idx).toBuilder().addAllRows(rs.getRowsList()).build());
171+
}
167172
}
168-
}
169-
});
173+
});
170174

171175

172-
return future.thenApply(res -> {
173-
if (!res.isSuccess()) {
174-
return res.map(v -> null);
175-
}
176-
QueryStats stats = res.getValue().getStats();
177-
String txId = txRef.get();
178-
Status status = res.getStatus().withIssues(issues.toArray(new Issue[0]));
179-
DataQueryResult value = new DataQueryResult(txId, results, stats != null ? stats.toProtobuf() : null);
180-
return Result.success(value, status);
181-
});
176+
return future.thenApply(res -> {
177+
if (!res.isSuccess()) {
178+
return res.map(v -> null);
179+
}
180+
QueryStats stats = res.getValue().getStats();
181+
String txId = txRef.get();
182+
Status status = res.getStatus().withIssues(issues.toArray(new Issue[0]));
183+
DataQueryResult value = new DataQueryResult(
184+
txId, results, stats != null ? stats.toProtobuf() : null);
185+
return Result.success(value, status);
186+
});
187+
}
182188
}
183189

184190
@Override
@@ -214,7 +220,9 @@ protected CompletableFuture<Status> commitTransactionInternal(String txId, Commi
214220
.withTraceId(settings.getTraceId())
215221
.withRequestTimeout(settings.getTimeoutDuration())
216222
.build();
217-
return Span.endOnStatus(span, querySession.commitById(txId, querySettings, span));
223+
try (Scope ignored = span.makeCurrent()) {
224+
return Span.endOnStatus(span, querySession.commitById(txId, querySettings, span));
225+
}
218226
}
219227

220228
@Override
@@ -224,7 +232,9 @@ protected CompletableFuture<Status> rollbackTransactionInternal(String txId, Rol
224232
.withTraceId(settings.getTraceId())
225233
.withRequestTimeout(settings.getTimeoutDuration())
226234
.build();
227-
return Span.endOnStatus(span, querySession.rollbackById(txId, querySettings, span));
235+
try (Scope ignored = span.makeCurrent()) {
236+
return Span.endOnStatus(span, querySession.rollbackById(txId, querySettings, span));
237+
}
228238
}
229239

230240
private final class TracedTableTransaction implements TableTransaction {

0 commit comments

Comments
 (0)