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

@@ -126,13 +126,14 @@ public <ReqT, RespT> CompletableFuture<Result<RespT>> unaryCall(
 
             ClientCall<ReqT, RespT> call = channel.getReadyChannel().newCall(method, options);
             ChannelStatusHandler handler = new ChannelStatusHandler(channel, settings);
-
+            Span contextSpan = settings.getSpan();
             if (logger.isTraceEnabled()) {
                 logger.trace("UnaryCall[{}] with method {} and endpoint {} created",
                         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, contextSpan)
+                    .startCall(request, metadata);
         } catch (UnexpectedResultException ex) {
             logger.warn("UnaryCall[{}] got unexpected status {}", traceId, ex.getStatus());
             return CompletableFuture.completedFuture(Result.fail(ex));
@@ -164,6 +165,7 @@ public <ReqT, RespT> GrpcReadStream<RespT> readStreamCall(
 
             ClientCall<ReqT, RespT> call = channel.getReadyChannel().newCall(method, options);
             ChannelStatusHandler handler = new ChannelStatusHandler(channel, settings);
+            Span contextSpan = settings.getSpan();
 
             if (logger.isTraceEnabled()) {
                 logger.trace("ReadStreamCall[{}] with method {} and endpoint {} created",
@@ -173,7 +175,9 @@ 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<>(
+                    endpoint.getHostAndPort(), call, flowCtrl, request, metadata, handler, contextSpan
+            );
         } catch (UnexpectedResultException ex) {
             logger.warn("ReadStreamCall[{}] got unexpected status {}", traceId, ex.getStatus());
             return new EmptyStream<>(ex.getStatus());
@@ -254,7 +258,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.Span;
+import tech.ydb.core.tracing.SpanScope;
 
 /**
  *
@@ -29,11 +31,12 @@
 public class ReadStreamCall<ReqT, RespT> extends ClientCall.Listener<RespT> implements GrpcReadStream<RespT> {
     private static final Logger logger = LoggerFactory.getLogger(GrpcTransport.class);
 
-    private final String traceId;
     private final String endpoint;
     private final ClientCall<ReqT, RespT> call;
     private final Lock callLock = new ReentrantLock();
     private final GrpcStatusHandler statusConsumer;
+    private final Span parentSpan;
+    private final Span span;
     private final ReqT request;
     private final Metadata headers;
     private final GrpcFlowControl.Call flow;
@@ -42,14 +45,15 @@ public class ReadStreamCall<ReqT, RespT> extends ClientCall.Listener<RespT> impl
 
     private Observer<RespT> consumer;
 
-    public ReadStreamCall(String traceId, String endpoint, ClientCall<ReqT, RespT> call, GrpcFlowControl flowCtrl,
-            ReqT req, Metadata headers, GrpcStatusHandler statusHandler) {
-        this.traceId = traceId;
+    public ReadStreamCall(String endpoint, ClientCall<ReqT, RespT> call, GrpcFlowControl flowCtrl, ReqT req,
+            Metadata headers, GrpcStatusHandler statusHandler, Span span) {
         this.endpoint = endpoint;
         this.call = call;
         this.request = req;
         this.headers = headers;
         this.statusConsumer = statusHandler;
+        this.parentSpan = span.getParentSpan();
+        this.span = span;
         this.flow = flowCtrl.newCall(this::nextRequest);
     }
 
@@ -67,7 +71,7 @@ public CompletableFuture<Status> start(Observer<RespT> observer) {
             consumer = observer;
             call.start(this, headers);
             if (logger.isTraceEnabled()) {
-                logger.trace("ReadStreamCall[{}] --> {}", traceId, TextFormat.shortDebugString((Message) request));
+                logger.trace("ReadStreamCall[{}] --> {}", endpoint, TextFormat.shortDebugString((Message) request));
             }
             call.sendMessage(request);
             // close stream by client side
@@ -80,7 +84,7 @@ public CompletableFuture<Status> start(Observer<RespT> observer) {
             try {
                 call.cancel(null, th);
             } catch (Throwable ex) {
-                logger.error("ReadStreamCall[{}] got exception while canceling", traceId, ex);
+                logger.error("ReadStreamCall[{}] got exception while canceling", endpoint, ex);
             }
         } finally {
             callLock.unlock();
@@ -111,42 +115,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 (SpanScope ignored = span.makeCurrent()) {
             try {
-                callLock.lock();
+                if (logger.isTraceEnabled()) {
+                    logger.trace("ReadStreamCall[{}] <-- {}", endpoint, 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", endpoint, 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 (SpanScope ignored = parentSpan.makeCurrent()) {
+            if (logger.isTraceEnabled()) {
+                logger.trace("ReadStreamCall[{}] closed with status {}", endpoint, 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.Span;
+import tech.ydb.core.tracing.SpanScope;
 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 parentSpan;
 
     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 contextSpan) {
         this.traceId = traceId;
         this.endpoint = endpoint;
         this.call = call;
         this.statusConsumer = statusConsumer;
+        this.parentSpan = contextSpan.getParentSpan();
     }
 
     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 (SpanScope ignored = parentSpan.makeCurrent()) {
+            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;
     }
 }