Tracing in topic write session (#12649)

Author: UgnineSirdis

ydb/core/client/server/msgbus_server_pq_metacache.cpp

@@ -5,6 +5,7 @@
 #include <ydb/public/lib/deprecated/kicli/kicli.h>
 
 #include <ydb/library/persqueue/topic_parser/topic_parser.h>
+#include <ydb/library/wilson_ids/wilson.h>
 #include <ydb/core/tx/tx_proxy/proxy.h>
 #include <ydb/core/tx/schemeshard/schemeshard.h>
 #include <ydb/core/tx/scheme_board/cache.h>
@@ -357,16 +358,18 @@ class TPersQueueMetaCacheActor : public TActorBootstrapped<TPersQueueMetaCacheAc
         bool FirstRequestDone = false;
         bool SyncVersion;
         bool ShowPrivate;
+        NWilson::TSpan Span;
 
         TWaiter(const TActorId& waiterId, const TString& dbRoot, bool syncVersion, bool showPrivate,
-                const TVector<NPersQueue::TDiscoveryConverterPtr>& topics, EWaiterType type)
+                const TVector<NPersQueue::TDiscoveryConverterPtr>& topics, EWaiterType type, NWilson::TSpan span = {})
 
             : WaiterId(waiterId)
             , DbRoot(dbRoot)
             , Topics(topics)
             , Type(type)
             , SyncVersion(syncVersion)
             , ShowPrivate(showPrivate)
+            , Span(std::move(span))
         {}
 
         bool ApplyResult(std::shared_ptr<TSchemeCacheNavigate>& result) {
@@ -457,7 +460,7 @@ class TPersQueueMetaCacheActor : public TActorBootstrapped<TPersQueueMetaCacheAc
         }
         SendSchemeCacheRequest(
                 std::make_shared<TWaiter>(ev->Sender, DbRoot, msg.SyncVersion, msg.ShowPrivate, ev->Get()->Topics,
-                                          EWaiterType::DescribeCustomTopics),
+                                          EWaiterType::DescribeCustomTopics, NWilson::TSpan(TWilsonTopic::TopicDetailed, NWilson::TTraceId(ev->TraceId), "Topic.SchemeCacheRequest", NWilson::EFlags::AUTO_END)),
                 ctx
         );
     }
@@ -532,7 +535,7 @@ class TPersQueueMetaCacheActor : public TActorBootstrapped<TPersQueueMetaCacheAc
 
         LOG_DEBUG_S(ctx, NKikimrServices::PQ_METACACHE, "send request for " << (waiter->Type == EWaiterType::DescribeAllTopics ? " all " : "") << waiter->GetTopics().size() << " topics, got " << DescribeTopicsWaiters.size() << " requests infly");
 
