Make IPoller contract stricter to avoid use-after-free

commit_hash:ae9dabf22e3b1705320c339e78c4c09e7990d84f

2 files changed, 11 insertions, 11 deletions

diff --git a/yt/yt/core/concurrency/poller.h b/yt/yt/core/concurrency/poller.h
index 2fdd16aaf47..b492ad819fe 100644
--- a/yt/yt/core/concurrency/poller.h
+++ b/yt/yt/core/concurrency/poller.h
@@ -78,12 +78,10 @@ struct IPoller
 
     //! Unregisters the previously registered entity.
     /*!
-     *  If the pollable entity was previously armed, one should unarm it first
-     *  before unregistering. Not doing so is OK for the poller, however
-     *  in this case #IPollable::OnShutdown and #IPollable::OnEvent could be invoked concurrently.
+     *  If the pollable entity was previously armed, one MUST unarm it first
+     *  before unregistering.
      *
-     *  At the same time, if the poller was properly unarmed before unregistering,
-     *  it is guaranteed that #IPollable::OnShutdown and #IPollable::OnEvent will
+     *  It is guaranteed that #IPollable::OnShutdown and #IPollable::OnEvent will
      *  not be run concurrently.
      *
      *  The entity gets unregistered asynchronously.
diff --git a/yt/yt/core/concurrency/thread_pool_poller.cpp b/yt/yt/core/concurrency/thread_pool_poller.cpp
index f42f6b22157..183edb93f7c 100644
--- a/yt/yt/core/concurrency/thread_pool_poller.cpp
+++ b/yt/yt/core/concurrency/thread_pool_poller.cpp
@@ -381,7 +381,7 @@ private:
     TNotificationHandle WakeupHandle_;
     TMpscStack<IPollablePtr> RegisterQueue_;
     TMpscStack<IPollablePtr> UnregisterQueue_;
-    THashSet<IPollablePtr> Pollables_;
+    THashMap<IPollable*, IPollablePtr> Pollables_;
 
     std::array<TPollerImpl::TEvent, MaxEventsPerPoll> PooledImplEvents_;
 
@@ -436,12 +436,14 @@ private:
                 continue;
             }
 
+            // These can trigger if user forgot to unarm the pollable.
+            YT_VERIFY(Pollables_.contains(pollable));
+            YT_VERIFY(pollable->GetRefCount() > 0);
+
             YT_LOG_TRACE("Got pollable event (Pollable: %v, Control: %v)",
                 pollable->GetLoggingTag(),
                 control);
 
-            YT_VERIFY(pollable->GetRefCount() > 0);
-
             ScheduleEvent(pollable, control);
         }
     }
@@ -470,13 +472,13 @@ private:
                 });
 
                 RegisterQueue_.DequeueAll(false, [&] (auto& pollable) {
-                    InsertOrCrash(Pollables_, std::move(pollable));
+                    EmplaceOrCrash(Pollables_, pollable.Get(), pollable);
                 });
 
                 HandleEvents(eventCount);
 
                 for (const auto& pollable : unregisterItems) {
-                    EraseOrCrash(Pollables_, pollable);
+                    EraseOrCrash(Pollables_, pollable.Get());
                 }
 
                 unregisterItems.clear();
@@ -486,7 +488,7 @@ private:
                         break;
                     }
                     // Need to unregister pollables when stopping to break reference cycles between pollables and poller.
-                    for (const auto& pollable : Pollables_) {
+                    for (const auto& [rawPollable, pollable] : Pollables_) {
                         DoUnregister(pollable);
                     }
                 }