sharded.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) 2015 Cloudius Systems, Ltd.
 */
 
#pragma once
 
#include <seastar/core/smp.hh>
#include <seastar/core/loop.hh>
#include <seastar/core/map_reduce.hh>
#include <seastar/core/internal/run_in_background.hh>
#include <seastar/util/is_smart_ptr.hh>
#include <seastar/util/tuple_utils.hh>
#include <seastar/core/do_with.hh>
#include <seastar/util/log.hh>
#include <seastar/util/modules.hh>
 
#ifndef SEASTAR_MODULE
#include <boost/iterator/counting_iterator.hpp>
#include <concepts>
#include <functional>
#endif
 
 
namespace seastar {
 
SEASTAR_MODULE_EXPORT_BEGIN
 
template <typename Func, typename... Param>
class sharded_parameter;
 
template <typename Service>
class sharded;
 
SEASTAR_MODULE_EXPORT_END
 
namespace internal {
 
template <typename Func, typename... Param>
auto unwrap_sharded_arg(sharded_parameter<Func, Param...> sp);
 
using on_each_shard_func = std::function<future<> (unsigned shard)>;
 
future<> sharded_parallel_for_each(unsigned nr_shards, on_each_shard_func on_each_shard) noexcept(std::is_nothrow_move_constructible_v<on_each_shard_func>);
 
template <typename Service>
class either_sharded_or_local {
    sharded<Service>& _sharded;
public:
    either_sharded_or_local(sharded<Service>& s) : _sharded(s) {}
    operator sharded<Service>& ();
    operator Service& ();
};
 
template <typename T>
struct sharded_unwrap {
    using evaluated_type = T;
    using type = T;
};
 
template <typename T>
struct sharded_unwrap<std::reference_wrapper<sharded<T>>> {
    using evaluated_type = T&;
    using type = either_sharded_or_local<T>;
};
 
template <typename Func, typename... Param>
struct sharded_unwrap<sharded_parameter<Func, Param...>> {
    using type = std::invoke_result_t<Func, Param...>;
};
 
template <typename T>
using sharded_unwrap_evaluated_t = typename sharded_unwrap<T>::evaluated_type;
 
template <typename T>
using sharded_unwrap_t = typename sharded_unwrap<T>::type;
 
} // internal
 
 
 
SEASTAR_MODULE_EXPORT_BEGIN
 
template <typename T>
class sharded;
 
template<typename T>
class async_sharded_service : public enable_shared_from_this<T> {
    promise<> _freed;
protected:
    async_sharded_service() noexcept = default;
    virtual ~async_sharded_service() {
        _freed.set_value();
    }
    future<> freed() noexcept {
        return _freed.get_future();
    }
    template <typename Service> friend class sharded;
};
 
 
template <typename Service>
class peering_sharded_service {
    sharded<Service>* _container = nullptr;
private:
    template <typename T> friend class sharded;
    void set_container(sharded<Service>* container) noexcept { _container = container; }
public:
    peering_sharded_service() noexcept = default;
    peering_sharded_service(peering_sharded_service<Service>&&) noexcept = default;
    peering_sharded_service(const peering_sharded_service<Service>&) = delete;
    peering_sharded_service& operator=(const peering_sharded_service<Service>&) = delete;
    sharded<Service>& container() noexcept { return *_container; }
    const sharded<Service>& container() const noexcept { return *_container; }
};
 
 
class no_sharded_instance_exception : public std::exception {
    sstring _msg;
public:
    no_sharded_instance_exception() : _msg("sharded instance does not exist") {}
    explicit no_sharded_instance_exception(sstring type_info)
        : _msg("sharded instance does not exist: " + type_info) {}
    virtual const char* what() const noexcept override {
        return _msg.c_str();
    }
};
 
template <typename Service>
class sharded {
    struct entry {
        shared_ptr<Service> service;
 
        future<> track_deletion() noexcept {
            // do not wait for instance to be deleted if it is not going to notify us
            if constexpr (std::is_base_of_v<async_sharded_service<Service>, Service>) {
                if (service) {
                    return service->freed();
                }
            }
            return make_ready_future<>();
        }
    };
    std::vector<entry> _instances;
private:
    using invoke_on_all_func_type = std::function<future<> (Service&)>;
private:
    template <typename U, bool async>
    friend struct shared_ptr_make_helper;
 
