From 6404599b1564f67ee3a7fe2eb7a34e9012cfae55 Mon Sep 17 00:00:00 2001
From: cherepashka <cherepashka@yandex-team.com>
Date: Thu, 18 Jul 2024 11:35:50 +0300
Subject: YT-22281: Fix trace context in RPC
 a85568257402bf5d8f6e439d71b99321654e9289

---
 yt/yt/core/actions/bind-inl.h     | 2 ++
 yt/yt/core/rpc/service_detail.cpp | 6 ++++--
 yt/yt/core/rpc/service_detail.h   | 2 +-
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/yt/yt/core/actions/bind-inl.h b/yt/yt/core/actions/bind-inl.h
index f8ff2d016d5..5af380fadc0 100644
--- a/yt/yt/core/actions/bind-inl.h
+++ b/yt/yt/core/actions/bind-inl.h
@@ -684,6 +684,8 @@ auto Bind(
         THelper::template GetInvokeFunction<TState>()};
 }
 
+// NB: This specialization doesn't act proper to `Propagate` flag,
+// it just copies propagating policy from given callback.
 template <
     bool Propagate,
 #ifdef YT_ENABLE_BIND_LOCATION_TRACKING
diff --git a/yt/yt/core/rpc/service_detail.cpp b/yt/yt/core/rpc/service_detail.cpp
index b1db24e7186..46bdf7ad27c 100644
--- a/yt/yt/core/rpc/service_detail.cpp
+++ b/yt/yt/core/rpc/service_detail.cpp
@@ -1515,9 +1515,11 @@ void TRequestQueue::RunRequest(TServiceBase::TServiceContextPtr context)
     options.SetHeavy(RuntimeInfo_->Heavy.load(std::memory_order::relaxed));
 
     if (options.Heavy) {
-        BIND(RuntimeInfo_->Descriptor.HeavyHandler)
+        BIND([this, this_ = MakeStrong(this), context, options] {
+            return RuntimeInfo_->Descriptor.HeavyHandler.Run(context, options);
+        })
             .AsyncVia(TDispatcher::Get()->GetHeavyInvoker())
-            .Run(context, options)
+            .Run()
             .Subscribe(BIND(&TServiceBase::TServiceContext::CheckAndRun, std::move(context)));
     } else {
         context->Run(RuntimeInfo_->Descriptor.LiteHandler);
diff --git a/yt/yt/core/rpc/service_detail.h b/yt/yt/core/rpc/service_detail.h
index de2f40c81f6..f731ed84545 100644
--- a/yt/yt/core/rpc/service_detail.h
+++ b/yt/yt/core/rpc/service_detail.h
@@ -426,7 +426,7 @@ protected:
             return ::NYT::NRpc::TServiceBase::TLiteHandler(); \
         } \
         return \
-            BIND_NO_PROPAGATE([=, this] ( \
+            BIND([=, this] ( \
                 const ::NYT::NRpc::IServiceContextPtr&, \
                 const ::NYT::NRpc::THandlerInvocationOptions&) \
             { \
-- 
cgit v1.3