From 86174c3e5fddc2bf0e583dc7ed39984bfd61891e Mon Sep 17 00:00:00 2001
From: Chris Hopman <cjhopman@meta.com>
Date: Fri, 20 Dec 2024 11:14:46 -0800
Subject: [PATCH] Rearrange code + add simple documentation

Summary:
The main thing in this module is going to end up being three views of
this execution context, moving them to the top of the file.

Reviewed By: JakobDegen

Differential Revision: D65362045

fbshipit-source-id: 501240ccd55a358ea59b80a90fd169c39752df6b
---
 app/buck2_futures/src/details/shared_state.rs | 151 +++++++++---------
 1 file changed, 78 insertions(+), 73 deletions(-)

diff --git a/app/buck2_futures/src/details/shared_state.rs b/app/buck2_futures/src/details/shared_state.rs
index 764ca6a2fa35a..01594fb65493b 100644
--- a/app/buck2_futures/src/details/shared_state.rs
+++ b/app/buck2_futures/src/details/shared_state.rs
@@ -24,6 +24,84 @@ use slab::Slab;
 
 use crate::cancellation::CriticalSectionGuard;
 
+/// Shared cancellation related execution context for a cancellable task.
+///
+/// This is the "outer" context, held by the task's spawned future impl within
+/// its poll loop.
+pub(crate) struct ExecutionContextOuter {
+    shared: Arc<Mutex<ExecutionContextData>>,
+}
+
+impl ExecutionContextOuter {
+    pub(crate) fn new() -> (ExecutionContextOuter, ExecutionContextInner) {
+        let shared = Arc::new(Mutex::new(ExecutionContextData {
+            cancellation_notification: {
+                CancellationNotificationData {
+                    inner: Arc::new(CancellationNotificationDataInner {
+                        notified: Default::default(),
+                        wakers: Mutex::new(Some(Default::default())),
+                    }),
+                }
+            },
+            prevent_cancellation: 0,
+        }));
+        (
+            ExecutionContextOuter {
+                shared: shared.dupe(),
+            },
+            ExecutionContextInner { shared },
+        )
+    }
+
+    pub(crate) fn notify_cancelled(&self) -> bool {
+        let mut lock = self.shared.lock();
+        if lock.can_exit() {
+            true
+        } else {
+            lock.notify_cancelled();
+            false
+        }
+    }
+
+    pub(crate) fn can_exit(&self) -> bool {
+        self.shared.lock().can_exit()
+    }
+}
+
+/// Shared cancellation related execution context for a cancellable task.
+///
+/// This is the "inner" context, used by a CancellationContext to observe cancellation
+/// and enter critical sections.
+pub(crate) struct ExecutionContextInner {
+    shared: Arc<Mutex<ExecutionContextData>>,
+}
+
+impl ExecutionContextInner {
+    pub(crate) fn enter_structured_cancellation(&self) -> CriticalSectionGuard {
+        let mut shared = self.shared.lock();
+
+        let notification = shared.enter_structured_cancellation();
+
+        CriticalSectionGuard::new_explicit(self, notification)
+    }
+
+    pub(crate) fn try_to_disable_cancellation(&self) -> bool {
+        let mut shared = self.shared.lock();
+        if shared.try_to_disable_cancellation() {
+            true
+        } else {
+            // couldn't prevent cancellation, so release our hold onto the counter
+            shared.exit_prevent_cancellation();
+            false
+        }
+    }
+
+    pub(crate) fn exit_prevent_cancellation(&self) -> bool {
+        let mut shared = self.shared.lock();
+        shared.exit_prevent_cancellation()
+    }
+}
+
 pub(crate) struct SharedState {
     inner: Arc<SharedStateData>,
 }
@@ -182,79 +260,6 @@ impl ExecutionContextData {
     }
 }
 
-/// Context relating to execution of the `poll` of the future. This will contain the information
-/// required for the `CancellationContext` that the future holds to enter critical sections and
-/// structured cancellations.
-pub(crate) struct ExecutionContextOuter {
-    shared: Arc<Mutex<ExecutionContextData>>,
-}
-
-pub(crate) struct ExecutionContextInner {
-    shared: Arc<Mutex<ExecutionContextData>>,
-}
-
-impl ExecutionContextOuter {
-    pub(crate) fn new() -> (ExecutionContextOuter, ExecutionContextInner) {
-        let shared = Arc::new(Mutex::new(ExecutionContextData {
-            cancellation_notification: {
-                CancellationNotificationData {
-                    inner: Arc::new(CancellationNotificationDataInner {
-                        notified: Default::default(),
-                        wakers: Mutex::new(Some(Default::default())),
-                    }),
-                }
-            },
-            prevent_cancellation: 0,
-        }));
-        (
-            ExecutionContextOuter {
-                shared: shared.dupe(),
-            },
-            ExecutionContextInner { shared },
-        )
-    }
-
-    pub(crate) fn notify_cancelled(&self) -> bool {
-        let mut lock = self.shared.lock();
-        if lock.can_exit() {
-            true
-        } else {
-            lock.notify_cancelled();
-            false
-        }
-    }
-
-    pub(crate) fn can_exit(&self) -> bool {
-        self.shared.lock().can_exit()
-    }
-}
-
-impl ExecutionContextInner {
-    pub(crate) fn enter_structured_cancellation(&self) -> CriticalSectionGuard {
-        let mut shared = self.shared.lock();
-
-        let notification = shared.enter_structured_cancellation();
-
-        CriticalSectionGuard::new_explicit(self, notification)
-    }
-
-    pub(crate) fn try_to_disable_cancellation(&self) -> bool {
-        let mut shared = self.shared.lock();
-        if shared.try_to_disable_cancellation() {
-            true
-        } else {
-            // couldn't prevent cancellation, so release our hold onto the counter
-            shared.exit_prevent_cancellation();
-            false
-        }
-    }
-
-    pub(crate) fn exit_prevent_cancellation(&self) -> bool {
-        let mut shared = self.shared.lock();
-        shared.exit_prevent_cancellation()
-    }
-}
-
 pub(crate) enum CancellationNotificationStatus {
     /// no notifications yet. maps to '0'
     Pending,