    template <typename T>
    std::enable_if_t<std::is_base_of_v<peering_sharded_service<T>, T>>
    set_container(T& service) noexcept {
        service.set_container(this);
    }
 
    template <typename T>
    std::enable_if_t<!std::is_base_of_v<peering_sharded_service<T>, T>>
    set_container(T&) noexcept {
    }
 
    future<>
    sharded_parallel_for_each(internal::on_each_shard_func func) noexcept(std::is_nothrow_move_constructible_v<internal::on_each_shard_func>) {
        return internal::sharded_parallel_for_each(_instances.size(), std::move(func));
    }
public:
    sharded() noexcept {}
    sharded(const sharded& other) = delete;
    sharded& operator=(const sharded& other) = delete;
    sharded(sharded&& other) = delete;
    sharded& operator=(sharded&& other) = delete;
    ~sharded();
 
    template <typename... Args>
    future<> start(Args&&... args) noexcept;
 
    template <typename... Args>
    future<> start_single(Args&&... args) noexcept;
 
    future<> stop() noexcept;
 
    future<> invoke_on_all(smp_submit_to_options options, std::function<future<> (Service&)> func) noexcept;
 
    future<> invoke_on_all(std::function<future<> (Service&)> func) noexcept {
      try {
        return invoke_on_all(smp_submit_to_options{}, std::move(func));
      } catch (...) {
        return current_exception_as_future();
      }
    }
 
    template <typename Func, typename... Args>
    requires std::invocable<Func, Service&, internal::sharded_unwrap_t<Args>...>
    future<> invoke_on_all(smp_submit_to_options options, Func func, Args... args) noexcept;
 
    template <typename Func, typename... Args>
    requires std::invocable<Func, Service&, internal::sharded_unwrap_t<Args>...>
    future<> invoke_on_all(Func func, Args... args) noexcept {
      try {
        return invoke_on_all(smp_submit_to_options{}, std::move(func), std::move(args)...);
      } catch (...) {
        return current_exception_as_future();
      }
    }
 
    template <typename Func, typename... Args>
    requires std::invocable<Func, Service&, Args...>
    future<> invoke_on_others(smp_submit_to_options options, Func func, Args... args) noexcept;
 
    template <typename Func, typename... Args>
    requires std::invocable<Func, Service&, Args...>
    future<> invoke_on_others(Func func, Args... args) noexcept {
      try {
        return invoke_on_others(smp_submit_to_options{}, std::move(func), std::move(args)...);
      } catch (...) {
        return current_exception_as_future();
      }
    }
 
    template <typename Reducer, typename Func, typename... Args>
    inline
    auto map_reduce(Reducer&& r, Func&& func, Args&&... args) -> typename reducer_traits<Reducer>::future_type
    {
        return ::seastar::map_reduce(boost::make_counting_iterator<unsigned>(0),
                            boost::make_counting_iterator<unsigned>(_instances.size()),
            [this, func = std::forward<Func>(func), args = std::make_tuple(std::forward<Args>(args)...)] (unsigned c) mutable {
                return smp::submit_to(c, [this, &func, args] () mutable {
                    return std::apply([this, &func] (Args&&... args) mutable {
                        auto inst = get_local_service();
                        return std::invoke(func, *inst, std::forward<Args>(args)...);
                    }, std::move(args));
                });
            }, std::forward<Reducer>(r));
    }
 
    template <typename Reducer, typename Func, typename... Args>
    inline
    auto map_reduce(Reducer&& r, Func&& func, Args&&... args) const -> typename reducer_traits<Reducer>::future_type
    {
        return ::seastar::map_reduce(boost::make_counting_iterator<unsigned>(0),
                            boost::make_counting_iterator<unsigned>(_instances.size()),
            [this, func = std::forward<Func>(func), args = std::make_tuple(std::forward<Args>(args)...)] (unsigned c) {
                return smp::submit_to(c, [this, &func, args] () {
                    return std::apply([this, &func] (Args&&... args) {
                        auto inst = get_local_service();
                        return std::invoke(func, *inst, std::forward<Args>(args)...);
                    }, std::move(args));
                });
            }, std::forward<Reducer>(r));
    }
 
