execution_stage.hh

/*
 * This file is open source software, licensed to you under the terms
 * of the Apache License, Version 2.0 (the "License").  See the NOTICE file
 * distributed with this work for additional information regarding copyright
 * ownership.  You may not use this file except in compliance with the License.
 *
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
/*
 * Copyright (C) 2017 ScyllaDB Ltd.
 */
 
#pragma once
 
#include <seastar/core/future.hh>
#include <seastar/core/chunked_fifo.hh>
#include <seastar/core/function_traits.hh>
#include <seastar/core/sstring.hh>
#include <seastar/core/metrics.hh>
#include <seastar/core/scheduling.hh>
#include <seastar/util/reference_wrapper.hh>
#include <seastar/util/noncopyable_function.hh>
#include <seastar/util/tuple_utils.hh>
#include <seastar/util/std-compat.hh>
#include <seastar/util/modules.hh>
#ifndef SEASTAR_MODULE
#include <fmt/format.h>
#include <vector>
#include <boost/range/irange.hpp>
#include <boost/range/adaptor/transformed.hpp>
#include <boost/container/static_vector.hpp>
#endif
 
namespace seastar {
 
 
 
namespace internal {
 
// Execution wraps lreferences in reference_wrapper so that the caller is forced
// to use seastar::ref(). Then when the function is actually called the
// reference is unwrapped. However, we need to distinguish between functions
// which argument is lvalue reference and functions that take
// reference_wrapper<> as an argument and not unwrap the latter. To solve this
// issue reference_wrapper_for_es type is used for wrappings done automatically
// by execution stage.
template<typename T>
struct reference_wrapper_for_es : reference_wrapper<T> {
    reference_wrapper_for_es(reference_wrapper <T> rw) noexcept
        : reference_wrapper<T>(std::move(rw)) {}
};
 
template<typename T>
struct wrap_for_es {
    using type = T;
};
 
template<typename T>
struct wrap_for_es<T&> {
    using type = reference_wrapper_for_es<T>;
};
 
template<typename T>
struct wrap_for_es<T&&> {
    using type = T;
};
 
template<typename T>
decltype(auto) unwrap_for_es(T&& object) {
    return std::forward<T>(object);
}
 
template<typename T>
std::reference_wrapper<T> unwrap_for_es(reference_wrapper_for_es<T> ref) {
    return std::reference_wrapper<T>(ref.get());
}
 
}
 
SEASTAR_MODULE_EXPORT
class execution_stage {
public:
    struct stats {
        uint64_t tasks_scheduled = 0;
        uint64_t tasks_preempted = 0;
        uint64_t function_calls_enqueued = 0;
        uint64_t function_calls_executed = 0;
    };
protected:
    bool _empty = true;
    bool _flush_scheduled = false;
    scheduling_group _sg;
    stats _stats;
    sstring _name;
    metrics::metric_group _metric_group;
protected:
    virtual void do_flush() noexcept = 0;
public:
    explicit execution_stage(const sstring& name, scheduling_group sg = {});
    virtual ~execution_stage();
 
    execution_stage(const execution_stage&) = delete;
 
    execution_stage(execution_stage&&);
 
    const sstring& name() const noexcept { return _name; }
 
    const stats& get_stats() const noexcept { return _stats; }
 
    bool flush() noexcept;
 
    bool poll() const noexcept {
        return !_empty;
    }
};
 
namespace internal {
 
class execution_stage_manager {
    std::vector<execution_stage*> _execution_stages;
    std::unordered_map<sstring, execution_stage*> _stages_by_name;
private:
    execution_stage_manager() = default;
    execution_stage_manager(const execution_stage_manager&) = delete;
    execution_stage_manager(execution_stage_manager&&) = delete;
public:
    void register_execution_stage(execution_stage& stage);
    void unregister_execution_stage(execution_stage& stage) noexcept;
    void update_execution_stage_registration(execution_stage& old_es, execution_stage& new_es) noexcept;
    execution_stage* get_stage(const sstring& name);
    bool flush() noexcept;
    bool poll() const noexcept;
public:
    static execution_stage_manager& get() noexcept;
};
 
}
 
template<typename ReturnType, typename... Args>
requires std::is_nothrow_move_constructible_v<std::tuple<Args...>>
class concrete_execution_stage final : public execution_stage {
    using args_tuple = std::tuple<Args...>;
    static_assert(std::is_nothrow_move_constructible_v<args_tuple>,
                  "Function arguments need to be nothrow move constructible");
 
    static constexpr size_t flush_threshold = 128;
    static constexpr size_t max_queue_length = 1024;
 
    using return_type = futurize_t<ReturnType>;
    using promise_type = typename return_type::promise_type;
    using input_type = typename tuple_map_types<internal::wrap_for_es, args_tuple>::type;
 
    struct work_item {
        input_type _in;
        promise_type _ready;
 
        work_item(typename internal::wrap_for_es<Args>::type... args) : _in(std::move(args)...) { }
 
        work_item(work_item&& other) = delete;
        work_item(const work_item&) = delete;
        work_item(work_item&) = delete;
    };
    chunked_fifo<work_item, flush_threshold> _queue;
 
    noncopyable_function<ReturnType (Args...)> _function;
private:
    auto unwrap(input_type&& in) {
        return tuple_map(std::move(in), [] (auto&& obj) {
            return internal::unwrap_for_es(std::forward<decltype(obj)>(obj));
        });
    }
 
