aboutsummaryrefslogtreecommitdiffstats
path: root/contrib/clickhouse/src/Common/Fiber.h
blob: f48ace149f492eb34f62d1fce3c7e7cc4710e71f (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
#pragma once
/// defines.h should be included before fiber.hpp
/// BOOST_USE_ASAN, BOOST_USE_TSAN and BOOST_USE_UCONTEXT should be correctly defined for sanitizers.
#include <base/defines.h>
#include <boost/context/fiber.hpp>
#include <map>

/// Class wrapper for boost::context::fiber.
/// It tracks current executing fiber for thread and
/// supports storing fiber-specific data
/// that will be destroyed on fiber destructor.
class Fiber
{
private:
    using Impl = boost::context::fiber;
    using FiberPtr = Fiber *;
    template <typename T> friend class FiberLocal;

public:
    template< typename StackAlloc, typename Fn>
    Fiber(StackAlloc && salloc, Fn && fn) : impl(std::allocator_arg_t(), std::forward<StackAlloc>(salloc), RoutineImpl(std::forward<Fn>(fn)))
    {
    }

    Fiber() = default;

    Fiber(Fiber && other) = default;
    Fiber & operator=(Fiber && other) = default;

    Fiber(const Fiber &) = delete;
    Fiber & operator =(const Fiber &) = delete;

    explicit operator bool() const
    {
        return impl.operator bool();
    }

    void resume()
    {
        /// Update information about current executing fiber.
        FiberPtr & current_fiber = getCurrentFiber();
        FiberPtr parent_fiber = current_fiber;
        current_fiber = this;
        impl = std::move(impl).resume();
        /// Restore parent fiber.
        current_fiber = parent_fiber;
    }

private:
    template <typename Fn>
    struct RoutineImpl
    {
        struct SuspendCallback
        {
            Impl & impl;

            void operator()()
            {
                impl = std::move(impl).resume();
            }
        };

        explicit RoutineImpl(Fn && fn_) : fn(std::move(fn_))
        {
        }

        Impl operator()(Impl && sink)
        {
            SuspendCallback suspend_callback{sink};
            fn(suspend_callback);
            return std::move(sink);
        }

        Fn fn;
    };

    static FiberPtr & getCurrentFiber()
    {
        thread_local static FiberPtr current_fiber;
        return current_fiber;
    }

    /// Special wrapper to store data in uniquer_ptr.
    struct DataWrapper
    {
        virtual ~DataWrapper() = default;
    };

    using DataPtr = std::unique_ptr<DataWrapper>;

    /// Get reference to fiber-specific data by key
    /// (the pointer to the structure that uses this data).
    DataPtr & getLocalData(void * key)
    {
        return local_data[key];
    }

    Impl && release()
    {
        return std::move(impl);
    }

    Impl impl;
    std::map<void *, DataPtr> local_data;
};

/// Implementation for fiber local variable.
/// If we are in fiber, it returns fiber local data,
/// otherwise it returns it's single field.
/// Fiber local data is destroyed in Fiber destructor.
/// Implementation is similar to boost::fiber::fiber_specific_ptr
/// (we cannot use it because we don't use boost::fiber API.
template <typename T>
class FiberLocal
{
public:
    T & operator*()
    {
        return get();
    }

    T * operator->()
    {
        return &get();
    }

private:
    struct DataWrapperImpl : public Fiber::DataWrapper
    {
        T impl;
    };

    T & get()
    {
        Fiber * current_fiber = Fiber::getCurrentFiber();
        if (!current_fiber)
            return main_instance;

        Fiber::DataPtr & ptr = current_fiber->getLocalData(this);
        /// Initialize instance on first request.
        if (!ptr)
            ptr = std::make_unique<DataWrapperImpl>();

        return dynamic_cast<DataWrapperImpl *>(ptr.get())->impl;
    }

    T main_instance;
};