schedulable.h

Go to the documentation of this file.

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */
 
#pragma once
 
#include <cassert>
#include <chrono>
#include <condition_variable>
#include <mutex>
#include <thread>
 
#include <dispenso/detail/completion_event_impl.h>
#include <dispenso/task_set.h>
 
namespace dispenso {
 
class ImmediateInvoker {
 public:
  template <typename F>
  DISPENSO_REQUIRES(OnceCallableFunc<F>)
  void schedule(F&& f) const {
    f();
  }
 
  template <typename F>
  DISPENSO_REQUIRES(OnceCallableFunc<F>)
  void schedule(F&& f, ForceQueuingTag) const {
    f();
  }
};
 
constexpr ImmediateInvoker kImmediateInvoker;
 
class NewThreadInvoker {
 public:
  template <typename F>
  DISPENSO_REQUIRES(OnceCallableFunc<F>)
  void schedule(F&& f) const {
    schedule(std::forward<F>(f), ForceQueuingTag());
  }
  template <typename F>
  DISPENSO_REQUIRES(OnceCallableFunc<F>)
  void schedule(F&& f, ForceQueuingTag) const {
    auto* waiter = getWaiter();
    waiter->add();
    std::thread thread([f = std::move(f), waiter]() {
      // RAII so remove() runs even if f() throws. In practice an uncaught exception out of a
      // std::thread function calls std::terminate -> abort, which kills the process before the
      // atexit waiter could block, but the guard documents the invariant and covers exotic
      // terminate-handler setups.
      RemoveGuard guard{waiter};
      f();
    });
    thread.detach();
  }
 
 private:
  // DO NOT REMOVE WITHOUT READING THE FOLLOWING:
  //
  // schedule() above spawns a detached std::thread per call. On Windows shared-lib builds, if the
  // process exits while a detached thread is mid-execution inside dispenso's code, the OS unmaps
  // dispenso.dll's pages during DLL_PROCESS_DETACH while the thread is still executing them →
  // EXCEPTION_ACCESS_VIOLATION (0xC0000005). On any platform, the detached thread also has a
  // non-trivial window between Future::status notify(kReady) and the final
  // decRefCountMaybeDestroy + thread-local-storage teardown, which extends past the user-visible
  // future.get() return.
  //
  // Two mitigations, both needed:
  //   1. pinModuleForNewThread() (schedulable.cpp, Windows-only) pins the module that contains
  //      dispenso's code so it is never unmapped while a detached thread is still executing in it.
  //      This removes the access-violation regardless of how long the thread runs, and also covers
  //      a host FreeLibrary() at runtime, which the atexit path cannot. (Only Windows eagerly
  //      unmaps a module's code under a running thread; POSIX/other loaders don't, so no analog.)
  //   2. ThreadWaiter gives outstanding threads a BOUNDED grace period at exit to reach remove(),
  //      narrowing the post-future.get() teardown window. It is bounded (never blocks shutdown
  //      indefinitely) precisely because (1) makes a thread that runs past the deadline non-fatal.
  // The SmallBufferAllocator controlled-leak fix (small_buffer_allocator.cpp) addresses a separate
  // hazard (returning small buffers to a destroyed central store) and is also needed.
  //
  // The relevant tests are future_test_sans_exceptions and future_shared_test, the
  // Future.AsyncNotAsyncSpecifyNewThread / NewThreadInvoker / AsyncSpecifyNewThread cases.
  //
  // The ThreadWaiter object itself is intentionally leaked (controlled-leak singleton in
  // schedulable.cpp). Deleting it on shutdown would create a UAF window for any post-atexit
  // schedule() call (external thread, static destructor) that hits a freed waiter.
  struct ThreadWaiter {
    int count_ = 0;
    std::mutex mtx_;
    std::condition_variable cond_;
 
    void add() DISPENSO_NO_THREAD_SAFETY_ANALYSIS {
      std::lock_guard<std::mutex> lk(mtx_);
      ++count_;
    }
 
    void remove() DISPENSO_NO_THREAD_SAFETY_ANALYSIS {
      std::lock_guard<std::mutex> lk(mtx_);
      assert(count_ > 0 && "remove() called without matching add()");
      if (--count_ == 0) {
        cond_.notify_one();
      }
    }
 
    void wait() {
      std::unique_lock<std::mutex> lk(mtx_);
      // Bounded best-effort: pinModuleForNewThread() keeps our code mapped, so a thread still
      // running past this deadline cannot fault on unload — this only narrows the teardown window
      // and must never wedge shutdown, so it times out rather than blocking forever.
      cond_.wait_for(lk, std::chrono::seconds(2), [this]() { return count_ == 0; });
    }
  };
 
  struct RemoveGuard {
    ThreadWaiter* w;
    ~RemoveGuard() {
      w->remove();
    }
  };
 
  DISPENSO_DLL_ACCESS static ThreadWaiter* getWaiter();
};
 
constexpr NewThreadInvoker kNewThreadInvoker;
 
} // namespace dispenso