fiber.h

// Copyright (c) 2020-now by the Zeek Project. See LICENSE for details.
 
#pragma once
 
#include <csetjmp>
#include <memory>
#include <optional>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>
 
#include <hilti/rt/any.h>
#include <hilti/rt/autogen/config.h>
#include <hilti/rt/configuration.h>
#include <hilti/rt/exception.h>
#include <hilti/rt/types/reference.h>
#include <hilti/rt/util.h>
 
struct Fiber;
 
// Fiber entry point for execution of fiber payload.
extern "C" void __fiber_run_trampoline(void* args);
 
// Fiber entry point for stack switch trampoline.
extern "C" void __fiber_switch_trampoline(void* args);
 
namespace hilti::rt {
 
namespace detail {
class Fiber;
} // namespace detail
 
namespace resumable {
using Handle = detail::Fiber;
} // namespace resumable
 
namespace detail {
 
extern void trackStack();
 
struct FiberContext {
    FiberContext();
    ~FiberContext();
 
    FiberContext(const FiberContext&) = delete;
    FiberContext(FiberContext&&) = default;
 
    FiberContext& operator=(const FiberContext&) = delete;
    FiberContext& operator=(FiberContext&&) = default;
 
    std::unique_ptr<detail::Fiber> main;
 
    std::unique_ptr<detail::Fiber> switch_trampoline;
 
    detail::Fiber* current = nullptr;
 
    std::unique_ptr<::Fiber> shared_stack;
 
    std::vector<std::unique_ptr<Fiber>> cache;
};
 
struct StackBuffer {
    StackBuffer(const ::Fiber* fiber) : _fiber(fiber) {}
 
    ~StackBuffer();
 
    std::pair<char*, char*> activeRegion() const;
 
    std::pair<char*, char*> allocatedRegion() const;
 
    size_t activeSize() const;
 
    size_t allocatedSize() const { return static_cast<size_t>(allocatedRegion().second - allocatedRegion().first); }
 
    size_t liveRemainingSize() const;
 
    void save();
 
    void restore() const;
 
private:
    const ::Fiber* _fiber;
    void* _buffer = nullptr; // allocated memory holding swapped out stack content
    size_t _buffer_size = 0; // amount currently allocated for `_buffer`
};
 
// Render stack region for use in debug output.
inline std::ostream& operator<<(std::ostream& out, const StackBuffer& s) {
    out << fmt("%p-%p:%zu", s.activeRegion().first, s.activeRegion().second, s.activeSize());
    return out;
}
 
class Callback {
public:
    template<typename F>
    Callback(F f)
        : _f(std::move(f)), _invoke([](const hilti::rt::any& f, resumable::Handle* h) -> hilti::rt::any {
              return hilti::rt::any_cast<F>(f)(h);
          }) {}
 
    Callback(const Callback&) = default;
    Callback(Callback&&) = default;
 
    Callback& operator=(const Callback&) = default;
    Callback& operator=(Callback&&) = default;
 
    hilti::rt::any operator()(resumable::Handle* h) const { return _invoke(_f, h); }
 
private:
    hilti::rt::any _f; //< Type-erased storage for the concrete callback.
    hilti::rt::any (*_invoke)(const hilti::rt::any& f, resumable::Handle* h); //< Invoke type-erased callback.
};
 
class Fiber {
public:
    enum class Type : int64_t {
        IndividualStack, 
        SharedStack,     
        Main,             
        SwitchTrampoline, 
    };
 
    Fiber(Type type);
    ~Fiber();
 
    Fiber(const Fiber&) = delete;
    Fiber(Fiber&&) = delete;
    Fiber& operator=(const Fiber&) = delete;
    Fiber& operator=(Fiber&&) = delete;
 
    void init(Callback f) {
        _result = {};
        _exception = nullptr;
        _function = std::move(f);
    }
 
    auto type() { return _type; }
 
    const auto& stackBuffer() const { return _stack_buffer; }
 
    void run();
    void yield();
    void resume();
    void abort();
 
    bool isMain() const { return _type == Type::Main; }
 
    bool isDone() {
        switch ( _state ) {
            case State::Running:
            case State::Yielded: return false;
 
            case State::Aborting:
            case State::Finished:
            case State::Idle:
            case State::Init:
                // All these mean we didn't recently run a function that could have
                // produced a result still pending.
                return true;
        }
        cannot_be_reached(); // For you, GCC.
    }
 
    auto&& result() { return std::move(_result); }
    std::exception_ptr exception() const { return _exception; }
 
    const char* location() const { return _location; }
 
    void setLocation(const char* location = nullptr) { _location = location; }
 
    std::string tag() const;
 
    static std::unique_ptr<Fiber> create();
    static void destroy(std::unique_ptr<Fiber> f);
    static void primeCache();
    static void reset();
 
    struct Statistics {
        uint64_t total;
        uint64_t current;
        uint64_t cached;
        uint64_t max;
        uint64_t max_stack_size;
        uint64_t initialized;
    };
 
