scheduling_specific.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) 2019 Scylla DB Ltd
 */
 
#ifndef SEASTAR_MODULE
#include <seastar/core/scheduling.hh>
#include <seastar/core/map_reduce.hh>
#include <seastar/util/assert.hh>
#include <seastar/util/modules.hh>
#include <array>
#include <cstdlib>
#include <map>
#include <typeindex>
#include <vector>
#include <ranges>
#endif
 
#pragma once
 
namespace seastar {
 
namespace internal {
 
struct scheduling_group_specific_thread_local_data {
    using val_ptr = std::unique_ptr<void, void (*)(void*) noexcept>;
    using cfg_ptr = lw_shared_ptr<scheduling_group_key_config>;
 
    struct specific_val {
        val_ptr valp;
        cfg_ptr cfg;
 
        static void free(void* ptr) noexcept {
            std::free(ptr);
        }
 
        specific_val() : valp(nullptr, &free), cfg(nullptr) {}
 
        specific_val(val_ptr&& valp_, const cfg_ptr& cfg_) : valp(std::move(valp_)), cfg(cfg_) {
            if (valp && cfg->constructor) {
                cfg->constructor(valp.get());
            }
        }
 
        ~specific_val() {
            if (valp && cfg->destructor) {
                cfg->destructor(valp.get());
            }
        }
 
        specific_val(const specific_val& other) = delete;
        specific_val& operator=(const specific_val& other) = delete;
 
        specific_val(specific_val&& other) : valp(std::move(other.valp)), cfg(std::move(other.cfg)) {}
 
        specific_val& operator=(specific_val&& other) {
            if (this != &other) {
                valp = std::move(other.valp);
                cfg = std::move(other.cfg);
            }
            return *this;
        }
 
        void* get() { return valp.get(); }
 
        void rename() {
            if (valp && cfg->rename) {
                cfg->rename(valp.get());
            }
        }
    };
 
    struct per_scheduling_group {
        bool queue_is_initialized = false;
        std::vector<specific_val> specific_vals;
 
        void rename() {
            for (auto& v : specific_vals) {
                v.rename();
            }
        }
    };
    std::array<per_scheduling_group, max_scheduling_groups()> per_scheduling_group_data;
    std::map<unsigned long, cfg_ptr> scheduling_group_key_configs;
};
 
#ifdef SEASTAR_BUILD_SHARED_LIBS
scheduling_group_specific_thread_local_data** get_scheduling_group_specific_thread_local_data_ptr() noexcept;
#else
inline
scheduling_group_specific_thread_local_data** get_scheduling_group_specific_thread_local_data_ptr() noexcept {
    static thread_local scheduling_group_specific_thread_local_data* data;
    return &data;
}
#endif
inline
scheduling_group_specific_thread_local_data& get_scheduling_group_specific_thread_local_data() noexcept {
    return **get_scheduling_group_specific_thread_local_data_ptr();
}
 
[[noreturn]] void no_such_scheduling_group(scheduling_group sg);
 
template<typename T>
T* scheduling_group_get_specific_ptr(scheduling_group sg, scheduling_group_key key) noexcept {
    auto& data = internal::get_scheduling_group_specific_thread_local_data();
#ifdef SEASTAR_DEBUG
    SEASTAR_ASSERT(std::type_index(typeid(T)) == data.scheduling_group_key_configs[key.id()]->type_index);
#endif
    auto sg_id = internal::scheduling_group_index(sg);
    if (__builtin_expect(sg_id < data.per_scheduling_group_data.size() &&
            data.per_scheduling_group_data[sg_id].queue_is_initialized, true)) {
        return reinterpret_cast<T*>(data.per_scheduling_group_data[sg_id].specific_vals[key.id()].get());
    }
    return nullptr;
}
 
}
 
SEASTAR_MODULE_EXPORT_BEGIN
 
template<typename T>
T& scheduling_group_get_specific(scheduling_group sg, scheduling_group_key key) {
    T* p = internal::scheduling_group_get_specific_ptr<T>(sg, std::move(key));
    if (!p) {
        internal::no_such_scheduling_group(sg);
    }
    return *p;
}
 
template<typename T>
T& scheduling_group_get_specific(scheduling_group_key key) noexcept {
    // Unlike internal::scheduling_group_get_specific_ptr, this can
    // return a reference to an element whose queue_is_initialized is
    // false.
    auto& data = internal::get_scheduling_group_specific_thread_local_data();
    SEASTAR_ASSERT(std::type_index(typeid(T)) == data.scheduling_group_key_configs[key.id()]->type_index);
    auto sg_id = internal::scheduling_group_index(current_scheduling_group());
    return *reinterpret_cast<T*>(data.per_scheduling_group_data[sg_id].specific_vals[key.id()].get());
}
 
template<typename SpecificValType, typename Mapper, typename Reducer, typename Initial>
requires requires(SpecificValType specific_val, Mapper mapper, Reducer reducer, Initial initial) {
    {reducer(initial, mapper(specific_val))} -> std::convertible_to<Initial>;
}
future<typename function_traits<Reducer>::return_type>
map_reduce_scheduling_group_specific(Mapper mapper, Reducer reducer,
        Initial initial_val, scheduling_group_key key) {
    using per_scheduling_group = internal::scheduling_group_specific_thread_local_data::per_scheduling_group;
    auto& data = internal::get_scheduling_group_specific_thread_local_data();
    auto wrapped_mapper = [key, mapper] (per_scheduling_group& psg) {
        auto id = internal::scheduling_group_key_id(key);
        return make_ready_future<typename function_traits<Mapper>::return_type>
            (mapper(*reinterpret_cast<SpecificValType*>(psg.specific_vals[id].get())));
    };
 
    return map_reduce(
            data.per_scheduling_group_data
            | std::views::filter(std::mem_fn(&per_scheduling_group::queue_is_initialized)),
            wrapped_mapper, std::move(initial_val), reducer);
}
 
template<typename SpecificValType, typename Reducer, typename Initial>
requires requires(SpecificValType specific_val, Reducer reducer, Initial initial) {
    {reducer(initial, specific_val)} -> std::convertible_to<Initial>;
}
future<typename function_traits<Reducer>::return_type>
reduce_scheduling_group_specific(Reducer reducer, Initial initial_val, scheduling_group_key key) {
    using per_scheduling_group = internal::scheduling_group_specific_thread_local_data::per_scheduling_group;
    auto& data = internal::get_scheduling_group_specific_thread_local_data();
 
    auto mapper = [key] (per_scheduling_group& psg) {
        auto id = internal::scheduling_group_key_id(key);
        return make_ready_future<SpecificValType>(*reinterpret_cast<SpecificValType*>(psg.specific_vals[id].get()));
    };
 
    return map_reduce(
            data.per_scheduling_group_data
            | std::views::filter(std::mem_fn(&per_scheduling_group::queue_is_initialized)),
            mapper, std::move(initial_val), reducer);
}
 
SEASTAR_MODULE_EXPORT_END
 
}