diff --git a/query/src/main/java/tech/ydb/query/impl/SessionImpl.java b/query/src/main/java/tech/ydb/query/impl/SessionImpl.java index 3334f3e05..1f3cbaa54 100644 --- a/query/src/main/java/tech/ydb/query/impl/SessionImpl.java +++ b/query/src/main/java/tech/ydb/query/impl/SessionImpl.java @@ -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; @@ -126,10 +127,14 @@ public CompletableFuture> 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) { @@ -331,14 +336,16 @@ GrpcReadStream 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> delete(DeleteSessionSettings settings) { @@ -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 @@ -534,22 +543,25 @@ public CompletableFuture> 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 @@ -568,20 +580,22 @@ public CompletableFuture 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))); + }); + } } } } diff --git a/query/src/main/java/tech/ydb/query/impl/SessionPool.java b/query/src/main/java/tech/ydb/query/impl/SessionPool.java index a7743c079..497fb4fa0 100644 --- a/query/src/main/java/tech/ydb/query/impl/SessionPool.java +++ b/query/src/main/java/tech/ydb/query/impl/SessionPool.java @@ -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; @@ -286,20 +287,25 @@ public CompletableFuture 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); } diff --git a/query/src/main/java/tech/ydb/query/impl/TableClientImpl.java b/query/src/main/java/tech/ydb/query/impl/TableClientImpl.java index 27ec97e6f..93401cb5b 100644 --- a/query/src/main/java/tech/ydb/query/impl/TableClientImpl.java +++ b/query/src/main/java/tech/ydb/query/impl/TableClientImpl.java @@ -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; @@ -136,49 +137,54 @@ public CompletableFuture> executeDataQueryInternal( final List 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> 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> 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 @@ -214,7 +220,9 @@ protected CompletableFuture 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 @@ -224,7 +232,9 @@ protected CompletableFuture 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 {