feat: add RetryCtx span propagation

core/src/main/java/tech/ydb/core/grpc/GrpcRequestSettings.java

@@ -6,8 +6,6 @@
 import java.util.function.BooleanSupplier;
 import java.util.function.Consumer;
 
-import javax.annotation.Nullable;
-
 import io.grpc.Metadata;
 
 import tech.ydb.core.impl.call.GrpcFlows;
@@ -87,7 +85,6 @@ public GrpcFlowControl getFlowControl() {
         return flowControl;
     }
 
-    @Nullable
     public Span getSpan() {
         return span;
     }
@@ -103,7 +100,7 @@ public static final class Builder {
         private Consumer<Metadata> trailersHandler = null;
         private BooleanSupplier pessimizationHook = null;
         private GrpcFlowControl flowControl = GrpcFlows.SIMPLE_FLOW;
-        private Span span = null;
+        private Span span = Span.NOOP;
 
         /**
          * Returns a new {@code Builder} with a deadline, based on the running Java Virtual Machine's
@@ -185,7 +182,7 @@ public Builder withPessimizationHook(BooleanSupplier pessimizationHook) {
         }
 
         public Builder withSpan(Span span) {
-            this.span = span;
+            this.span = span == null ? Span.NOOP : span;
             return this;
         }
 

core/src/main/java/tech/ydb/core/grpc/GrpcTransport.java

@@ -43,7 +43,7 @@ <ReqT, RespT> GrpcReadWriteStream<RespT, ReqT> readWriteStreamCall(
     ScheduledExecutorService getScheduler();
 
     default Tracer getTracer() {
-        return NoopTracer.INSTANCE;
+        return NoopTracer.getInstance();
     }
 
     @Override

core/src/main/java/tech/ydb/core/grpc/GrpcTransportBuilder.java

@@ -30,7 +30,6 @@
 import tech.ydb.core.tracing.Tracer;
 import tech.ydb.core.utils.Version;
 
-
 /**
  *
  * @author Aleksandr Gorshenin

core/src/main/java/tech/ydb/core/impl/BaseGrpcTransport.java

@@ -132,7 +132,8 @@ public <ReqT, RespT> CompletableFuture<Result<RespT>> unaryCall(
                         traceId, method.getFullMethodName(), endpoint.getHostAndPort());
             }
             Metadata metadata = makeMetadataFromSettings(settings, endpoint);
-            return new UnaryCall<>(traceId, endpoint.getHostAndPort(), call, handler).startCall(request, metadata);
+            return new UnaryCall<>(traceId, endpoint.getHostAndPort(), call, handler, settings.getSpan())
+                    .startCall(request, metadata);
         } catch (UnexpectedResultException ex) {
             logger.warn("UnaryCall[{}] got unexpected status {}", traceId, ex.getStatus());
             return CompletableFuture.completedFuture(Result.fail(ex));
@@ -173,7 +174,8 @@ public <ReqT, RespT> GrpcReadStream<RespT> readStreamCall(
 
             Metadata metadata = makeMetadataFromSettings(settings, endpoint);
             GrpcFlowControl flowCtrl = settings.getFlowControl();
-            return new ReadStreamCall<>(traceId, endpoint.getHostAndPort(), call, flowCtrl, request, metadata, handler);
+            return new ReadStreamCall<>(traceId, endpoint.getHostAndPort(), call, flowCtrl, request, metadata, handler,
+                    settings.getSpan());
         } catch (UnexpectedResultException ex) {
             logger.warn("ReadStreamCall[{}] got unexpected status {}", traceId, ex.getStatus());
             return new EmptyStream<>(ex.getStatus());
@@ -254,7 +256,7 @@ private Metadata makeMetadataFromSettings(GrpcRequestSettings settings, Endpoint
         }
 
         Span span = settings.getSpan();
-        if (span != null) {
+        if (span.isValid()) {
             span.setAttribute("db.system.name", "ydb");
             span.setAttribute("db.namespace", getDatabase());
             span.setAttribute("server.address", serverEndpoint.getHost());

core/src/main/java/tech/ydb/core/impl/call/ReadStreamCall.java

@@ -19,6 +19,8 @@
 import tech.ydb.core.grpc.GrpcReadStream;
 import tech.ydb.core.grpc.GrpcStatuses;
 import tech.ydb.core.grpc.GrpcTransport;
+import tech.ydb.core.tracing.Scope;
+import tech.ydb.core.tracing.Span;
 
 /**
  *
@@ -34,6 +36,7 @@ public class ReadStreamCall<ReqT, RespT> extends ClientCall.Listener<RespT> impl
     private final ClientCall<ReqT, RespT> call;
     private final Lock callLock = new ReentrantLock();
     private final GrpcStatusHandler statusConsumer;
+    private final Span callSpan;
     private final ReqT request;
     private final Metadata headers;
     private final GrpcFlowControl.Call flow;
@@ -42,14 +45,16 @@ public class ReadStreamCall<ReqT, RespT> extends ClientCall.Listener<RespT> impl
 
     private Observer<RespT> consumer;
 
+    @SuppressWarnings("checkstyle:ParameterNumber")
     public ReadStreamCall(String traceId, String endpoint, ClientCall<ReqT, RespT> call, GrpcFlowControl flowCtrl,
-            ReqT req, Metadata headers, GrpcStatusHandler statusHandler) {
+                          ReqT req, Metadata headers, GrpcStatusHandler statusHandler, Span callSpan) {
         this.traceId = traceId;
         this.endpoint = endpoint;
         this.call = call;
         this.request = req;
         this.headers = headers;
         this.statusConsumer = statusHandler;
+        this.callSpan = callSpan;
         this.flow = flowCtrl.newCall(this::nextRequest);
     }
 
@@ -111,42 +116,45 @@ private void nextRequest(int count) {
 
     @Override
     public void onMessage(RespT message) {
-        try {
-            if (logger.isTraceEnabled()) {
-                logger.trace("ReadStreamCall[{}] <-- {}", traceId, TextFormat.shortDebugString((Message) message));
-            }
-            consumer.onNext(message);
-            flow.onMessageRead();
-        } catch (Exception ex) {
-            statusFuture.completeExceptionally(ex);
-
+        try (Scope ignored = callSpan.makeCurrent()) {
             try {
-                callLock.lock();
+                if (logger.isTraceEnabled()) {
+                    logger.trace("ReadStreamCall[{}] <-- {}", traceId, TextFormat.shortDebugString((Message) message));
+                }
+                consumer.onNext(message);
+                flow.onMessageRead();
+            } catch (Exception ex) {
+                statusFuture.completeExceptionally(ex);
+
                 try {
-                    call.cancel("Canceled by exception from observer", ex);
-                } finally {
-                    callLock.unlock();
+                    callLock.lock();
+                    try {
+                        call.cancel("Canceled by exception from observer", ex);
+                    } finally {
+                        callLock.unlock();
+                    }
+                } catch (Throwable th) {
+                    logger.error("ReadStreamCall[{}] got exception while canceling", traceId, th);
                 }
-            } catch (Throwable th) {
-                logger.error("ReadStreamCall[{}] got exception while canceling", traceId, th);
             }
         }
     }
 
     @Override
     public void onClose(io.grpc.Status status, @Nullable Metadata trailers) {
-        if (logger.isTraceEnabled()) {
-            logger.trace("ReadStreamCall[{}] closed with status {}", traceId, status);
-        }
+        try (Scope ignored = callSpan.restoreContext()) {
+            if (logger.isTraceEnabled()) {
+                logger.trace("ReadStreamCall[{}] closed with status {}", traceId, status);
+            }
 
-        statusConsumer.accept(status, trailers);
+            statusConsumer.accept(status, trailers);
 
-        if (status.isOk()) {
-            statusFuture.complete(Status.SUCCESS);
-        } else {
-            statusFuture.complete(GrpcStatuses.toStatus(status, endpoint));
+            if (status.isOk()) {
+                statusFuture.complete(Status.SUCCESS);
+            } else {
+                statusFuture.complete(GrpcStatuses.toStatus(status, endpoint));
+            }
+            statusConsumer.postComplete();
         }
-
-        statusConsumer.postComplete();
     }
 }

core/src/main/java/tech/ydb/core/impl/call/UnaryCall.java

@@ -18,6 +18,8 @@
 import tech.ydb.core.StatusCode;
 import tech.ydb.core.grpc.GrpcStatuses;
 import tech.ydb.core.grpc.GrpcTransport;
+import tech.ydb.core.tracing.Scope;
+import tech.ydb.core.tracing.Span;
 import tech.ydb.proto.auth.YdbAuth;
 
 /**
@@ -39,15 +41,18 @@ public class UnaryCall<ReqT, RespT> extends ClientCall.Listener<RespT> {
     private final String endpoint;
     private final ClientCall<ReqT, RespT> call;
     private final GrpcStatusHandler statusConsumer;
+    private final Span callSpan;
 
     private final CompletableFuture<Result<RespT>> future = new CompletableFuture<>();
     private final AtomicReference<RespT> value = new AtomicReference<>();
 
-    public UnaryCall(String traceId, String endpoint, ClientCall<ReqT, RespT> call, GrpcStatusHandler statusConsumer) {
+    public UnaryCall(String traceId, String endpoint, ClientCall<ReqT, RespT> call,
+                     GrpcStatusHandler statusConsumer, Span callSpan) {
         this.traceId = traceId;
         this.endpoint = endpoint;
         this.call = call;
         this.statusConsumer = statusConsumer;
+        this.callSpan = callSpan;
     }
 
     public CompletableFuture<Result<RespT>> startCall(ReqT request, Metadata headers) {
@@ -91,23 +96,25 @@ public void onMessage(RespT value) {
 
     @Override
     public void onClose(io.grpc.Status status, @Nullable Metadata trailers) {
-        statusConsumer.accept(status, trailers);
-        if (logger.isTraceEnabled()) {
-            logger.trace("UnaryCall[{}] closed with status {}", traceId, status);
-        }
+        try (Scope ignored = callSpan.restoreContext()) {
+            statusConsumer.accept(status, trailers);
+            if (logger.isTraceEnabled()) {
+                logger.trace("UnaryCall[{}] closed with status {}", traceId, status);
+            }
 
-        if (status.isOk()) {
-            RespT snapshotValue = value.get();
+            if (status.isOk()) {
+                RespT snapshotValue = value.get();
 
-            if (snapshotValue == null) {
-                future.complete(Result.fail(NO_VALUE));
+                if (snapshotValue == null) {
+                    future.complete(Result.fail(NO_VALUE));
+                } else {
+                    future.complete(Result.success(snapshotValue));
+                }
             } else {
-                future.complete(Result.success(snapshotValue));
+                future.complete(GrpcStatuses.toResult(status, endpoint));
             }
-        } else {
-            future.complete(GrpcStatuses.toResult(status, endpoint));
-        }
 
-        statusConsumer.postComplete();
+            statusConsumer.postComplete();
+        }
     }
 }

core/src/main/java/tech/ydb/core/tracing/NoopTracer.java

@@ -7,7 +7,7 @@
  * singletons, no allocations per call (except caller code).
  */
 public final class NoopTracer implements Tracer {
-    public static final NoopTracer INSTANCE = new NoopTracer();
+    private static final NoopTracer INSTANCE = new NoopTracer();
 
     private NoopTracer() {
     }
@@ -18,6 +18,6 @@ public static NoopTracer getInstance() {
 
     @Override
     public Span startSpan(String spanName, SpanKind spanKind) {
-        return null;
+        return Span.NOOP;
     }
 }

core/src/main/java/tech/ydb/core/tracing/Scope.java

@@ -0,0 +1,12 @@
+package tech.ydb.core.tracing;
+
+/**
+ * Closeable scope returned by context-switching span operations.
+ *
+ * <p>Use with try-with-resources to avoid leaking context across async callbacks.
+ */
+public interface Scope extends AutoCloseable {
+
+    @Override
+    void close();
+}