Refactor IDirectSession to support it in v1 (#46292)

7 files changed, 215 insertions, 42 deletions

diff --git a/ydb/library/actors/interconnect/events/events.h b/ydb/library/actors/interconnect/events/events.h
index 285edbe0dd5..a54f3710781 100644
--- a/ydb/library/actors/interconnect/events/events.h
+++ b/ydb/library/actors/interconnect/events/events.h
@@ -65,6 +65,9 @@ namespace NActors {
         EvUringUnregister,
         EvUringRegisterFailed,
 
+        // wake for the direct-session (v1) lock-free registration queue
+        EvProcessDirectSessionQueue,
+
         ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
         // nonlocal messages; their indices must be preserved in order to work properly while doing rolling update
         ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/ydb/library/actors/interconnect/interconnect_direct_session.h b/ydb/library/actors/interconnect/interconnect_direct_session.h
index 95c241904d1..43bf8a1f424 100644
--- a/ydb/library/actors/interconnect/interconnect_direct_session.h
+++ b/ydb/library/actors/interconnect/interconnect_direct_session.h
@@ -1,110 +1,248 @@
 #pragma once
 
 #include <ydb/library/actors/core/actorid.h>
+#include <ydb/library/actors/core/actorsystem.h>
 #include <ydb/library/actors/core/event.h>
+#include <ydb/library/actors/util/funnel_queue.h>
+
+#include "events/events.h"
 
 #include <util/generic/hash.h>
+#include <util/generic/ptr.h>
 #include <util/system/mutex.h>
 
-#include <functional>
+#include <atomic>
 #include <memory>
 
 namespace NActors {
 
-    // Thread-safe handle to an established interconnect session (v2) that allows sending and receiving
+    // Refcounted receiver for incoming events delivered outside the actor system. An instance is
+    // registered for a local ActorId; the session invokes Receive() (on its input thread) for every
+    // incoming event addressed to that ActorId. Held via TIntrusivePtr, so the session co-owns the
+    // receiver for as long as it stays registered.
+    class IReceiveCallback : public TThrRefBase {
+    public:
+        virtual void Receive(TAutoPtr<IEventHandle> ev) = 0;
+    };
+
+    // Thread-safe handle to an established interconnect session that allows sending and receiving
     // events to/from the remote peer without going through the actor system. It is delivered to
-    // subscribers inside TEvInterconnect::TEvNodeConnected as a std::shared_ptr.
+    // subscribers inside TEvInterconnect::TEvNodeConnected as a std::shared_ptr and is universal: the
+    // same interface is exposed by both the classic (v1) and the newer (v2) session implementations,
+    // so consumers do not need to change when upgrading to a newer protocol.
     //
     // Lifecycle / races: both the session actor and every user hold a std::shared_ptr to the same
-    // object. When the connection is lost the session calls Shutdown(), which atomically transitions
-    // the handle into a disconnected state (Send() starts returning false and all registered callbacks
-    // are dropped). All methods are safe to call concurrently from arbitrary threads.
+    // object. When the connection is lost the session calls Shutdown(), which transitions the handle
+    // into a disconnected state (Send() starts returning false and all registered callbacks are
+    // dropped). All methods are safe to call concurrently from arbitrary threads.
     class IDirectSession {
     public:
-        // Invoked (on the session actor thread) for an incoming event addressed to a registered ActorId.
-        using TReceiveCallback = std::function<void(TAutoPtr<IEventHandle>)>;
-
         virtual ~IDirectSession() = default;
 
         // Sends an event to the remote peer. The event's Recipient identifies the destination actor on
-        // the peer node. Returns false if the session is no longer connected. Thread-safe.
-        virtual bool Send(TAutoPtr<IEventHandle> ev) = 0;
+        // the peer node. If replyCallback is set, it is registered for ev->Sender BEFORE the event is
+        // dispatched (equivalent to RegisterReceiveCallback(ev->Sender, replyCallback) followed by
+        // Send(ev), but saving one virtual call and guaranteeing the reply cannot race ahead of the
+        // registration). Returns false if the session is no longer connected. Thread-safe.
+        virtual bool Send(TAutoPtr<IEventHandle> ev, TIntrusivePtr<IReceiveCallback> replyCallback = nullptr) = 0;
 
         // Registers (or replaces) a callback invoked for incoming events addressed to localActorId.
         // Has no effect once the session is disconnected. Thread-safe.
-        virtual void RegisterReceiveCallback(const TActorId& localActorId, TReceiveCallback callback) = 0;
+        virtual void RegisterReceiveCallback(const TActorId& localActorId, TIntrusivePtr<IReceiveCallback> callback) = 0;
 
         // Removes a previously registered callback. Thread-safe.
         virtual void UnregisterReceiveCallback(const TActorId& localActorId) = 0;
     };
 
-    // Concrete implementation of IDirectSession backing a TInterconnectSessionTCPv2 connection.
-    // The data-plane transfer (actual serialization / socket handoff of outgoing events and decoding
-    // of incoming ones) is intentionally left as a stub; the thread-safe registration table and the
-    // connect/disconnect lifecycle are fully implemented.
-    class TDirectSession final : public IDirectSession {
+    // Direct session backing a classic (v1) TInterconnectSessionTCP connection.
+    //
+    // Design: the producer side (Send / Register / Unregister, callable from arbitrary threads) never
+    // takes a lock. Registration changes are pushed to a lock-free MPSC queue; when the queue turns
+    // non-empty the (stable) session actor is woken to drain it. The callback table itself is owned by
+    // the session/input-session mailbox thread -- the session and its TInputSessionTCP share one
+    // mailbox (RegisterWithSameMailbox), so the table is applied and looked up from a single thread
+    // without any synchronization. The queue is drained on the wake event and also right before an
+    // incoming event is dispatched, so a reply can never overtake its registration.
+    //
+    // Outgoing events are routed through the actor system, which delivers them to the peer via the
+    // existing interconnect machinery.
+    class TDirectSessionV1 final : public IDirectSession {
     public:
-        bool Send(TAutoPtr<IEventHandle> ev) override {
-            TGuard<TMutex> guard(Mutex);
-            if (!Connected) {
+        TDirectSessionV1(TActorSystem* actorSystem, const TActorId& sessionId, ui32 peerNodeId)
+            : ActorSystem(actorSystem)
+            , SessionId(sessionId)
+            , PeerNodeId(peerNodeId)
+        {}
+
+        // ---- producer side (arbitrary threads), lock-free ----
+
+        bool Send(TAutoPtr<IEventHandle> ev, TIntrusivePtr<IReceiveCallback> replyCallback = nullptr) override {
+            if (!Connected.load(std::memory_order_acquire)) {
+                return false;
+            }
+            if (replyCallback) {
+                EnqueueCommand(ev->Sender, std::move(replyCallback));
+            }
+            Y_ABORT_UNLESS(ev->GetRecipientRewrite().NodeId() == PeerNodeId);
+            ActorSystem->Send(ev);
+            return true;
+        }
+
+        void RegisterReceiveCallback(const TActorId& localActorId, TIntrusivePtr<IReceiveCallback> callback) override {
+            if (Connected.load(std::memory_order_acquire)) {
+                EnqueueCommand(localActorId, std::move(callback));
+            }
+        }
+
+        void UnregisterReceiveCallback(const TActorId& localActorId) override {
+            if (Connected.load(std::memory_order_acquire)) {
+                EnqueueCommand(localActorId, nullptr);
+            }
+        }
+
+        // ---- consumer side (session/input-session mailbox thread only) ----
+
+        // Applies every queued registration command into the table. Called on the
+        // EvProcessDirectSessionQueue wake and before dispatching incoming events.
+        void ApplyPendingCommands(bool drop = false) {
+            while (!CommandQueue.IsEmpty()) {
+                TCommand& cmd = CommandQueue.Top();
+                if (!drop) {
+                    if (cmd.Callback) {
+                        Callbacks[cmd.ActorId] = std::move(cmd.Callback);
+                    } else {
+                        Callbacks.erase(cmd.ActorId);
+                    }
+                }
+                CommandQueue.Pop();
+            }
+        }
+
+        // Delivers an incoming event to the callback registered for its Recipient, if any. Returns true
+        // if a callback consumed the event (ev is moved out), false otherwise (ev is left untouched and
+        // the caller should deliver the event through the actor system).
+        bool DeliverIncoming(TAutoPtr<IEventHandle>& ev) {
+            ApplyPendingCommands();
+            if (Callbacks.empty()) {
                 return false;
             }
+            // just received -> Recipient and RecipientRewrite are guaranteed to match
+            Y_DEBUG_ABORT_UNLESS(ev->Recipient == ev->GetRecipientRewrite());
+            if (const auto it = Callbacks.find(ev->Recipient); it != Callbacks.end()) {
+                it->second->Receive(std::move(ev));
+                return true;
+            }
+            return false;
+        }
+
+        // Transitions the handle into the disconnected state. Called by the session on teardown, on the
+        // same mailbox thread that owns the table, so it can safely drop the callbacks.
+        void Shutdown() {
+            Connected.store(false, std::memory_order_release);
+            ApplyPendingCommands(true); // drain the queue to release its entries
+            Callbacks.clear();
+        }
+
+    private:
+        struct TCommand {
+            TActorId ActorId;
+            TIntrusivePtr<IReceiveCallback> Callback; // null => unregister
+        };
+
+        void EnqueueCommand(const TActorId& actorId, TIntrusivePtr<IReceiveCallback> callback) {
+            if (CommandQueue.Push(TCommand{actorId, std::move(callback)})) {
+                // queue went empty -> non-empty: wake the stable session actor to drain the queue
+                ActorSystem->Send(new IEventHandle(static_cast<ui32>(ENetwork::EvProcessDirectSessionQueue),
+                    0, SessionId, TActorId(), nullptr, 0));
+            }
+        }
+
+    private:
+        TActorSystem* const ActorSystem;
+        const TActorId SessionId;
+        const ui32 PeerNodeId;
+        std::atomic<bool> Connected = true;
+        TFunnelQueue<TCommand> CommandQueue; // MPSC: many producers, single consumer (mailbox thread)
+        // touched only on the session/input-session mailbox thread
+        THashMap<TActorId, TIntrusivePtr<IReceiveCallback>, TActorId::THash> Callbacks;
+    };
+
+    // Direct session backing a newer (v2) TInterconnectSessionTCPv2 connection. The outgoing data plane
+    // (serialization / socket handoff) is intentionally left as a stub for now, so this implementation
+    // keeps a simple mutex-guarded registration table. An atomic counter lets DeliverIncoming skip the
+    // mutex entirely when no receivers are registered.
+    class TDirectSessionV2 final : public IDirectSession {
+    public:
+        bool Send(TAutoPtr<IEventHandle> ev, TIntrusivePtr<IReceiveCallback> replyCallback = nullptr) override {
+            {
+                TGuard<TMutex> guard(Mutex);
+                if (!Connected) {
+                    return false;
+                }
+                if (replyCallback) {
+                    RegisterLocked(ev->Sender, std::move(replyCallback));
+                }
+            }
             // TODO(v2 data plane): hand the event off to the session's outgoing stream.
             Y_UNUSED(ev);
             return true;
         }
 
-        void RegisterReceiveCallback(const TActorId& localActorId, TReceiveCallback callback) override {
+        void RegisterReceiveCallback(const TActorId& localActorId, TIntrusivePtr<IReceiveCallback> callback) override {
             TGuard<TMutex> guard(Mutex);
             if (!Connected) {
                 return;
             }
-            Callbacks[localActorId] = std::move(callback);
+            RegisterLocked(localActorId, std::move(callback));
         }
 
         void UnregisterReceiveCallback(const TActorId& localActorId) override {
             TGuard<TMutex> guard(Mutex);
-            Callbacks.erase(localActorId);
+            if (Callbacks.erase(localActorId)) {
+                CallbackCount.store(Callbacks.size(), std::memory_order_release);
+            }
         }
 
-        // Delivers an incoming event to the callback registered for its Recipient, if any. Intended to
-        // be called on the session actor thread by the (future) v2 data plane. Returns true if a
-        // callback consumed the event. Thread-safe.
-        bool DeliverIncoming(TAutoPtr<IEventHandle> ev) {
-            TReceiveCallback callback;
+        bool DeliverIncoming(TAutoPtr<IEventHandle>& ev) {
+            if (CallbackCount.load(std::memory_order_acquire) == 0) {
+                return false;
+            }
+            TIntrusivePtr<IReceiveCallback> callback;
             {
                 TGuard<TMutex> guard(Mutex);
                 if (!Connected) {
                     return false;
                 }
-                const auto it = Callbacks.find(ev->GetRecipientRewrite());
+                // just received -> Recipient and RecipientRewrite are guaranteed to match
+                const auto it = Callbacks.find(ev->Recipient);
                 if (it == Callbacks.end()) {
                     return false;
                 }
                 callback = it->second;
             }
-            // invoke outside the lock to avoid re-entrancy/deadlocks
-            callback(std::move(ev));
+            callback->Receive(std::move(ev));
             return true;
         }
 
-        // Transitions the handle into the disconnected state. Called by the session on teardown.
-        // After this returns, Send() fails and no callbacks remain registered. Thread-safe.
         void Shutdown() {
             TGuard<TMutex> guard(Mutex);
             Connected = false;
             Callbacks.clear();
+            CallbackCount.store(0, std::memory_order_release);
         }
 
-        bool IsConnected() const {
-            TGuard<TMutex> guard(Mutex);
-            return Connected;
+    private:
+        // Mutex must be held and Connected must be true.
+        void RegisterLocked(const TActorId& localActorId, TIntrusivePtr<IReceiveCallback> callback) {
+            Callbacks[localActorId] = std::move(callback);
+            CallbackCount.store(Callbacks.size(), std::memory_order_release);
         }
 
     private:
         mutable TMutex Mutex;
         bool Connected = true;
-        THashMap<TActorId, TReceiveCallback, TActorId::THash> Callbacks;
+        std::atomic<size_t> CallbackCount = 0;
+        THashMap<TActorId, TIntrusivePtr<IReceiveCallback>, TActorId::THash> Callbacks;
     };
 
 }
diff --git a/ydb/library/actors/interconnect/interconnect_tcp_input_session.cpp b/ydb/library/actors/interconnect/interconnect_tcp_input_session.cpp
index 44269070717..83c7fdcb4e7 100644
--- a/ydb/library/actors/interconnect/interconnect_tcp_input_session.cpp
+++ b/ydb/library/actors/interconnect/interconnect_tcp_input_session.cpp
@@ -1232,7 +1232,12 @@ namespace NActors {
                 ev.reset();
             }
             if (ev) {
-                TActivationContext::Send(ev.release());
+                // Give the direct-session interface a chance to consume the event via a registered
+                // receiver; otherwise fall back to normal actor-system delivery.
+                TAutoPtr<IEventHandle> evPtr(ev.release());
+                if (!Context->DirectSession->DeliverIncoming(evPtr)) {
+                    TActivationContext::Send(evPtr.Release());
+                }
             }
         }
     }
diff --git a/ydb/library/actors/interconnect/interconnect_tcp_session.cpp b/ydb/library/actors/interconnect/interconnect_tcp_session.cpp
index fff6619d322..2e0c20beb5e 100644
--- a/ydb/library/actors/interconnect/interconnect_tcp_session.cpp
+++ b/ydb/library/actors/interconnect/interconnect_tcp_session.cpp
@@ -69,6 +69,8 @@ namespace NActors {
         Params = params;
         Proxy->Metrics->SetPeerScopeId(Params.PeerScopeId);
         Proxy->Metrics->SetConnected(0);
+        DirectSession = std::make_shared<TDirectSessionV1>(TActivationContext::ActorSystem(), SelfId(), Proxy->PeerNodeId);
+        ReceiveContext->DirectSession = DirectSession;
         auto destroyCallback = [as = TActivationContext::ActorSystem(), id = Proxy->Common->DestructorId](THolder<IEventBase> event) {
             as->Send(id, event.Release());
         };
@@ -81,6 +83,12 @@ namespace NActors {
         SendUpdateToWhiteboard();
     }
 
+    void TInterconnectSessionTCP::HandleProcessDirectSessionQueue() {
+        // Runs on the shared session/input-session mailbox thread, so applying registration commands
+        // here does not race with incoming-event dispatch.
+        DirectSession->ApplyPendingCommands();
+    }
+
     std::optional<ui8> TInterconnectSessionTCP::GetXDCFlags() const noexcept {
         std::optional<ui8> flags;
         using NInterconnect::NRdma::TQueuePair;
@@ -141,6 +149,10 @@ namespace NActors {
         IActor::InvokeOtherActor(*Proxy, &TInterconnectProxyTCP::UnregisterSession, this);
         ShutdownSocket(std::move(reason));
 
+        // disconnect the direct interface so racing user threads observe a clean shutdown; the handle
+        // itself stays published in ReceiveContext for the input session's whole lifetime
+        DirectSession->Shutdown();
+
         ReceiveContext->Terminated = true; // prevent further message generation by receiving actor
         for (const auto& kv : Subscribers) {
             Send(kv.first, new TEvInterconnect::TEvNodeDisconnected(Proxy->PeerNodeId), 0, kv.second.Cookie);
@@ -258,7 +270,7 @@ namespace NActors {
         LOG_DEBUG_IC_SESSION("ICS12", "subscribe for session state for %s", msg->Event->Sender.ToString().data());
         UpdateSubscriber(msg->Event->Sender, msg->Event->Cookie, msg->ActivityIndex, std::move(msg->EventTypeName),
             std::move(msg->StackTrace));
-        Send(msg->Event->Sender, new TEvInterconnect::TEvNodeConnected(Proxy->PeerNodeId), 0, msg->Event->Cookie);
+        Send(msg->Event->Sender, new TEvInterconnect::TEvNodeConnected(Proxy->PeerNodeId, DirectSession), 0, msg->Event->Cookie);
 
         EnqueueForward(TAutoPtr<IEventHandle>(msg->Event.Release()));
     }
@@ -291,7 +303,7 @@ namespace NActors {
     void TInterconnectSessionTCP::Subscribe(STATEFN_SIG) {
         LOG_DEBUG_IC_SESSION("ICS04", "subscribe for session state for %s", ev->Sender.ToString().data());
         UpdateSubscriber(ev->Sender, ev->Cookie);
-        Send(ev->Sender, new TEvInterconnect::TEvNodeConnected(Proxy->PeerNodeId), 0, ev->Cookie);
+        Send(ev->Sender, new TEvInterconnect::TEvNodeConnected(Proxy->PeerNodeId, DirectSession), 0, ev->Cookie);
     }
 
     void TInterconnectSessionTCP::Unsubscribe(STATEFN_SIG) {
diff --git a/ydb/library/actors/interconnect/interconnect_tcp_session.h b/ydb/library/actors/interconnect/interconnect_tcp_session.h
index f016d60885f..14bc354a6f9 100644
--- a/ydb/library/actors/interconnect/interconnect_tcp_session.h
+++ b/ydb/library/actors/interconnect/interconnect_tcp_session.h
@@ -38,6 +38,7 @@
 #include "channel_scheduler.h"
 #include "outgoing_stream.h"
 #include "interconnect_session_iface.h"
+#include "interconnect_direct_session.h"
 
 #include <unordered_set>
 #include <unordered_map>
@@ -188,6 +189,11 @@ namespace NActors {
         ui64 LastProcessedSerial = 0;
         bool Terminated = false;
 
+        // Direct send/receive handle shared with users via TEvNodeConnected. The input session
+        // consults it before dispatching an incoming event into the actor system; if a receiver is
+        // registered for the recipient, the event is delivered directly instead.
+        std::shared_ptr<TDirectSessionV1> DirectSession;
+
         TReceiveContext() {
             GetTimeFast(&StartTime);
         }
@@ -580,6 +586,7 @@ namespace NActors {
                 hFunc(TEvUringRegisterFailed, Handle)
                 hFunc(TEvUringWriteComplete, Handle)
                 hFunc(TEvUringSendZcNotif, Handle)
+                cFunc(static_cast<ui32>(ENetwork::EvProcessDirectSessionQueue), HandleProcessDirectSessionQueue)
             )
             UpdateUtilization();
         }
@@ -795,6 +802,14 @@ namespace NActors {
         TActorId ReceiverId;
         TDuration Ping;
 
+        // Direct send/receive handle handed out to subscribers via TEvNodeConnected; also published
+        // into ReceiveContext so the input session can intercept incoming events.
+        std::shared_ptr<TDirectSessionV1> DirectSession;
+
+        // Drains the direct session's lock-free registration queue on the session/input-session mailbox
+        // thread (woken via EvProcessDirectSessionQueue by producers).
+        void HandleProcessDirectSessionQueue();
+
         ui64 ConfirmPacketsForcedBySize = 0;
         ui64 ConfirmPacketsForcedByTimeout = 0;
 
diff --git a/ydb/library/actors/interconnect/interconnect_tcp_session_v2.cpp b/ydb/library/actors/interconnect/interconnect_tcp_session_v2.cpp
index abef5b4b2b9..fea8d1ca9e8 100644
--- a/ydb/library/actors/interconnect/interconnect_tcp_session_v2.cpp
+++ b/ydb/library/actors/interconnect/interconnect_tcp_session_v2.cpp
@@ -18,7 +18,7 @@ namespace NActors {
         Proxy->Metrics->SetConnected(0);
         SetPrefix(Sprintf("SessionV2 %s [node %" PRIu32 "]", SelfId().ToString().data(), Proxy->PeerNodeId));
         LOG_INFO_IC_SESSION("ICS90", "v2 session created");
-        DirectSession = std::make_shared<TDirectSession>();
+        DirectSession = std::make_shared<TDirectSessionV2>();
     }
 
     void TInterconnectSessionTCPv2::SetNewConnection(TEvHandshakeDone::TPtr& ev) {
diff --git a/ydb/library/actors/interconnect/interconnect_tcp_session_v2.h b/ydb/library/actors/interconnect/interconnect_tcp_session_v2.h
index 59ed4c1f499..06e02020260 100644
--- a/ydb/library/actors/interconnect/interconnect_tcp_session_v2.h
+++ b/ydb/library/actors/interconnect/interconnect_tcp_session_v2.h
@@ -91,7 +91,7 @@ namespace NActors {
         TIntrusivePtr<NInterconnect::TStreamSocket> XdcSocket;
 
         // thread-safe direct send/receive handle shared with users via TEvNodeConnected
-        std::shared_ptr<TDirectSession> DirectSession;
+        std::shared_ptr<TDirectSessionV2> DirectSession;
 
         // subscribers awaiting connection state notifications (actor id -> cookie)
         std::unordered_map<TActorId, ui64, TActorId::THash> Subscribers;