-        ctx.Send(SchemeCacheId, new TEvTxProxySchemeCache::TEvNavigateKeySet(schemeCacheRequest.release()));
+        ctx.Send(SchemeCacheId, new TEvTxProxySchemeCache::TEvNavigateKeySet(schemeCacheRequest.release()), 0, 0, waiter->Span.GetTraceId());
     }
 
     void HandleSchemeCacheResponse(TEvTxProxySchemeCache::TEvNavigateKeySetResult::TPtr& ev, const TActorContext& ctx) {

ydb/core/grpc_services/base/base.h

@@ -906,7 +906,7 @@ class TGRpcRequestBiStreamWrapper
     void SetDiskQuotaExceeded(bool) override {
     }
 
-    void RefreshToken(const TString& token, const TActorContext& ctx, TActorId id);
+    void RefreshToken(const TString& token, const TActorContext& ctx, TActorId id, NWilson::TTraceId traceId = {});
 
     void SetRespHook(TRespHook&&) override {
         /* cannot add hook to bidirect streaming */

ydb/core/grpc_services/grpc_request_proxy.cpp

@@ -77,7 +77,8 @@ class TGRpcRequestProxyImpl
     void HandleSchemeBoard(TSchemeBoardEvents::TEvNotifyDelete::TPtr& ev);
     void ReplayEvents(const TString& databaseName, const TActorContext& ctx);
 
-    void MaybeStartTracing(IRequestProxyCtx& ctx);
+    template<class TEvent>
+    void MaybeStartTracing(TAutoPtr<TEventHandle<TEvent>>& event);
 
     static bool IsAuthStateOK(const IRequestProxyCtx& ctx);
 
@@ -153,7 +154,7 @@ class TGRpcRequestProxyImpl
             return;
         }
 
-        MaybeStartTracing(*requestBaseCtx);
+        MaybeStartTracing(event);
 
         if (IsAuthStateOK(*requestBaseCtx)) {
             Handle(event, ctx);
@@ -436,7 +437,9 @@ bool TGRpcRequestProxyImpl::IsAuthStateOK(const IRequestProxyCtx& ctx) {
     return false;
 }
 
-void TGRpcRequestProxyImpl::MaybeStartTracing(IRequestProxyCtx& ctx) {
+template<class TEvent>
+void TGRpcRequestProxyImpl::MaybeStartTracing(TAutoPtr<TEventHandle<TEvent>>& event) {
+    IRequestProxyCtx& ctx = *event->Get();
     auto isTracingDecided = ctx.IsTracingDecided();
     if (!isTracingDecided) {
         return;
@@ -445,8 +448,12 @@ void TGRpcRequestProxyImpl::MaybeStartTracing(IRequestProxyCtx& ctx) {
         return;
     }
 
-    TMaybe<TString> traceparentHeader = ctx.GetPeerMetaValues(NYdb::OTEL_TRACE_HEADER);
-    NWilson::TTraceId traceId = NJaegerTracing::HandleTracing(ctx.GetRequestDiscriminator(), traceparentHeader);
+    NWilson::TTraceId traceId = NWilson::TTraceId(event->TraceId); // Can be not empty in case of internal subrequests // In this case it is part of the big request
+    if (!traceId) {
+        TMaybe<TString> traceparentHeader = ctx.GetPeerMetaValues(NYdb::OTEL_TRACE_HEADER);
+        traceId = NJaegerTracing::HandleTracing(ctx.GetRequestDiscriminator(), traceparentHeader);
+    }
+
     if (traceId) {
         NWilson::TSpan grpcRequestProxySpan(TWilsonGrpc::RequestProxy, std::move(traceId), "GrpcRequestProxy");
         if (auto database = ctx.GetDatabaseName()) {

ydb/core/grpc_services/local_rate_limiter.cpp

@@ -13,7 +13,8 @@ TActorId RateLimiterAcquireUseSameMailbox(
     const TDuration& duration,
     std::function<void()>&& onSuccess,
     std::function<void()>&& onTimeout,
-    const TActorContext& ctx)
+    const TActorContext& ctx,
+    NWilson::TTraceId traceId)
 {
     auto cb = [onSuccess{std::move(onSuccess)}, onTimeout{std::move(onTimeout)}]
         (const Ydb::RateLimiter::AcquireResourceResponse& resp)
@@ -43,18 +44,20 @@ TActorId RateLimiterAcquireUseSameMailbox(
         fullPath.DatabaseName,
         fullPath.Token,
         std::move(cb),
-        ctx);
+        ctx,
+        std::move(traceId));
 }
 
 TActorId RateLimiterAcquireUseSameMailbox(
     Ydb::RateLimiter::AcquireResourceRequest&& request,
     const TString& database,
     const TString& token,
     std::function<void(Ydb::RateLimiter::AcquireResourceResponse resp)>&& cb,
-    const TActorContext& ctx)
+    const TActorContext& ctx,
+    NWilson::TTraceId traceId)
 {
     return DoLocalRpcSameMailbox<NKikimr::NGRpcService::TEvAcquireRateLimiterResource>(
-        std::move(request), std::move(cb), database, token, ctx);
+        std::move(request), std::move(cb), database, token, ctx, false, std::move(traceId));
 }
 
 TActorId RateLimiterAcquireUseSameMailbox(
@@ -63,7 +66,8 @@ TActorId RateLimiterAcquireUseSameMailbox(
     const TDuration& duration,
     std::function<void()>&& onSuccess,
     std::function<void()>&& onTimeout,
-    const TActorContext& ctx)
+    const TActorContext& ctx,
+    NWilson::TTraceId traceId)
 {
     if (const auto maybeRlPath = reqCtx.GetRlPath()) {
         auto cb = [onSuccess{std::move(onSuccess)}, onTimeout{std::move(onTimeout)}]
@@ -95,7 +99,8 @@ TActorId RateLimiterAcquireUseSameMailbox(
             reqCtx.GetDatabaseName().GetOrElse(""),
             reqCtx.GetSerializedToken(),
             std::move(cb),
-            ctx);
+            ctx,
+            std::move(traceId));
     }
     return {};
 }

ydb/core/grpc_services/local_rate_limiter.h

@@ -41,20 +41,23 @@ TActorId RateLimiterAcquireUseSameMailbox(
     const TDuration& duration,
     std::function<void()>&& onSuccess,
     std::function<void()>&& onTimeout,
-    const TActorContext& ctx);
+    const TActorContext& ctx,
+    NWilson::TTraceId traceId = {});
 
 TActorId RateLimiterAcquireUseSameMailbox(
     Ydb::RateLimiter::AcquireResourceRequest&& request,
     const TString& database,
     const TString& token,
-    std::function<void(Ydb::RateLimiter::AcquireResourceResponse resp)>&& cb, const TActorContext &ctx);
+    std::function<void(Ydb::RateLimiter::AcquireResourceResponse resp)>&& cb, const TActorContext& ctx,
+    NWilson::TTraceId traceId = {});
 
 TActorId RateLimiterAcquireUseSameMailbox(const NGRpcService::IRequestCtxBase& reqCtx,
     ui64 required,
     const TDuration& duration,
     std::function<void()>&& onSuccess,
     std::function<void()>&& onFail,
-    const NActors::TActorContext &ctx);
+    const NActors::TActorContext& ctx,
+    NWilson::TTraceId traceId = {});
 
 struct TRlConfig {
     struct TOnReqAction {

ydb/core/grpc_services/local_rpc/local_rpc.h

@@ -5,6 +5,9 @@
 
 #include <ydb/core/base/appdata.h>
 
+#include <ydb/library/actors/wilson/wilson_span.h>
+#include <ydb/library/wilson_ids/wilson.h>
+
 #include <library/cpp/threading/future/future.h>
 
 namespace NKikimr {
@@ -48,9 +51,12 @@ class TLocalRpcCtxImpl<TRpc, TCbWrapper, false> : public NGRpcService::IRequestN
     using TBase = TLocalRpcCtxImplData<TRpc, TCbWrapper>;
 
     template<typename TCb>
-    TLocalRpcCtxImpl(TCb&& cb)
+    TLocalRpcCtxImpl(TCb&& cb, NWilson::TTraceId = {})
         : TBase(std::forward<TCb>(cb))
     {}
+
+protected:
+    NWilson::TSpan Span;
 };
 
 template<typename TRpc, typename TCbWrapper>
@@ -59,8 +65,9 @@ class TLocalRpcCtxImpl<TRpc, TCbWrapper, true> : public NGRpcService::IRequestOp
     using TBase = TLocalRpcCtxImplData<TRpc, TCbWrapper>;
 
     template<typename TCb>
-    TLocalRpcCtxImpl(TCb&& cb)
+    TLocalRpcCtxImpl(TCb&& cb, NWilson::TTraceId traceId = {})
         : TBase(std::forward<TCb>(cb))
+        , Span(TWilsonGrpc::RequestProxy, std::move(traceId), "LocalRpc")
     {}
 
 public:
@@ -77,6 +84,7 @@ class TLocalRpcCtxImpl<TRpc, TCbWrapper, true> : public NGRpcService::IRequestOp
         NYql::IssuesToMessage(TBase::IssueManager.GetIssues(), deferred->mutable_issues());
         auto data = deferred->mutable_result();
         data->PackFrom(result);
+        EndSpan(status);
         TBase::CbWrapper(resp);
     }
 
@@ -93,14 +101,25 @@ class TLocalRpcCtxImpl<TRpc, TCbWrapper, true> : public NGRpcService::IRequestOp
         }
         auto data = deferred->mutable_result();
         data->PackFrom(result);
