thread_pool.h

Go to the documentation of this file.

// Copyright Microsoft and CHERIoT Contributors.
// SPDX-License-Identifier: MIT
 
#pragma once
#include <__cheri_sealed.h>
#include <cdefs.h>
 
/**
 * \file
 *
 * APIs for an example thread pool.  This does not provide any availability
 * guarantees (work dispatched to it can block a thread indefinitely) so this
 * should be used with care.
 *
 * The thread pool exposed by this header is implemented in the `thread_pool`
 * compartment, which must be linked to the final firmware image.
 */
 
/**
 * The type of a thread-pool callback.  This is a CHERI callback so that it can
 * run in the compartment that schedule the work to run, but a different
 * thread.
 */
typedef __cheri_callback void (*ThreadPoolCallback)(CHERI_SEALED(void *));
 
/**
 * Message to a thread pool.  This encapsulates a simple closure that will be
 * invoked in the next available thread in a thread pool.
 */
struct ThreadPoolMessage
{
    /**
     * The function pointer that will be called in the thread pool.
     */
    ThreadPoolCallback invoke;
    /**
     * The data associated with the asynchronous invocation.
     */
    CHERI_SEALED(void *) data;
};
 
__BEGIN_DECLS
 
/**
 * Invoke a function that takes a `void*` argument in another thread.  The
 * function will be invoked with `data` as the argument.  The `data` argument
 * must be sealed.
 *
 * This can block indefinitely until the thread pool is able to process a
 * message.
 *
 * FIXME: We should add a non-blocking variant of this.
 *
 * Returns 0 on success, -EINVAL if either of the arguments are invalid.
 */
int __cheri_compartment("thread_pool")
  thread_pool_async(ThreadPoolCallback fn, CHERI_SEALED(void *) data);
 
/**
 * Run a thread pool.  This does not return, despite the claimed type, and can
 * be used as a thread entry point.
 */
int __cheri_compartment("thread_pool") thread_pool_run(void);
__END_DECLS
 
#ifdef __cplusplus
#   include <stdlib.h>
#   include <token.h>
#   include <type_traits>
#   include <cheri.hh>
/**
 * For C++ programmers, we provide some more user-friendly wrappers.
 */
namespace thread_pool
{
    /**
     * Implementation details, should not be used outside of this header.
     */
    namespace detail
    {
        /**
         * Helper that generates a different sealing key per type using the
         * allocator's token mechanism.
         */
        template<typename T>
        inline TokenKey sealing_key_for_type()
        {
            static TokenKey key = token_key_new();
            return key;
        }
 
        /**
         * Helper that provides a callback function for invoking and then
         * freeing a sealed heap-allocated lambda.  The lambda is sealed with
         * the software-defined sealing key constructed by the
         * `sealing_key_for_type` helper.
         */
        template<typename T>
        __cheri_callback void wrap_callback_lambda(CHERI_SEALED(void *) rawFn)
        {
            auto key = sealing_key_for_type<T>();
            auto fn  = token_unseal(
              key, Sealed<T>(static_cast<CHERI_SEALED(T *)>(rawFn)));
            if (fn == nullptr)
            {
                return;
            }
            (*fn)();
            /*
             * This fails only if the thread pool runner compartment can't make
             * cross-compartment calls to the allocator at all, since we're
             * in its initial trusted activation frame and near the beginning
             * (highest address) of its stack.
             *
             * Well, mostly.  It can also fail if the allocator has stack usage
             * checks compiled in and we have forgotten to bump that value.
             */
            (void)token_obj_destroy(MALLOC_CAPABILITY, key, rawFn);
        }
 
        /**
         * Helper to wrap pure C function (including a stateless lambda) as a
         * callback.
         */
        template<typename Fn>
        __cheri_callback void wrap_callback_function(CHERI_SEALED(void *))
        {
            // C++ objects are guaranteed to have unique addresses and so a
            // zero-sized object is actually 1 byte and using `sizeof(Fn) == 1`
            // would not let us tell the difference between empty lambdas and
            // ones with a one-byte capture.  To ensure that this object is
            // really stateless, we make it a field of another type, marking it
            // as not requiring a unique address, and then check that the
            // offset of the following field is 0.
            struct Test
            {
                // The stateless callable object
                [[no_unique_address]] Fn f;
                // An field that can share the same address as `f` if `f` is
                // truly stateless.
                char b;
            };
            static_assert(offsetof(Test, b) == 0,
                          "Stateless callable object is not stateless");
            // Invoke an empty instance of this stateless callable object.
            Fn{}();
        }
 
    } // namespace detail
 
    /**
     * Asynchronously invoke a lambda.  This moves the lambda to the heap and
     * passes it to the thread pool's queue.  If the lambda copies any stack
     * objects by reference then the copy will fail.
     *
     * Returns 0 on success, or compartment-call failures (ENOTENOUGHSTACK,
     * ENOTENOUGHTRUSTEDSTACK) if the thread pool cannot be invoked.
     */
    template<typename T>
    int async(T &&lambda)
    {
        // If this is a stateless function, just send a callback function
        // pointer, don't copy zero bytes of state to the heap.
        if constexpr (std::is_convertible_v<T, void (*)()>)
        {
            return thread_pool_async(
              &detail::wrap_callback_function<std::remove_cvref_t<T>>, nullptr);
        }
        else
        {
            // If this is a stateful lambda, move it to the heap, create a
            // callback function that will invoke it and free it, and send a
            // pointer to the wrapper and the heap-allocated object.
            void *buffer;
            using LambdaType = std::remove_reference_t<decltype(lambda)>;
            // Allocate a new sealed object with a key that is unique to this
            // type.
            Timeout t{UnlimitedTimeout};
            CHERI_SEALED(void *)
            sealed = token_sealed_unsealed_alloc(
              &t,
              MALLOC_CAPABILITY,
              detail::sealing_key_for_type<LambdaType>(),
              sizeof(lambda),
              &buffer);
            /*
             * Copy the lambda into the new allocation.
             *
             * If the above allocation has failed, this will trap.
             *
             * Note: We silence a warning here because we *do* want to
             * explicitly move, not forward.
             */
            T *copy = new (buffer) T(
              std::move(lambda)); // NOLINT(bugprone-move-forwarding-reference)
            // Create the wrapper that will unseal and invoke the lambda.
            ThreadPoolCallback invoke =
              &detail::wrap_callback_lambda<LambdaType>;
            // Dispatch it.
            return thread_pool_async(invoke, sealed);
        }
    }
} // namespace thread_pool
 
#endif