    template <typename Mapper, typename Initial, typename Reduce>
    inline
    future<Initial>
    map_reduce0(Mapper map, Initial initial, Reduce reduce) {
        auto wrapped_map = [this, map] (unsigned c) {
            return smp::submit_to(c, [this, map] {
                auto inst = get_local_service();
                return std::invoke(map, *inst);
            });
        };
        return ::seastar::map_reduce(smp::all_cpus().begin(), smp::all_cpus().end(),
                            std::move(wrapped_map),
                            std::move(initial),
                            std::move(reduce));
    }
 
    template <typename Mapper, typename Initial, typename Reduce>
    inline
    future<Initial>
    map_reduce0(Mapper map, Initial initial, Reduce reduce) const {
        auto wrapped_map = [this, map] (unsigned c) {
            return smp::submit_to(c, [this, map] {
                auto inst = get_local_service();
                return std::invoke(map, *inst);
            });
        };
        return ::seastar::map_reduce(smp::all_cpus().begin(), smp::all_cpus().end(),
                            std::move(wrapped_map),
                            std::move(initial),
                            std::move(reduce));
    }
 
    template <typename Mapper, typename Future = futurize_t<std::invoke_result_t<Mapper,Service&>>, typename return_type = decltype(internal::untuple(std::declval<typename Future::tuple_type>()))>
    inline future<std::vector<return_type>> map(Mapper mapper) {
        return do_with(std::vector<return_type>(), std::move(mapper),
                [this] (std::vector<return_type>& vec, Mapper& mapper) mutable {
            vec.resize(_instances.size());
            return parallel_for_each(boost::irange<unsigned>(0, _instances.size()), [this, &vec, &mapper] (unsigned c) {
                return smp::submit_to(c, [this, &mapper] {
                    auto inst = get_local_service();
                    return mapper(*inst);
                }).then([&vec, c] (auto&& res) {
                    vec[c] = std::move(res);
                });
            }).then([&vec] {
                return make_ready_future<std::vector<return_type>>(std::move(vec));
            });
        });
    }
 
    template <typename Func, typename... Args, typename Ret = futurize_t<std::invoke_result_t<Func, Service&, Args...>>>
    requires std::invocable<Func, Service&, Args&&...>
    Ret
    invoke_on(unsigned id, smp_submit_to_options options, Func&& func, Args&&... args) {
        return smp::submit_to(id, options, [this, func = std::forward<Func>(func), args = std::tuple(std::move(args)...)] () mutable {
            auto inst = get_local_service();
            return std::apply(std::forward<Func>(func), std::tuple_cat(std::forward_as_tuple(*inst), std::move(args)));
        });
    }
 
    template <typename Func, typename... Args, typename Ret = futurize_t<std::invoke_result_t<Func, Service&, Args&&...>>>
    requires std::invocable<Func, Service&, Args&&...>
    Ret
    invoke_on(unsigned id, Func&& func, Args&&... args) {
        return invoke_on(id, smp_submit_to_options(), std::forward<Func>(func), std::forward<Args>(args)...);
    }
 
    const Service& local() const noexcept;
 
    Service& local() noexcept;
 
    shared_ptr<Service> local_shared() noexcept;
 
    bool local_is_initialized() const noexcept;
 
private:
    template <typename... Args>
    shared_ptr<Service> create_local_service(Args&&... args) {
        auto s = ::seastar::make_shared<Service>(std::forward<Args>(args)...);
        set_container(*s);
        return s;
    }
 
    shared_ptr<Service> get_local_service() {
        auto inst = _instances[this_shard_id()].service;
        if (!inst) {
            throw no_sharded_instance_exception(pretty_type_name(typeid(Service)));
        }
        return inst;
    }
 