    static Statistics statistics();
 
private:
    friend void ::__fiber_run_trampoline(void* argsp);
    friend void ::__fiber_switch_trampoline(void* argsp);
    friend void detail::trackStack();
 
    enum class State { Init, Running, Aborting, Yielded, Idle, Finished };
 
    void _yield(const char* tag);
    void _activate(const char* tag);
 
    static void _startSwitchFiber(const char* tag, detail::Fiber* to);
 
    static void _finishSwitchFiber(const char* tag);
 
    static void _executeSwitch(const char* tag, detail::Fiber* from, detail::Fiber* to);
 
    Type _type;
    State _state{State::Init};
    std::optional<Callback> _function;
    std::optional<hilti::rt::any> _result;
    std::exception_ptr _exception;
 
    std::unique_ptr<::Fiber> _fiber;
 
    Fiber* _caller = nullptr;
 
    StackBuffer _stack_buffer;
 
    const char* _location = nullptr;
 
#ifdef _WIN32
    struct {
        void* stack_base = nullptr;
        void* stack_limit = nullptr;
        void* deallocation_stack = nullptr;
    } _teb;
#endif
 
#ifdef HILTI_HAVE_ASAN
    struct {
        const void* stack = nullptr;
        size_t stack_size = 0;
        void* fake_stack = nullptr;
    } _asan;
#endif
 
    // TODO: Usage of these isn't thread-safe. Should become "atomic" and
    // move into global state.
    HILTI_JIT_IMPORT_OR_INLINE static uint64_t _total_fibers;
    HILTI_JIT_IMPORT_OR_INLINE static uint64_t _current_fibers;
    HILTI_JIT_IMPORT_OR_INLINE static uint64_t _cached_fibers;
    HILTI_JIT_IMPORT_OR_INLINE static uint64_t _max_fibers;
    HILTI_JIT_IMPORT_OR_INLINE static uint64_t _max_stack_size;
    HILTI_JIT_IMPORT_OR_INLINE static uint64_t _initialized; // number of trampolines run
};
 
std::ostream& operator<<(std::ostream& out, const Fiber& fiber);
 
extern void yield();
 
} // namespace detail
 
class Resumable {
public:
    template<typename Function>
    Resumable(Function f)
        requires(std::is_invocable_v<Function, resumable::Handle*>)
        : _fiber(detail::Fiber::create()) {
        _fiber->init(std::move(f));
    }
 
    Resumable() = default;
    Resumable(const Resumable& r) = delete;
    Resumable(Resumable&& r) noexcept = default;
    Resumable& operator=(const Resumable& other) = delete;
    Resumable& operator=(Resumable&& other) noexcept = default;
 
    ~Resumable() {
        if ( _fiber )
            try {
                detail::Fiber::destroy(std::move(_fiber));
            } catch ( ... ) {
                cannot_be_reached();
            }
    }
 
    void run();
 
    void resume();
 
    void abort();
 
    resumable::Handle* handle() { return _fiber.get(); }
 
    bool hasResult() const { return _done && _result.has_value(); }
 
    template<typename Result>
    const Result& get() const {
        assert(static_cast<bool>(_result));
 
        if constexpr ( std::is_same_v<Result, void> )
            return {};
        else {
            try {
                return hilti::rt::any_cast<const Result&>(*_result);
            } catch ( const hilti::rt::bad_any_cast& ) {
                throw InvalidArgument("mismatch in result type");
            }
        }
    }
 
    explicit operator bool() const { return _done; }
 
private:
    void yielded();
 
    void checkFiber(const char* location) const {
        if ( ! _fiber )
            throw std::logic_error(std::string("fiber not set in ") + location);
    }
 
    std::unique_ptr<detail::Fiber> _fiber;
    bool _done = false;
    std::optional<hilti::rt::any> _result;
};
 
namespace resumable::detail {
 
template<typename T>
auto copyArg(T t) {
    // In general, we can't move references to the heap.
    static_assert(! std::is_reference_v<T>, "copyArg() does not accept references other than ValueReference<T>.");
    return t;
}
 
// Special case: We don't want to (nor need to) deep-copy value references.
// Their payload already resides on the heap, so reuse that.
template<typename T>
ValueReference<T> copyArg(const ValueReference<T>& t) {
    return ValueReference<T>(t.asSharedPtr());
}
 
// Special case: We don't want to (nor need to) deep-copy value references.
// Their payload already resides on the heap, so reuse that.
template<typename T>
ValueReference<T> copyArg(ValueReference<T>& t) {
    return ValueReference<T>(t.asSharedPtr());
}
 
} // namespace resumable::detail
 
namespace fiber {
 
template<typename Function>
auto execute(Function f) {
    Resumable r(std::move(f));
    r.run();
    return r;
}
 
} // namespace fiber
} // namespace hilti::rt