Improve coroutine actors with pre-leave and pre-enter hooks

2 files changed, 43 insertions, 0 deletions

diff --git a/library/cpp/actors/core/actor_coroutine.cpp b/library/cpp/actors/core/actor_coroutine.cpp
index 51e21e2950..726c644016 100644
--- a/library/cpp/actors/core/actor_coroutine.cpp
+++ b/library/cpp/actors/core/actor_coroutine.cpp
@@ -120,7 +120,13 @@ namespace NActors {
     void TActorCoroImpl::ReturnToActorSystem() {
         TExceptionSafeContext* returnContext = std::exchange(ActorSystemContext, nullptr);
         Y_VERIFY(returnContext);
+        if (StoreTlsState) {
+            StoreTlsState(this);
+        }
         FiberContext.SwitchTo(returnContext);
+        if (RestoreTlsState) {
+            RestoreTlsState(this);
+        }
         if (!PendingEvent) {
             // we have returned from the actor system and it kindly asks us to terminate the coroutine as it is being
             // stopped
diff --git a/library/cpp/actors/core/actor_coroutine.h b/library/cpp/actors/core/actor_coroutine.h
index 7a7844cc85..5e42247374 100644
--- a/library/cpp/actors/core/actor_coroutine.h
+++ b/library/cpp/actors/core/actor_coroutine.h
@@ -25,6 +25,43 @@ namespace NActors {
         TActorIdentity SelfActorId = TActorIdentity(TActorId());
         TActorId ParentActorId;
 
+        // Pre-leave and pre-enter hook functions are called by coroutine actor code to conserve the state of required TLS
+        // variables.
+        //
+        // They are called in the following order:
+        //
+        // 1. coroutine executes WaitForEvent
+        // 2. StoreTlsState() is called
+        // 3. control is returned to the actor system
+        // 4. some event is received, handler called (now in different thread, unconserved TLS variables are changed!)
+        // 5. handler transfers control to the coroutine
+        // 6. RestoreTlsState() is called
+        //
+        // These hooks may be used in the following way:
+        //
+        // thread_local TMyClass *MyObject = nullptr;
+        //
+        // class TMyCoroImpl : public TActorCoroImpl {
+        //     TMyClass *SavedMyObject;
+        //     ...
+        // public:
+        //     TMyCoroImpl()
+        //         : TActorCoroImpl(...)
+        //     {
+        //         StoreTlsState = RestoreTlsState = &TMyCoroImpl::ConserveState;
+        //     }
+        //
+        //     static void ConserveState(TActorCoroImpl *p) {
+        //         TMyCoroImpl *my = static_cast<TMyCoroImpl*>(p);
+        //         std::swap(my->SavedMyObject, MyObject);
+        //     }
+        //
+        //     ...
+        // }
+        void (*StoreTlsState)(TActorCoroImpl*) = nullptr;
+        void (*RestoreTlsState)(TActorCoroImpl*) = nullptr;
+
+
     private:
         template <typename TFirstEvent, typename... TOtherEvents>
         struct TIsOneOf: public TIsOneOf<TOtherEvents...> {