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/modules.hh>
#include <array>
#include <typeindex>
#include <vector>
#include <ranges>
#endif
 
#pragma once
 
namespace seastar {
 
namespace internal {
 
struct scheduling_group_specific_thread_local_data {
    struct per_scheduling_group {
        bool queue_is_initialized = false;
        std::vector<void*> specific_vals;
    };
    std::array<per_scheduling_group, max_scheduling_groups()> per_scheduling_group_data;
    std::vector<scheduling_group_key_config> 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
    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()]);
    }
    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();
    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()]);
}
 
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])));
    };
 
    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]));
    };
 
    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
 
}