Correctly fork orbits with multiple active shuttles in LWTrace lib

Original fix merged to the nbs repo a while ago: https://github.com/ydb-platform/nbs/pull/4479

Can you bring to the upstream?

---

Pull Request resolved: https://github.com/ytsaurus/ytsaurus/pull/1771

Co-authored-by: Alexander Smirnov <[email protected]>
commit_hash:5bf628be8a9c7b591f290ed05754e206b23efc34

3 files changed, 8 insertions, 2 deletions

diff --git a/library/cpp/lwtrace/log_shuttle.cpp b/library/cpp/lwtrace/log_shuttle.cpp
index 695aa90b316..a6563efb5b9 100644
--- a/library/cpp/lwtrace/log_shuttle.cpp
+++ b/library/cpp/lwtrace/log_shuttle.cpp
@@ -22,7 +22,9 @@ namespace NLWTrace {
 
     template <class TDepot>
     bool TLogShuttle<TDepot>::DoFork(TShuttlePtr& child) {
-        if (child = Executor->RentShuttle()) {
+        if (auto shuttle = Executor->RentShuttle()) {
+            shuttle->SetNext(child);
+            child = shuttle;
             child->SetParentSpanId(GetSpanId());
             Executor->Cast(child)->SetIgnore(true);
             TParams params;
diff --git a/library/cpp/lwtrace/shuttle.h b/library/cpp/lwtrace/shuttle.h
index 6950b3877a4..c4519cbb876 100644
--- a/library/cpp/lwtrace/shuttle.h
+++ b/library/cpp/lwtrace/shuttle.h
@@ -309,7 +309,9 @@ namespace NLWTrace {
                         if (shuttle->IsDead()) {
                             *ref = shuttle->Drop();
                         } else {
-                            result = result && shuttle->Fork(cHead);
+                            if (Y_UNLIKELY(!shuttle->Fork(cHead))) {
+                                result = false;
+                            }
                             ref = &shuttle->GetNext();
                         }
                     }
diff --git a/library/cpp/lwtrace/trace_ut.cpp b/library/cpp/lwtrace/trace_ut.cpp
index 9a29923e28c..5c81a6419a7 100644
--- a/library/cpp/lwtrace/trace_ut.cpp
+++ b/library/cpp/lwtrace/trace_ut.cpp
@@ -789,6 +789,7 @@ Y_UNIT_TEST_SUITE(LWTraceTrace) {
 
         UNIT_ASSERT(parsed);
         mngr.New("Query1", q);
+        mngr.New("Query2", q);
 
         {
             TOrbit a, b, c, d;
@@ -849,6 +850,7 @@ Y_UNIT_TEST_SUITE(LWTraceTrace) {
             }
         } reader;
         mngr.ReadDepot("Query1", reader);
+        mngr.ReadDepot("Query2", reader);
     }
 
     Y_UNIT_TEST(TrackForkError) {