    shared_ptr<const Service> get_local_service() const {
        auto inst = _instances[this_shard_id()].service;
        if (!inst) {
            throw no_sharded_instance_exception(pretty_type_name(typeid(Service)));
        }
        return inst;
    }
};
 
 
template <typename Func, typename... Params>
class sharded_parameter {
    Func _func;
    std::tuple<Params...> _params;
public:
    explicit sharded_parameter(Func func, Params... params)
            requires std::invocable<Func, internal::sharded_unwrap_evaluated_t<Params>...>
            : _func(std::move(func)), _params(std::make_tuple(std::move(params)...)) {
    }
private:
    auto evaluate() const;
 
    template <typename Func_, typename... Param_>
    friend auto internal::unwrap_sharded_arg(sharded_parameter<Func_, Param_...> sp);
};
 
SEASTAR_MODULE_EXPORT_END
 
template <typename Service>
sharded<Service>::~sharded() {
        assert(_instances.empty());
}
 
namespace internal {
 
template <typename T>
inline
T&&
unwrap_sharded_arg(T&& arg) {
    return std::forward<T>(arg);
}
 
template <typename Service>
either_sharded_or_local<Service>
unwrap_sharded_arg(std::reference_wrapper<sharded<Service>> arg) {
    return either_sharded_or_local<Service>(arg);
}
 
template <typename Func, typename... Param>
auto
unwrap_sharded_arg(sharded_parameter<Func, Param...> sp) {
    return sp.evaluate();
}
 
template <typename Service>
either_sharded_or_local<Service>::operator sharded<Service>& () { return _sharded; }
 
template <typename Service>
either_sharded_or_local<Service>::operator Service& () { return _sharded.local(); }
 
}
 
template <typename Func, typename... Param>
auto
sharded_parameter<Func, Param...>::evaluate() const {
    auto unwrap_params_and_invoke = [this] (const auto&... params) {
        return std::invoke(_func, internal::unwrap_sharded_arg(params)...);
    };
    return std::apply(unwrap_params_and_invoke, _params);
}
 
template <typename Service>
template <typename... Args>
future<>
sharded<Service>::start(Args&&... args) noexcept {
  try {
    _instances.resize(smp::count);
    return sharded_parallel_for_each(
        [this, args = std::make_tuple(std::forward<Args>(args)...)] (unsigned c) mutable {
            return smp::submit_to(c, [this, args] () mutable {
                _instances[this_shard_id()].service = std::apply([this] (Args... args) {
                    return create_local_service(internal::unwrap_sharded_arg(std::forward<Args>(args))...);
                }, args);
            });
    }).then_wrapped([this] (future<> f) {
        try {
            f.get();
            return make_ready_future<>();
        } catch (...) {
            return this->stop().then([e = std::current_exception()] () mutable {
                std::rethrow_exception(e);
            });
        }
    });
  } catch (...) {
    return current_exception_as_future();
  }
}
 
template <typename Service>
template <typename... Args>
future<>
sharded<Service>::start_single(Args&&... args) noexcept {
  try {
    assert(_instances.empty());
    _instances.resize(1);
    return smp::submit_to(0, [this, args = std::make_tuple(std::forward<Args>(args)...)] () mutable {
        _instances[0].service = std::apply([this] (Args... args) {
            return create_local_service(internal::unwrap_sharded_arg(std::forward<Args>(args))...);
        }, args);
    }).then_wrapped([this] (future<> f) {
        try {
            f.get();
            return make_ready_future<>();
        } catch (...) {
            return this->stop().then([e = std::current_exception()] () mutable {
                std::rethrow_exception(e);
            });
        }
    });
  } catch (...) {
    return current_exception_as_future();
  }
}
 
namespace internal {
 
// Helper check if Service::stop exists
 
struct sharded_has_stop {
    // If a member names "stop" exists, try to call it, even if it doesn't
    // have the correct signature. This is so that we don't ignore a function
    // named stop() just because the signature is incorrect, and instead
    // force the user to resolve the ambiguity.
    template <typename Service>
    constexpr static auto check(int) -> std::enable_if_t<(sizeof(&Service::stop) >= 0), bool> {
        return true;
    }
 
