sharded.hh: add invoke_on variant for a shard range

Add a convenient method to invoke a function on a range of shards, to have better support for the cold control path on more complex compute models like shard groups for different tasks. A workload can benefit from such a model if inter-task cooperation is better when grouping from perf and QoS perspectives, for example. Or a computation requires a only a subset of shards due to internal concurrency limits.

include/seastar/core/sharded.hh

@@ -35,6 +35,8 @@
 #include <boost/iterator/counting_iterator.hpp>
 #include <concepts>
 #include <functional>
+#include <ranges>
+#include <type_traits>
 #endif
 
 /// \defgroup smp-module Multicore
@@ -100,6 +102,10 @@ using sharded_unwrap_evaluated_t = typename sharded_unwrap<T>::evaluated_type;
 template <typename T>
 using sharded_unwrap_t = typename sharded_unwrap<T>::type;
 
+template<typename R>
+concept unsigned_range = std::ranges::range<R>
+    && std::is_unsigned_v<std::ranges::range_value_t<R>>;
+
 } // internal
 
 
@@ -190,7 +196,7 @@ class sharded {
     };
     std::vector<entry> _instances;
 private:
-    using invoke_on_all_func_type = std::function<future<> (Service&)>;
+    using invoke_on_multiple_func_type = std::function<future<> (Service&)>;
 private:
     template <typename U, bool async>
     friend struct shared_ptr_make_helper;
@@ -338,6 +344,42 @@ public:
       }
     }
 
+    /// Invoke a callable on a range of instances of `Service`.
+    ///
+    /// \param range std::ranges::range of unsigned integers
+    /// \param options the options to forward to the \ref smp::submit_to()
+    ///         called behind the scenes.
+    /// \param func a callable with signature `Value (Service&, Args...)` or
+    ///        `future<Value> (Service&, Args...)` (for some `Value` type), or a pointer
+    ///        to a member function of Service
+    /// \param args parameters to the callable; will be copied or moved. To pass by reference,
+    ///              use std::ref().
+    /// \return Future that becomes ready once all calls have completed
+    template <typename R, typename Func, typename... Args>
+    requires std::invocable<Func, Service&, Args...>
+        && std::is_same_v<futurize_t<std::invoke_result_t<Func, Service&, internal::sharded_unwrap_t<Args>...>>, future<>>
+        && internal::unsigned_range<R>
+    future<>
+    invoke_on(R range, smp_submit_to_options options, Func func, Args... args) noexcept;
+
+    /// Invoke a callable on a range of instances of `Service`.
+    /// Passes the default \ref smp_submit_to_options to the
+    /// \ref smp::submit_to() called behind the scenes.
+    ///
+    /// \param range std::ranges::range of unsigned integers
+    /// \param func a callable with signature `Value (Service&, Args...)` or
+    ///        `future<Value> (Service&, Args...)` (for some `Value` type), or a pointer
+    ///        to a member function of Service
+    /// \param args parameters to the callable; will be copied or moved. To pass by reference,
+    ///              use std::ref().
+    /// \return Future that becomes ready once all calls have completed
+    template <typename R, typename Func, typename... Args>
+    requires std::invocable<Func, Service&, Args...>
+        && std::is_same_v<futurize_t<std::invoke_result_t<Func, Service&, internal::sharded_unwrap_t<Args>...>>, future<>>
+        && internal::unsigned_range<R>
+    future<>
+    invoke_on(R range, Func func, Args... args) noexcept;
+
     /// Invoke a callable on all instances of `Service` and reduce the results using
     /// `Reducer`.
     ///
@@ -761,7 +803,7 @@ sharded<Service>::invoke_on_all(smp_submit_to_options options, Func func, Args..
     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 invoke_on_all(options, invoke_on_multiple_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));
@@ -788,6 +830,48 @@ sharded<Service>::invoke_on_others(smp_submit_to_options options, Func func, Arg
   }
 }
 
