[V4] Null checkpoints (#75)

* definition of the null state

* idea comment

* starting the null stack

* move initialization inside constructor

* move outside of frame

* don't set default

* static assert to catch bad catagory

* todo

* spell

* generalise lock/key

* use opaque checkpoint

* acq/rel

* move init to constructor

* get context/alloc helpers

* assume empty on acquire

* docs

* spell

* fix empty bug + expose

* strong exception safety

* add mission import

* spell

Author: ConorWilliams

benchmark/src/libfork_benchmark/fib/lf_parts.cpp

@@ -86,7 +86,6 @@ struct deque_ctx {
   void post(lf::await_handle<deque_ctx>) {}
 
   // TODO: try LF_NO_INLINE for final allocator
-  LF_NO_INLINE
   void push(handle_type handle) { work.push(handle); }
 
   auto pop() noexcept -> handle_type {

src/core/concepts.cxx

@@ -62,19 +62,23 @@ consteval auto constify(T &&x) noexcept -> std::add_const_t<T> &;
  *
  * // TODO: define if release is required before acquire?
  *
- * - After construction `this` is in the empty state and push is valid.
+ * - After construction `this` is empty and push is valid.
  * - Pop is valid provided the FILO order is respected.
  * - Push produces pointers aligned to __STDCPP_DEFAULT_NEW_ALIGNMENT__.
  * - Destruction is expected to only occur when the stack is empty.
  * - Result of `.checkpoint()` is expected to:
  *     - Be "cheap to copy".
- *     - Compare equal if and only if they belong to the same stack.
+ *     - Have a null state (default constructed) that only compares equal to itself.
+ *     - Is allowed to return null if push has never been called.
+ *     - Compare equal if and only if they belong to the same stack or are both null.
  *     - Have no preconditions about when it's called.
- *     - A default constructed checkpoint considered undefined
- * - Release detaches the current stack and leaves `this` in the empty state.
+ * - Prepare release puts the stack into a state which another thread can acquire it.
+ * - Release detaches the current stack and leaves `this` empty.
+ *     - This may be called concurrently with acquire
  * - Acquire attaches to the stack that the checkpoint came from:
- *     - This is a noop if the checkpoint is from the current stack.
- *     - Otherwise `this` is empty.
+ *     - It is only called the stack is empty.
+ *     - It is only called with a checkpoint not equal to the current checkpoint.
+ *     - It is called after prepare release (and no other functions in between)
  *
  * Fast-path operations: empty, push, pop, checkpoint
  * Slow-path operations: release, acquire

src/core/frame.cxx

@@ -7,6 +7,8 @@ import std;
 
 import :concepts;
 import :constants;
+import :utility;
+import :thread_locals;
 
 namespace lf {
 // =================== Cancellation =================== //
@@ -72,13 +74,16 @@ struct frame_type {
 
   // == Member variables == //
 
+  // TODO: add checked accessors for all the things (including except etc)
+
   union {
     parent_union parent;
     except_type *except;
   };
 
   cancellation *cancel;
 
+  // TODO: drop default constructible requirement?
   [[no_unique_address]]
   checkpoint_type stack_ckpt;
 
@@ -92,7 +97,7 @@ struct frame_type {
   // Explicitly post construction, this allows the compiler to emit a single
   // instruction for the zero init then an instruction for the joins init,
   // instead of three instructions.
-  constexpr frame_type() noexcept { joins = k_u16_max; }
+  constexpr frame_type(checkpoint_type &&ckpt) noexcept : stack_ckpt(std::move(ckpt)) { joins = k_u16_max; }
 
   [[nodiscard]]
   constexpr auto is_cancelled() const noexcept -> bool {

src/core/handles.cxx

@@ -3,15 +3,10 @@ export module libfork.core:handles;
 import std;
 
 import :frame;
+import :utility;
 
 namespace lf {
 
-// =================== Lock and key =================== //
-
-struct lock {};
-
-inline constexpr lock key = {};
-
 // =================== Frame =================== //
 
 // TODO: api + test this is lock-free
@@ -28,7 +23,7 @@ export template <typename T>
 class frame_handle {
  public:
   constexpr frame_handle() = default;
-  constexpr frame_handle(lock, frame_type<T> *ptr) noexcept : m_ptr{ptr} {}
+  constexpr frame_handle(key_t, frame_type<T> *ptr) noexcept : m_ptr{ptr} {}
 
   constexpr auto operator==(frame_handle const &) const noexcept -> bool = default;
 
@@ -48,7 +43,7 @@ export template <typename T>
 class await_handle {
  public:
   constexpr await_handle() = default;
-  constexpr await_handle(lock, frame_type<T> *ptr) noexcept : m_ptr{ptr} {}
+  constexpr await_handle(key_t, frame_type<T> *ptr) noexcept : m_ptr{ptr} {}
 
   constexpr auto operator==(await_handle const &) const noexcept -> bool = default;
 

src/core/promise.cxx

@@ -119,15 +119,20 @@ final_suspend_continue(Context *context, frame_type<Context> *parent) noexcept -
   // the stack so we must prepare it for release now.
   auto release_key = context->allocator().prepare_release();
 
+  // TODO: we could add an `if (owner)` around acquire below, then we could
+  // define that acquire is always called with null or not-self.
+
   // Register with parent we have completed this child task.
   if (parent->atomic_joins().fetch_sub(1, std::memory_order_release) == 1) {
     // Parent has reached join and we are the last child task to complete. We
     // are the exclusive owner of the parent and therefore, we must continue
     // parent. As we won the race, acquire all writes before resuming.
     std::atomic_thread_fence(std::memory_order_acquire);
 
-    // In case of scenario (2) we must acquire the parent's stack.
-    context->allocator().acquire(std::as_const(parent->stack_ckpt));
+    if (!owner) {
+      // In case of scenario (2) we must acquire the parent's stack.
+      context->allocator().acquire(std::as_const(parent->stack_ckpt));
+    }
 
     // Must reset parent's control block before resuming parent.
     parent->reset_counters();
@@ -192,11 +197,11 @@ constexpr auto final_suspend(frame_type<Context> *frame) noexcept -> coro<> {
       return parent->handle();
     case category::fork:
 
-      Context *context = not_null(thread_context<Context>);
+      Context *context = get_context<Context>();
 
       if (frame_handle last_pushed = context->pop()) {
         // No-one stole continuation, we are the exclusive owner of parent -> just keep ripping!
-        LF_ASSUME(last_pushed == frame_handle{key, parent});
+        LF_ASSUME(last_pushed == frame_handle{key(), parent});
         // This is not a join point so no state (i.e. counters) is guaranteed.
         return parent->handle();
       }
@@ -243,6 +248,8 @@ constexpr void stash_current_exception(frame_type<Context> *frame) noexcept {
 template <category Cat, worker_context Context>
 struct awaitable : std::suspend_always {
 
+  static_assert(Cat == category::call || Cat == category::fork, "Invalid category for awaitable");
+
   frame_type<Context> *child;
 
   /**
@@ -274,16 +281,21 @@ struct awaitable : std::suspend_always {
     // Propagate parent->child relationships
     self.child->parent.frame = &parent.promise().frame;
     self.child->cancel = parent.promise().frame.cancel;
-    self.child->stack_ckpt = not_null(thread_context<Context>)->allocator().checkpoint();
-    self.child->kind = Cat;
+
+    if constexpr (Cat == category::call) {
+      // Should be the default
+      LF_ASSUME(self.child->kind == category::call);
+    } else {
+      self.child->kind = Cat;
+    }
 
     if constexpr (Cat == category::fork) {
       // It is critical to pass self by-value here, after the call to push()
       // the object `*this` may be destroyed, if passing by ref it would be
       // use-after-free to then access self in the following line to fetch the
       // handle.
       LF_TRY {
-        not_null(thread_context<Context>)->push(frame_handle{key, &parent.promise().frame});
+        get_context<Context>()->push(frame_handle{key(), &parent.promise().frame});
       } LF_CATCH_ALL {
         return self.stash_and_resume(parent), parent;
       }
@@ -303,7 +315,7 @@ struct join_awaitable {
   frame_type<Context> *frame;
 
   constexpr auto take_stack_and_reset(this join_awaitable self) noexcept -> void {
-    Context *context = not_null(thread_context<Context>);
+    Context *context = get_context<Context>();
     LF_ASSUME(self.frame->stack_ckpt != context->allocator().checkpoint());
     context->allocator().acquire(std::as_const(self.frame->stack_ckpt));
     self.frame->reset_counters();
@@ -437,14 +449,12 @@ struct mixin_frame {
   // --- Allocation
 
   static auto operator new(std::size_t sz) -> void * {
-    void *ptr = not_null(thread_context<Ctx>)->allocator().push(sz);
+    void *ptr = get_allocator<Ctx>().push(sz);
     LF_ASSUME(is_aligned<k_new_align>(ptr));
     return std::assume_aligned<k_new_align>(ptr);
   }
 
-  static auto operator delete(void *p, std::size_t sz) noexcept -> void {
-    not_null(thread_context<Ctx>)->allocator().pop(p, sz);
-  }
+  static auto operator delete(void *p, std::size_t sz) noexcept -> void { get_allocator<Ctx>().pop(p, sz); }
 
   // --- Await transformations
 
@@ -506,7 +516,10 @@ struct mixin_frame {
 template <worker_context Context>
 struct promise_type<void, Context> : mixin_frame<Context> {
 
-  frame_type<Context> frame;
+  // Putting init here allows:
+  //  1. Frame not no need to know about the checkpoint type
+  //  2. Compiler merge double read of thread local here and in allocator
+  frame_type<Context> frame{get_allocator<Context>().checkpoint()};
 
   constexpr auto get_return_object() noexcept -> task<void> { return access::task(this); }
 
@@ -518,7 +531,10 @@ struct promise_type<void, Context> : mixin_frame<Context> {
 template <returnable T, worker_context Context>
 struct promise_type : mixin_frame<Context> {
 
-  frame_type<Context> frame;
+  // Putting init here allows:
+  //  1. Frame not no need to know about the checkpoint type
+  //  2. Compiler merge double read of thread local here and in allocator
+  frame_type<Context> frame{get_allocator<Context>().checkpoint()};
   T *return_address;
 
   constexpr auto get_return_object() noexcept -> task<T> { return access::task(this); }

src/core/schedule.cxx

@@ -45,6 +45,8 @@ constexpr auto schedule(Context *context, Fn &&fn, Args &&...args) noexcept -> a
 
   // TODO: make sure this is exception safe and correctly qualifed
 
+  LF_ASSUME(context != nullptr);
+
   // This is what the async function will return.
   using result_type = async_result_t<Fn, Context, Args...>;
 
@@ -54,15 +56,15 @@ constexpr auto schedule(Context *context, Fn &&fn, Args &&...args) noexcept -> a
   // TODO: make sure we're cancel safe
 
   // TODO: Before doing this we must be on a valid context.
-  LF_ASSUME(thread_context<Context> == context);
+  LF_ASSUME(get_context<Context>() == context);
 
   auto *promise = access::promise<Context>(
       std::invoke(std::forward<Fn>(fn), env<Context>{}, std::forward<Args>(args)...));
 
   // TODO: expose cancellable?
   promise->frame.parent.block = root_block.get();
   promise->frame.cancel = nullptr;
-  promise->frame.stack_ckpt = thread_context<Context>->allocator().checkpoint();
+  promise->frame.stack_ckpt = get_allocator<Context>().checkpoint();
   promise->frame.kind = lf::category::root;
 
   if constexpr (!std::is_void_v<result_type>) {