    // Fallback in case Service::stop doesn't exist.
    template<typename>
    static constexpr auto check(...) -> bool {
        return false;
    }
};
 
template <bool stop_exists>
struct sharded_call_stop {
    template <typename Service>
    static future<> call(Service& instance);
};
 
template <>
template <typename Service>
inline
future<> sharded_call_stop<true>::call(Service& instance) {
    return instance.stop();
}
 
template <>
template <typename Service>
inline
future<> sharded_call_stop<false>::call(Service&) {
    return make_ready_future<>();
}
 
template <typename Service>
inline
future<>
stop_sharded_instance(Service& instance) {
    constexpr bool has_stop = internal::sharded_has_stop::check<Service>(0);
    return internal::sharded_call_stop<has_stop>::call(instance);
}
 
}
 
template <typename Service>
future<>
sharded<Service>::stop() noexcept {
  try {
    return sharded_parallel_for_each([this] (unsigned c) mutable {
        return smp::submit_to(c, [this] () mutable {
            auto inst = _instances[this_shard_id()].service;
            if (!inst) {
                return make_ready_future<>();
            }
            return internal::stop_sharded_instance(*inst);
        });
    }).then_wrapped([this] (future<> fut) {
        return sharded_parallel_for_each([this] (unsigned c) {
            return smp::submit_to(c, [this] {
                auto fut = _instances[this_shard_id()].track_deletion();
                _instances[this_shard_id()].service = nullptr;
                return fut;
            });
        }).finally([this, fut = std::move(fut)] () mutable {
            _instances.clear();
            _instances = std::vector<sharded<Service>::entry>();
            return std::move(fut);
        });
    });
  } catch (...) {
    return current_exception_as_future();
  }
}
 
template <typename Service>
future<>
sharded<Service>::invoke_on_all(smp_submit_to_options options, std::function<future<> (Service&)> func) noexcept {
  try {
    return sharded_parallel_for_each([this, options, func = std::move(func)] (unsigned c) {
        return smp::submit_to(c, options, [this, func] {
            return func(*get_local_service());
        });
    });
  } catch (...) {
    return current_exception_as_future();
  }
}
 
template <typename Service>
template <typename Func, typename... Args>
requires std::invocable<Func, Service&, internal::sharded_unwrap_t<Args>...>
inline
future<>
sharded<Service>::invoke_on_all(smp_submit_to_options options, Func func, Args... args) noexcept {
    static_assert(std::is_same_v<futurize_t<std::invoke_result_t<Func, Service&, internal::sharded_unwrap_t<Args>...>>, future<>>,
                  "invoke_on_all()'s func must return void or future<>");
  try {
    return invoke_on_all(options, invoke_on_all_func_type([func = std::move(func), args = std::tuple(std::move(args)...)] (Service& service) mutable {
        return std::apply([&service, &func] (Args&&... args) mutable {
            return futurize_apply(func, std::tuple_cat(std::forward_as_tuple(service), std::tuple(internal::unwrap_sharded_arg(std::forward<Args>(args))...)));
        }, std::move(args));
    }));
  } catch (...) {
    return current_exception_as_future();
  }
}
 
template <typename Service>
template <typename Func, typename... Args>
requires std::invocable<Func, Service&, Args...>
inline
future<>
sharded<Service>::invoke_on_others(smp_submit_to_options options, Func func, Args... args) noexcept {
    static_assert(std::is_same_v<futurize_t<std::invoke_result_t<Func, Service&, Args...>>, future<>>,
                  "invoke_on_others()'s func must return void or future<>");
  try {
    return invoke_on_all(options, [orig = this_shard_id(), func = std::move(func), args = std::tuple(std::move(args)...)] (Service& s) mutable -> future<> {
        return this_shard_id() == orig ? make_ready_future<>() : futurize_apply(func, std::tuple_cat(std::forward_as_tuple(s), args));;
    });
  } catch (...) {
    return current_exception_as_future();
  }
}
 
template <typename Service>
const Service& sharded<Service>::local() const noexcept {
    assert(local_is_initialized());
    return *_instances[this_shard_id()].service;
}
 
template <typename Service>
Service& sharded<Service>::local() noexcept {
    assert(local_is_initialized());
    return *_instances[this_shard_id()].service;
}
 
template <typename Service>
shared_ptr<Service> sharded<Service>::local_shared() noexcept {
    assert(local_is_initialized());
    return _instances[this_shard_id()].service;
}
 
template <typename Service>
inline bool sharded<Service>::local_is_initialized() const noexcept {
    return _instances.size() > this_shard_id() &&
           _instances[this_shard_id()].service;
}
 
SEASTAR_MODULE_EXPORT_BEGIN
template <typename PtrType>
requires (!std::is_pointer_v<PtrType>)
class foreign_ptr {
private:
    PtrType _value;
    unsigned _cpu;
private:
    void destroy(PtrType p, unsigned cpu) noexcept {
        // `destroy()` is called from the destructor and other
        // synchronous methods (like `reset()`), that have no way to
        // wait for this future.
        auto f = destroy_on(std::move(p), cpu);
        if (!f.available() || f.failed()) {
            internal::run_in_background(std::move(f));
        }
    }
 