+        EndSpan(status);
         TBase::CbWrapper(resp);
     }
 
     void SendOperation(const Ydb::Operations::Operation& operation) override {
         TResp resp;
         resp.mutable_operation()->CopyFrom(operation);
+        EndSpan(operation.status());
         TBase::CbWrapper(resp);
     }
+
+protected:
+    void EndSpan(Ydb::StatusIds::StatusCode status) {
+        Span.Attribute("status", static_cast<int>(status));
+        Span.End();
+    }
+
+protected:
+    NWilson::TSpan Span;
 };
 
 template<typename TRpc, typename TCbWrapper, bool IsOperation = TRpc::IsOp>
@@ -116,8 +135,9 @@ class TLocalRpcCtx : public TLocalRpcCtxImpl<TRpc, TCbWrapper, IsOperation> {
             const TString& databaseName,
             const TMaybe<TString>& token,
             const TMaybe<TString>& requestType,
-            bool internalCall)
-        : TBase(std::forward<TCb>(cb))
+            bool internalCall,
+            NWilson::TTraceId traceId = {})
+        : TBase(std::forward<TCb>(cb), std::move(traceId))
         , Request(std::forward<TProto>(req))
         , DatabaseName(databaseName)
         , RequestType(requestType)
@@ -126,6 +146,11 @@ class TLocalRpcCtx : public TLocalRpcCtxImpl<TRpc, TCbWrapper, IsOperation> {
         if (token && !token->empty()) {
             InternalToken = new NACLib::TUserToken(*token);
         }
+
+        if (DatabaseName) {
+            this->Span.Attribute("database", DatabaseName);
+        }
+        this->Span.Attribute("request_type", GetRequestName());
     }
 
     bool HasClientCapability(const TString&) const override {
@@ -226,7 +251,7 @@ class TLocalRpcCtx : public TLocalRpcCtxImpl<TRpc, TCbWrapper, IsOperation> {
     }
 
     NWilson::TTraceId GetWilsonTraceId() const override {
-        return {};
+        return this->Span.GetTraceId();
     }
 
     TInstant GetDeadline() const override {
@@ -301,14 +326,14 @@ void SetRequestSyncOperationMode(TRequest&) {
 template<typename TRpc>
 NThreading::TFuture<typename TRpc::TResponse> DoLocalRpc(typename TRpc::TRequest&& proto, const TString& database,
         const TMaybe<TString>& token, const TMaybe<TString>& requestType,
-        TActorSystem* actorSystem, bool internalCall = false)
+        TActorSystem* actorSystem, bool internalCall = false, NWilson::TTraceId traceId = {})
 {
     auto promise = NThreading::NewPromise<typename TRpc::TResponse>();
 
     SetRequestSyncOperationMode(proto);
 
     using TCbWrapper = TPromiseWrapper<typename TRpc::TResponse>;
-    auto req = new TLocalRpcCtx<TRpc, TCbWrapper>(std::move(proto), TCbWrapper(promise), database, token, requestType, internalCall);
+    auto req = new TLocalRpcCtx<TRpc, TCbWrapper>(std::move(proto), TCbWrapper(promise), database, token, requestType, internalCall, std::move(traceId));
     auto actor = TRpc::CreateRpcActor(req);
     actorSystem->Register(actor, TMailboxType::HTSwap, actorSystem->AppData<TAppData>()->UserPoolId);
 
@@ -328,7 +353,8 @@ NThreading::TFuture<typename TRpc::TResponse> DoLocalRpc(
         const TMaybe<TString>& requestType,
         TActorSystem* actorSystem,
         const TMap<TString, TString>& peerMeta,
-        bool internalCall = false
+        bool internalCall = false,
+        NWilson::TTraceId traceId = {}
 )
 {
     auto promise = NThreading::NewPromise<typename TRpc::TResponse>();
@@ -342,7 +368,8 @@ NThreading::TFuture<typename TRpc::TResponse> DoLocalRpc(
         database,
         token,
         requestType,
-        internalCall
+        internalCall,
+        std::move(traceId)
     );
 
     for (const auto& [key, value] : peerMeta) {
@@ -358,18 +385,18 @@ NThreading::TFuture<typename TRpc::TResponse> DoLocalRpc(
 template<typename TRpc>
 TActorId DoLocalRpcSameMailbox(typename TRpc::TRequest&& proto, std::function<void(typename TRpc::TResponse)>&& cb,
         const TString& database, const TMaybe<TString>& token, const TMaybe<TString>& requestType,
-        const TActorContext& ctx, bool internalCall = false)
+        const TActorContext& ctx, bool internalCall = false, NWilson::TTraceId traceId = {})
 {
     SetRequestSyncOperationMode(proto);
 
-    auto req = new TLocalRpcCtx<TRpc, std::function<void(typename TRpc::TResponse)>>(std::move(proto), std::move(cb), database, token, requestType, internalCall);
+    auto req = new TLocalRpcCtx<TRpc, std::function<void(typename TRpc::TResponse)>>(std::move(proto), std::move(cb), database, token, requestType, internalCall, std::move(traceId));
     auto actor = TRpc::CreateRpcActor(req);
     return ctx.RegisterWithSameMailbox(actor);
 }
 
 template<typename TRpc>
-TActorId DoLocalRpcSameMailbox(typename TRpc::TRequest&& proto, std::function<void(typename TRpc::TResponse)>&& cb, const TString& database, const TMaybe<TString>& token, const TActorContext& ctx, bool internalCall = false) {
-    return DoLocalRpcSameMailbox<TRpc>(std::move(proto), std::move(cb), database, token, Nothing(), ctx, internalCall);
+TActorId DoLocalRpcSameMailbox(typename TRpc::TRequest&& proto, std::function<void(typename TRpc::TResponse)>&& cb, const TString& database, const TMaybe<TString>& token, const TActorContext& ctx, bool internalCall = false, NWilson::TTraceId traceId = {}) {
+    return DoLocalRpcSameMailbox<TRpc>(std::move(proto), std::move(cb), database, token, Nothing(), ctx, internalCall, std::move(traceId));
 }
 
 //// Streaming part

ydb/core/grpc_services/rpc_calls.cpp

@@ -75,8 +75,8 @@ std::pair<TString, TString> SplitPath(const TString& path) {
     return {path.substr(0, splitPos), path.substr(splitPos + 1)};
 }
 
-void RefreshTokenSendRequest(const TActorContext& ctx, IEventBase* refreshTokenRequest) {
-    ctx.Send(CreateGRpcRequestProxyId(), refreshTokenRequest);
+void RefreshTokenSendRequest(const TActorContext& ctx, IEventBase* refreshTokenRequest, NWilson::TTraceId traceId) {
+    ctx.Send(CreateGRpcRequestProxyId(), refreshTokenRequest, 0, 0, std::move(traceId));
 }
 
 void RefreshTokenReplyUnauthenticated(TActorId recipient, TActorId sender, NYql::TIssues&& issues) {

ydb/core/grpc_services/rpc_calls.h

@@ -117,15 +117,15 @@ struct TRefreshTokenTypeForRequest<TEvStreamTopicWriteRequest> {
 // RefreshToken Send/Reply interface hides lowlevel details.
 // Used to avoid unwanted compile time dependencies.
 //
-void RefreshTokenSendRequest(const TActorContext& ctx, IEventBase* refreshTokenRequest);
+void RefreshTokenSendRequest(const TActorContext& ctx, IEventBase* refreshTokenRequest, NWilson::TTraceId traceId);
 void RefreshTokenReplyUnauthenticated(TActorId recipient, TActorId sender, NYql::TIssues&& issues);
 void RefreshTokenReplyUnavailable(TActorId recipient, NYql::TIssues&& issues);
 
 template <ui32 TRpcId, typename TReq, typename TResp>
-void TGRpcRequestBiStreamWrapper<TRpcId, TReq, TResp>::RefreshToken(const TString& token, const TActorContext& ctx, TActorId id) {
+void TGRpcRequestBiStreamWrapper<TRpcId, TReq, TResp>::RefreshToken(const TString& token, const TActorContext& ctx, TActorId id, NWilson::TTraceId traceId) {
     using TSelf = typename std::remove_pointer<decltype(this)>::type;
     using TRefreshToken = typename TRefreshTokenTypeForRequest<TSelf>::type;
-    RefreshTokenSendRequest(ctx, new TRefreshToken(token, GetDatabaseName().GetOrElse(""), id));
+    RefreshTokenSendRequest(ctx, new TRefreshToken(token, GetDatabaseName().GetOrElse(""), id), std::move(traceId));
 }
 
 template <ui32 TRpcId>