+template <typename Service>
+template <typename R, typename Func, typename... Args>
+requires std::invocable<Func, Service&, Args...>
+    && std::is_same_v<futurize_t<std::invoke_result_t<Func, Service&, internal::sharded_unwrap_t<Args>...>>, future<>>
+    && internal::unsigned_range<R>
+inline
+future<>
+sharded<Service>::invoke_on(R range, smp_submit_to_options options, Func func, Args... args) noexcept {
+    try {
+        auto func_futurized = invoke_on_multiple_func_type([func = std::move(func), args = std::tuple(std::move(args)...)] (Service& service) mutable {
+            // Avoid false-positive unused-lambda-capture warning on Clang
+            (void)args;
+            return futurize_apply(func, std::tuple_cat(std::forward_as_tuple(service), std::tuple(internal::unwrap_sharded_arg(std::forward<Args>(args))...)));
+        });
+        return parallel_for_each(range, [this, options, func = std::move(func_futurized)] (unsigned s) {
+            if (s > smp::count - 1) {
+                throw std::invalid_argument(format("Invalid shard id in range: {}. Must be in range [0,{})", s, smp::count));
+            }
+            return smp::submit_to(s, options, [this, func] {
+                return func(*get_local_service());
+            });
+        });
+    } catch(...) {
+        return current_exception_as_future();
+    }
+}
+
+template <typename Service>
+template <typename R, typename Func, typename... Args>
+requires std::invocable<Func, Service&, Args...>
+    && std::is_same_v<futurize_t<std::invoke_result_t<Func, Service&, internal::sharded_unwrap_t<Args>...>>, future<>>
+    && internal::unsigned_range<R>
+inline
+future<> 
+sharded<Service>::invoke_on(R range, Func func, Args... args) noexcept {
+    try {
+        return invoke_on(std::forward<R>(range), smp_submit_to_options{}, std::move(func), std::move(args)...);
+    } catch(...) {
+        return current_exception_as_future();
+    }
+}
+
 template <typename Service>
 const Service& sharded<Service>::local() const noexcept {
     assert(local_is_initialized());

src/seastar.cc

@@ -75,6 +75,7 @@ module;
 #include <optional>
 #include <queue>
 #include <random>
+#include <ranges>
 #include <regex>
 #include <source_location>
 #include <sstream>

tests/unit/sharded_test.cc

@@ -21,7 +21,11 @@
 
 #include <seastar/testing/thread_test_case.hh>
 
+#include <seastar/core/shard_id.hh>
 #include <seastar/core/sharded.hh>
+#include <seastar/core/smp.hh>
+
+#include <ranges>
 
 using namespace seastar;
 
@@ -221,3 +225,80 @@ SEASTAR_THREAD_TEST_CASE(invoke_on_modifiers) {
 
     srv.stop().get();
 }
+
+class coordinator_synced_shard_map : public peering_sharded_service<coordinator_synced_shard_map> {
+    std::vector<unsigned> unsigned_per_shard;
+    unsigned coordinator_id;
+
+public:
+    coordinator_synced_shard_map(unsigned coordinator_id) : unsigned_per_shard(smp::count), coordinator_id(coordinator_id) {}
+
+    future<> sync(unsigned value) {
+        return container().invoke_on(coordinator_id, [shard_id = this_shard_id(), value] (coordinator_synced_shard_map& s) {
+            s.unsigned_per_shard[shard_id] = value;
+        });
+    }
+
+    unsigned get_synced(int shard_id) {
+        assert(this_shard_id() == coordinator_id);
+        return unsigned_per_shard[shard_id];
+    }
+};
+
+SEASTAR_THREAD_TEST_CASE(invoke_on_range_contiguous) {
+    sharded<coordinator_synced_shard_map> s;
+    auto coordinator_id = this_shard_id();
+    s.start(coordinator_id).get();
+
+    auto mid = smp::count / 2;
+    auto half1 = std::views::iota(0u, mid);
+    auto half1_id = 1;
+    auto half2 = std::views::iota(mid, smp::count);
+    auto half2_id = 2;
+
+    auto f1 = s.invoke_on(half1, [half1_id] (coordinator_synced_shard_map& s) { return s.sync(half1_id); });
+    auto f2 = s.invoke_on(half2, [half2_id] (coordinator_synced_shard_map& s) { return s.sync(half2_id); });
+    f1.get();
+    f2.get();
+
+    auto f3 = s.invoke_on(coordinator_id, [mid, half1_id, half2_id] (coordinator_synced_shard_map& s) { 
+        for (unsigned i = 0; i < mid; ++i) {
+            BOOST_REQUIRE_EQUAL(half1_id, s.get_synced(i));
+        }
+        for (unsigned i = mid; i < smp::count; ++i) {
+            BOOST_REQUIRE_EQUAL(half2_id, s.get_synced(i));
+        }
+    });
+    f3.get();
+
+    s.stop().get();
+}
+
+SEASTAR_THREAD_TEST_CASE(invoke_on_range_fragmented) {
+    sharded<coordinator_synced_shard_map> s;
+    auto coordinator_id = this_shard_id();
+    s.start(coordinator_id).get();
+
+    // TODO: migrate to C++23 std::views::stride
+    auto even = std::views::iota(0u, smp::count) | std::views::filter([](int i) { return i % 2 == 0; });
+    auto even_id = 1;
+    auto odd = std::views::iota(1u, smp::count) | std::views::filter([](int i) { return i % 2 == 1; });
+    auto odd_id = 2;
+
+    auto f1 = s.invoke_on(even, [even_id] (coordinator_synced_shard_map& s) { return s.sync(even_id); });
+    auto f2 = s.invoke_on(odd, [odd_id] (coordinator_synced_shard_map& s) { return s.sync(odd_id); });
+    f1.get();
+    f2.get();
+
+    auto f3 = s.invoke_on(coordinator_id, [even_id, odd_id] (coordinator_synced_shard_map& s) { 
+        for (unsigned i = 0; i < smp::count; i += 2) {
+            BOOST_REQUIRE_EQUAL(even_id, s.get_synced(i));
+        }
+        for (unsigned i = 1; i < smp::count; i += 2) {
+            BOOST_REQUIRE_EQUAL(odd_id, s.get_synced(i));
+        }
+    });
+    f3.get();
+
+    s.stop().get();
+}