    static future<> destroy_on(PtrType p, unsigned cpu) noexcept {
        if (p) {
            if (cpu != this_shard_id()) {
                return smp::submit_to(cpu, [v = std::move(p)] () mutable {
                    // Destroy the contained pointer. We do this explicitly
                    // in the current shard, because the lambda is destroyed
                    // in the shard that submitted the task.
                    v = {};
                });
            } else {
                p = {};
            }
        }
        return make_ready_future<>();
    }
public:
    using element_type = typename std::pointer_traits<PtrType>::element_type;
    using pointer = element_type*;
 
    foreign_ptr() noexcept(std::is_nothrow_default_constructible_v<PtrType>)
        : _value(PtrType())
        , _cpu(this_shard_id()) {
    }
    foreign_ptr(std::nullptr_t) noexcept(std::is_nothrow_default_constructible_v<foreign_ptr>) : foreign_ptr() {}
    foreign_ptr(PtrType value) noexcept(std::is_nothrow_move_constructible_v<PtrType>)
        : _value(std::move(value))
        , _cpu(this_shard_id()) {
    }
    // The type is intentionally non-copyable because copies
    // are expensive because each copy requires across-CPU call.
    foreign_ptr(const foreign_ptr&) = delete;
    foreign_ptr(foreign_ptr&& other) noexcept(std::is_nothrow_move_constructible_v<PtrType>) = default;
    ~foreign_ptr() {
        destroy(std::move(_value), _cpu);
    }
    future<foreign_ptr> copy() const noexcept {
        return smp::submit_to(_cpu, [this] () mutable {
            auto v = _value;
            return make_foreign(std::move(v));
        });
    }
    element_type& operator*() const noexcept(noexcept(*_value)) { return *_value; }
    element_type* operator->() const noexcept(noexcept(&*_value)) { return &*_value; }
    pointer get() const  noexcept(noexcept(&*_value)) { return &*_value; }
    unsigned get_owner_shard() const noexcept { return _cpu; }
    operator bool() const noexcept(noexcept(static_cast<bool>(_value))) { return static_cast<bool>(_value); }
    foreign_ptr& operator=(foreign_ptr&& other) noexcept(std::is_nothrow_move_constructible_v<PtrType>) {
         destroy(std::move(_value), _cpu);
        _value = std::move(other._value);
        _cpu = other._cpu;
        return *this;
    }
    PtrType release() noexcept(std::is_nothrow_default_constructible_v<PtrType>) {
        return std::exchange(_value, {});
    }
    void reset(PtrType new_ptr) noexcept(std::is_nothrow_move_constructible_v<PtrType>) {
        auto old_ptr = std::move(_value);
        auto old_cpu = _cpu;
 
        _value = std::move(new_ptr);
        _cpu = this_shard_id();
 
        destroy(std::move(old_ptr), old_cpu);
    }
    void reset(std::nullptr_t = nullptr) noexcept(std::is_nothrow_default_constructible_v<PtrType>) {
        reset(PtrType());
    }
 
    future<> destroy() noexcept {
        return destroy_on(std::move(_value), _cpu);
    }
};
 
template <typename T>
foreign_ptr<T> make_foreign(T ptr) {
    return foreign_ptr<T>(std::move(ptr));
}
 
 
template<typename T>
struct is_smart_ptr<foreign_ptr<T>> : std::true_type {};
 
SEASTAR_MODULE_EXPORT_END
 
}