    virtual void do_flush() noexcept override {
        while (!_queue.empty()) {
            auto& wi = _queue.front();
            auto wi_in = std::move(wi._in);
            auto wi_ready = std::move(wi._ready);
            _queue.pop_front();
            futurize<ReturnType>::apply(_function, unwrap(std::move(wi_in))).forward_to(std::move(wi_ready));
            _stats.function_calls_executed++;
 
            if (internal::scheduler_need_preempt()) {
                _stats.tasks_preempted++;
                break;
            }
        }
        _empty = _queue.empty();
    }
public:
    explicit concrete_execution_stage(const sstring& name, scheduling_group sg, noncopyable_function<ReturnType (Args...)> f)
        : execution_stage(name, sg)
        , _function(std::move(f))
    {
        _queue.reserve(flush_threshold);
    }
    explicit concrete_execution_stage(const sstring& name, noncopyable_function<ReturnType (Args...)> f)
        : concrete_execution_stage(name, scheduling_group(), std::move(f)) {
    }
 
    return_type operator()(typename internal::wrap_for_es<Args>::type... args) {
        if (_queue.size() >= max_queue_length) {
            do_flush();
        }
        _queue.emplace_back(std::move(args)...);
        _empty = false;
        _stats.function_calls_enqueued++;
        auto f = _queue.back()._ready.get_future();
        flush();
        return f;
    }
};
 
class inheriting_execution_stage {
public:
    struct per_scheduling_group_stats {
        scheduling_group sg;
        execution_stage::stats stats;
    };
    using stats = boost::container::static_vector<per_scheduling_group_stats, max_scheduling_groups()>;
};
 
template<typename ReturnType, typename... Args>
requires std::is_nothrow_move_constructible_v<std::tuple<Args...>>
class inheriting_concrete_execution_stage final : public inheriting_execution_stage {
    using return_type = futurize_t<ReturnType>;
    using args_tuple = std::tuple<Args...>;
    using per_group_stage_type = concrete_execution_stage<ReturnType, Args...>;
 
    static_assert(std::is_nothrow_move_constructible_v<args_tuple>,
                  "Function arguments need to be nothrow move constructible");
 
    sstring _name;
    noncopyable_function<ReturnType (Args...)> _function;
    std::vector<std::optional<per_group_stage_type>> _stage_for_group{max_scheduling_groups()};
private:
    per_group_stage_type make_stage_for_group(scheduling_group sg) {
        // We can't use std::ref(function), because reference_wrapper decays to noncopyable_function& and
        // that selects the noncopyable_function copy constructor. Use a lambda instead.
        auto wrapped_function = [&_function = _function] (Args... args) {
            return _function(std::forward<Args>(args)...);
        };
        auto name = fmt::format("{}.{}", _name, sg.name());
        return per_group_stage_type(name, sg, wrapped_function);
    }
public:
    inheriting_concrete_execution_stage(const sstring& name, noncopyable_function<ReturnType (Args...)> f)
        : _name(std::move(name)),_function(std::move(f)) {
    }
 
    return_type operator()(typename internal::wrap_for_es<Args>::type... args) {
        auto sg = current_scheduling_group();
        auto sg_id = internal::scheduling_group_index(sg);
        auto& slot = _stage_for_group[sg_id];
        if (!slot) {
            slot.emplace(make_stage_for_group(sg));
        }
        return (*slot)(std::move(args)...);
    }
 
    inheriting_execution_stage::stats get_stats() const noexcept {
        inheriting_execution_stage::stats summary;
        for (unsigned sg_id = 0; sg_id != _stage_for_group.size(); ++sg_id) {
            auto sg = internal::scheduling_group_from_index(sg_id);
            if (_stage_for_group[sg_id]) {
                summary.push_back({sg, _stage_for_group[sg_id]->get_stats()});
            }
        }
        return summary;
    }
};
 
 
namespace internal {
 
template <typename Ret, typename ArgsTuple>
struct concrete_execution_stage_helper;
 
template <typename Ret, typename... Args>
struct concrete_execution_stage_helper<Ret, std::tuple<Args...>> {
    using type = concrete_execution_stage<Ret, Args...>;
};
 
}
 
SEASTAR_MODULE_EXPORT
template<typename Function>
auto make_execution_stage(const sstring& name, scheduling_group sg, Function&& fn) {
    using traits = function_traits<Function>;
    using ret_type = typename traits::return_type;
    using args_as_tuple = typename traits::args_as_tuple;
    using concrete_execution_stage = typename internal::concrete_execution_stage_helper<ret_type, args_as_tuple>::type;
    return concrete_execution_stage(name, sg, std::forward<Function>(fn));
}
 
SEASTAR_MODULE_EXPORT
template<typename Function>
auto make_execution_stage(const sstring& name, Function&& fn) {
    return make_execution_stage(name, scheduling_group(), std::forward<Function>(fn));
}
 
SEASTAR_MODULE_EXPORT
template<typename Ret, typename Object, typename... Args>
concrete_execution_stage<Ret, Object*, Args...>
make_execution_stage(const sstring& name, scheduling_group sg, Ret (Object::*fn)(Args...)) {
    return concrete_execution_stage<Ret, Object*, Args...>(name, sg, std::mem_fn(fn));
}
 
template<typename Ret, typename Object, typename... Args>
concrete_execution_stage<Ret, const Object*, Args...>
make_execution_stage(const sstring& name, scheduling_group sg, Ret (Object::*fn)(Args...) const) {
    return concrete_execution_stage<Ret, const Object*, Args...>(name, sg, std::mem_fn(fn));
}
 
template<typename Ret, typename Object, typename... Args>
concrete_execution_stage<Ret, Object*, Args...>
make_execution_stage(const sstring& name, Ret (Object::*fn)(Args...)) {
    return make_execution_stage(name, scheduling_group(), fn);
}
 
template<typename Ret, typename Object, typename... Args>
concrete_execution_stage<Ret, const Object*, Args...>
make_execution_stage(const sstring& name, Ret (Object::*fn)(Args...) const) {
    return make_execution_stage(name, scheduling_group(), fn);
}
 
 
}