Merge 'when_all: add Sentinel support to when_all_succeed() ' from Ke…

…fu Chai

Extend `when_all_succeed()` and `iterator_range_estimate_vector_capacity()` to accept ranges with heterogeneous iterator
and sentinel types by adding a separate template parameter for the sentinel.

Closes #2565

* https://github.com/scylladb/seastar:
  when_all: add Sentinel support to when_all_succeed()
  loop: add Sentinel support to iterator_range_estimate_vector_capacity()

include/seastar/core/loop.hh

@@ -503,17 +503,20 @@ struct has_iterator_category : std::false_type {};
 template <typename T>
 struct has_iterator_category<T, std::void_t<typename std::iterator_traits<T>::iterator_category >> : std::true_type {};
 
-template <typename Iterator, typename Sentinel, typename IteratorCategory>
+template <typename Iterator, typename Sentinel>
 inline
 size_t
-iterator_range_estimate_vector_capacity(Iterator begin, Sentinel end, IteratorCategory) {
-    // May be linear time below random_access_iterator_tag, but still better than reallocation
-    if constexpr (std::is_base_of_v<std::forward_iterator_tag, IteratorCategory>) {
-        return std::distance(begin, end);
+iterator_range_estimate_vector_capacity(Iterator begin, Sentinel end) {
+    if constexpr (std::forward_iterator<Iterator> &&
+                  std::forward_iterator<Sentinel>) {
+        return std::ranges::distance(begin, end);
+    } else if constexpr (std::random_access_iterator<Iterator> &&
+                         std::random_access_iterator<Sentinel>) {
+        return std::ranges::distance(begin, end);
+    } else {
+        // For InputIterators we can't estimate needed capacity
+        return 0;
     }
-
-    // For InputIterators we can't estimate needed capacity
-    return 0;
 }
 
 } // namespace internal
@@ -577,12 +580,11 @@ parallel_for_each(Iterator begin, Sentinel end, Func&& func) noexcept {
         memory::scoped_critical_alloc_section _;
         if (!f.available() || f.failed()) {
             if (!s) {
-                using itraits = std::iterator_traits<Iterator>;
                 size_t n{0U};
                 if constexpr (internal::has_iterator_category<Iterator>::value) {
                     // We need if-constexpr here because there exist iterators for which std::iterator_traits
                     // does not have 'iterator_category' as member type
-                    n = (internal::iterator_range_estimate_vector_capacity(begin, end, typename itraits::iterator_category{}) + 1);
+                    n = (internal::iterator_range_estimate_vector_capacity(begin, end) + 1);
                 }
                 s = new parallel_for_each_state(n);
             }

include/seastar/core/when_all.hh

@@ -286,9 +286,9 @@ complete_when_all(std::vector<Future>&& futures, typename std::vector<Future>::i
     });
 }
 
-template<typename ResolvedVectorTransform, typename FutureIterator>
+template<typename ResolvedVectorTransform, typename FutureIterator, typename Sentinel>
 inline auto
-do_when_all(FutureIterator begin, FutureIterator end) noexcept {
+do_when_all(FutureIterator begin, Sentinel end) noexcept {
     using itraits = std::iterator_traits<FutureIterator>;
     auto make_values_vector = [] (size_t size) noexcept {
         memory::scoped_critical_alloc_section _;
@@ -297,10 +297,10 @@ do_when_all(FutureIterator begin, FutureIterator end) noexcept {
         return ret;
     };
     std::vector<typename itraits::value_type> ret =
-            make_values_vector(iterator_range_estimate_vector_capacity(begin, end, typename itraits::iterator_category()));
+            make_values_vector(iterator_range_estimate_vector_capacity(begin, end));
     // Important to invoke the *begin here, in case it's a function iterator,
     // so we launch all computation in parallel.
-    std::move(begin, end, std::back_inserter(ret));
+    std::ranges::move(begin, end, std::back_inserter(ret));
     return complete_when_all<ResolvedVectorTransform>(std::move(ret), ret.begin());
 }
 
@@ -508,14 +508,14 @@ inline auto when_all_succeed(FutOrFuncs&&... fut_or_funcs) noexcept {
 /// \param end an \c InputIterator designating the end of the range of futures
 /// \return an \c std::vector<> of all the valus in the input
 SEASTAR_MODULE_EXPORT
-template <typename FutureIterator, typename = typename std::iterator_traits<FutureIterator>::value_type>
+template <typename FutureIterator, typename Sentinel, typename = typename std::iterator_traits<FutureIterator>::value_type>
 requires requires (FutureIterator i) {
      *i++;
      { i != i } -> std::convertible_to<bool>;
      requires is_future<std::remove_reference_t<decltype(*i)>>::value;
 }
 inline auto
-when_all_succeed(FutureIterator begin, FutureIterator end) noexcept {
+when_all_succeed(FutureIterator begin, Sentinel end) noexcept {
     using itraits = std::iterator_traits<FutureIterator>;
     using result_transform = internal::extract_values_from_futures_vector<typename itraits::value_type>;
     try {

include/seastar/coroutine/parallel_for_each.hh

@@ -128,9 +128,8 @@ public:
             } else {
                 memory::scoped_critical_alloc_section _;
                 if (_futures.empty()) {
-                    using itraits = std::iterator_traits<Iterator>;
                     if constexpr (seastar::internal::has_iterator_category<Iterator>::value) {
-                        auto n = seastar::internal::iterator_range_estimate_vector_capacity(it, end, typename itraits::iterator_category{});
+                        auto n = seastar::internal::iterator_range_estimate_vector_capacity(it, end);
                         _futures.reserve(n);
                     }
                 }

tests/unit/futures_test.cc

@@ -538,17 +538,30 @@ SEASTAR_TEST_CASE(test_when_all_iterator_range) {
 
 template<typename Container>
 void test_iterator_range_estimate() {
-    using iter_traits = std::iterator_traits<typename Container::iterator>;
     Container container{1,2,3};
 
     BOOST_REQUIRE_EQUAL(internal::iterator_range_estimate_vector_capacity(
-        container.begin(), container.end(), typename iter_traits::iterator_category{}), 3);
+        container.begin(), container.end()), 3);
 }
 
 BOOST_AUTO_TEST_CASE(test_iterator_range_estimate_vector_capacity) {
     test_iterator_range_estimate<std::vector<int>>();
     test_iterator_range_estimate<std::list<int>>();
     test_iterator_range_estimate<std::forward_list<int>>();
+    {
+        int n = 42;
+        auto seq = std::views::iota(0, n);
+        BOOST_REQUIRE_EQUAL(internal::iterator_range_estimate_vector_capacity(
+            seq.begin(), seq.end()), n);
+    }
+    {
+        // for ranges that generate elements on-the-fly, advancing an iterator
+        // might actually consume or transform the underlying sequence, in this
+        // case, the function under test returns 0.
+        auto seq = std::views::iota(1);
+        BOOST_REQUIRE_EQUAL(internal::iterator_range_estimate_vector_capacity(
+            seq.begin(), seq.end()), 0);
+    }
 }
 
 // helper function for when_any tests