Make TForbidContextSwitchGuard cheap

2 files changed, 36 insertions, 13 deletions

diff --git a/yt/yt/core/concurrency/fiber_scheduler_thread.cpp b/yt/yt/core/concurrency/fiber_scheduler_thread.cpp
index 32112db71a..47036ef8f6 100644
--- a/yt/yt/core/concurrency/fiber_scheduler_thread.cpp
+++ b/yt/yt/core/concurrency/fiber_scheduler_thread.cpp
@@ -889,6 +889,24 @@ void SetCurrentFiberId(TFiberId id)
 
 ////////////////////////////////////////////////////////////////////////////////
 
+thread_local bool ContextSwitchForbidden;
+
+bool IsContextSwitchForbidden()
+{
+    return ContextSwitchForbidden;
+}
+
+TForbidContextSwitchGuard::TForbidContextSwitchGuard()
+    : OldValue_(std::exchange(ContextSwitchForbidden, true))
+{ }
+
+TForbidContextSwitchGuard::~TForbidContextSwitchGuard()
+{
+    ContextSwitchForbidden = OldValue_;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 bool CheckFreeStackSpace(size_t space)
 {
     auto* currentFiber = NDetail::TryGetCurrentFiber();
@@ -915,6 +933,7 @@ TFiberCanceler GetCurrentFiberCanceler()
 
 void WaitUntilSet(TFuture<void> future, IInvokerPtr invoker)
 {
+    YT_VERIFY(!IsContextSwitchForbidden());
     YT_VERIFY(future);
     YT_ASSERT(invoker);
 
@@ -1020,10 +1039,6 @@ TOneShotContextSwitchGuard::TOneShotContextSwitchGuard(TContextSwitchHandler out
     , Active_(true)
 { }
 
-TForbidContextSwitchGuard::TForbidContextSwitchGuard()
-    : TOneShotContextSwitchGuard([] { YT_ABORT(); })
-{ }
-
 ////////////////////////////////////////////////////////////////////////////////
 
 } //namespace NYT::NConcurrency
diff --git a/yt/yt/core/concurrency/scheduler_api.h b/yt/yt/core/concurrency/scheduler_api.h
index 9556538465..608dd805b2 100644
--- a/yt/yt/core/concurrency/scheduler_api.h
+++ b/yt/yt/core/concurrency/scheduler_api.h
@@ -38,6 +38,23 @@ void SetCurrentFiberId(TFiberId id);
 
 ////////////////////////////////////////////////////////////////////////////////
 
+//! Returns |true| if fiber context switch is currently forbidden.
+bool IsContextSwitchForbidden();
+
+class TForbidContextSwitchGuard
+{
+public:
+    TForbidContextSwitchGuard();
+    TForbidContextSwitchGuard(const TForbidContextSwitchGuard&) = delete;
+
+    ~TForbidContextSwitchGuard();
+
+private:
+    const bool OldValue_;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
 // NB: Use function pointer to minimize the overhead.
 using TGlobalContextSwitchHandler = void(*)();
 
@@ -70,15 +87,6 @@ private:
     bool Active_;
 };
 
-class TForbidContextSwitchGuard
-    : public TOneShotContextSwitchGuard
-{
-public:
-    TForbidContextSwitchGuard();
-};
-
-////////////////////////////////////////////////////////////////////////////////
-
 //! Blocks the current fiber until #future is set.
 //! The fiber is resceduled to #invoker.
 void WaitUntilSet(