couchbase 3.5.1 → 3.5.2

data/ext/couchbase/core/range_scan_orchestrator.cxx

@@ -49,271 +49,289 @@ namespace couchbase::core
 mutation_state_to_snapshot_requirements(const std::optional<mutation_state>& state)
   -> std::map<std::size_t, std::optional<range_snapshot_requirements>>
 {
-    if (!state) {
-        return {};
+  if (!state) {
+    return {};
+  }
+
+  std::map<std::size_t, std::optional<range_snapshot_requirements>> requirements;
+  for (const auto& token : state->tokens) {
+    auto& requirement = requirements[token.partition_id()];
+    if (!requirement.has_value() || requirement->sequence_number < token.sequence_number()) {
+      requirement.emplace(
+        range_snapshot_requirements{ token.partition_uuid(), token.sequence_number() });
     }
-
-    std::map<std::size_t, std::optional<range_snapshot_requirements>> requirements;
-    for (const auto& token : state->tokens) {
-        auto& requirement = requirements[token.partition_id()];
-        if (!requirement.has_value() || requirement->sequence_number < token.sequence_number()) {
-            requirement.emplace(range_snapshot_requirements{ token.partition_uuid(), token.sequence_number() });
-        }
-    }
-    return requirements;
+  }
+  return requirements;
 }
 
 // Sent by the vbucket scan stream when it either completes or fails with a fatal error
 struct scan_stream_end_signal {
-    std::uint16_t vbucket_id;
-    std::optional<std::error_code> error{};
+  std::uint16_t vbucket_id;
+  std::optional<std::error_code> error{};
 };
 
 class range_scan_stream : public std::enable_shared_from_this<range_scan_stream>
 {
-    // The stream has failed and should not be retried
-    struct failed {
-        std::error_code ec;
-        bool fatal{ true };
-    };
-
-    // The stream is currently running
-    struct running {
-        std::vector<std::byte> uuid;
-    };
-
-    // The stream has completed and the items have been retrieved
-    struct completed {
-    };
-
-  public:
-    range_scan_stream(asio::io_context& io,
-                      agent kv_provider,
-                      std::uint16_t vbucket_id,
-                      std::int16_t node_id,
-                      range_scan_create_options create_options,
-                      range_scan_continue_options continue_options,
-                      std::shared_ptr<scan_stream_manager> stream_manager)
-      : agent_{ std::move(kv_provider) }
-      , io_{ io }
-      , vbucket_id_{ vbucket_id }
-      , node_id_{ node_id }
-      , create_options_{ std::move(create_options) }
-      , continue_options_{ std::move(continue_options) }
-      , stream_manager_{ std::move(stream_manager) }
-    {
+  // The stream has failed and should not be retried
+  struct failed {
+    std::error_code ec;
+    bool fatal{ true };
+  };
+
+  // The stream is currently running
+  struct running {
+    std::vector<std::byte> uuid;
+  };
+
+  // The stream has completed and the items have been retrieved
+  struct completed {
+  };
+
+public:
+  range_scan_stream(asio::io_context& io,
+                    agent kv_provider,
+                    std::uint16_t vbucket_id,
+                    std::int16_t node_id,
+                    range_scan_create_options create_options,
+                    range_scan_continue_options continue_options,
+                    std::shared_ptr<scan_stream_manager> stream_manager)
+    : agent_{ std::move(kv_provider) }
+    , io_{ io }
+    , vbucket_id_{ vbucket_id }
+    , node_id_{ node_id }
+    , create_options_{ std::move(create_options) }
+    , continue_options_{ std::move(continue_options) }
+    , stream_manager_{ std::move(stream_manager) }
+  {
+  }
+
+  void start()
+  {
+    // Fail the stream if more time than the timeout has elapsed since the stream was first
+    // attempted (if this is a retry)
+    if (first_attempt_timestamp_.has_value()) {
+      if (std::chrono::steady_clock::now() - first_attempt_timestamp_.value() >
+          create_options_.timeout) {
+        CB_LOG_DEBUG(
+          "stream for vbucket_id {} cannot be retried because it has exceeded the timeout",
+          vbucket_id_);
+        state_ = failed{ errc::common::unambiguous_timeout, !is_sampling_scan() };
+        stream_manager_->stream_failed(
+          node_id_, vbucket_id_, errc::common::unambiguous_timeout, error_is_fatal());
+        return;
+      }
+    } else {
+      first_attempt_timestamp_ = std::chrono::steady_clock::now();
     }
 
-    void start()
-    {
-        // Fail the stream if more time than the timeout has elapsed since the stream was first attempted (if this is a retry)
-        if (first_attempt_timestamp_.has_value()) {
-            if (std::chrono::steady_clock::now() - first_attempt_timestamp_.value() > create_options_.timeout) {
-                CB_LOG_DEBUG("stream for vbucket_id {} cannot be retried because it has exceeded the timeout", vbucket_id_);
-                state_ = failed{ errc::common::unambiguous_timeout, !is_sampling_scan() };
-                stream_manager_->stream_failed(node_id_, vbucket_id_, errc::common::unambiguous_timeout, error_is_fatal());
-                return;
-            }
-        } else {
-            first_attempt_timestamp_ = std::chrono::steady_clock::now();
-        }
-
-        CB_LOG_TRACE("starting stream for vbucket {} in node {}", vbucket_id_, node_id_);
-
-        if (std::holds_alternative<range_scan>(create_options_.scan_type) && !last_seen_key_.empty()) {
-            std::get<range_scan>(create_options_.scan_type).from = scan_term{ last_seen_key_ };
-        }
-
-        agent_.range_scan_create(vbucket_id_, create_options_, [self = shared_from_this()](auto res, auto ec) {
-            if (ec) {
-                if (ec == errc::key_value::document_not_found) {
-                    // Benign error
-                    CB_LOG_TRACE("ignoring vbucket_id {} because no documents exist for it", self->vbucket_id_);
-                    self->state_ = failed{ ec, false };
-                    self->stream_manager_->stream_failed(self->node_id_, self->vbucket_id_, ec, self->error_is_fatal());
-                } else if (ec == errc::common::temporary_failure) {
-                    // Retryable error - server is overwhelmed, retry after reducing concurrency
-                    CB_LOG_DEBUG("received busy status during scan from vbucket with ID {} - reducing concurrency & retrying",
-                                 self->vbucket_id_);
-                    self->state_ = std::monostate{};
-                    self->stream_manager_->stream_start_failed_awaiting_retry(self->node_id_, self->vbucket_id_);
-                } else if (ec == errc::common::internal_server_failure || ec == errc::common::collection_not_found) {
-                    // Fatal errors
-                    self->state_ = failed{ ec, true };
-                    self->stream_manager_->stream_failed(self->node_id_, self->vbucket_id_, ec, self->error_is_fatal());
-                } else {
-                    // Unexpected errors
-                    CB_LOG_DEBUG("received unexpected error {} from stream for vbucket {} during range scan continue ({})",
-                                 ec.value(),
-                                 self->vbucket_id_,
-                                 ec.message());
-                    self->state_ = failed{ ec, true };
-                    self->stream_manager_->stream_failed(self->node_id_, self->vbucket_id_, ec, self->error_is_fatal());
-                }
-                return;
-            }
-
-            self->state_ = running{ std::move(res.scan_uuid) };
+    CB_LOG_TRACE("starting stream for vbucket {} in node {}", vbucket_id_, node_id_);
 
-            return self->resume();
-        });
-    }
-
-    void should_cancel()
-    {
-        should_cancel_ = true;
-    }
-
-    [[nodiscard]] auto node_id() const -> std::int16_t
-    {
-        return node_id_;
+    if (std::holds_alternative<range_scan>(create_options_.scan_type) && !last_seen_key_.empty()) {
+      std::get<range_scan>(create_options_.scan_type).from = scan_term{ last_seen_key_ };
     }
 
-  private:
-    void fail(std::error_code ec)
-    {
-        if (is_failed()) {
-            return;
-        }
-
-        bool fatal;
-        if (ec == errc::key_value::document_not_found || ec == errc::common::authentication_failure ||
-            ec == errc::common::collection_not_found || ec == errc::common::request_canceled) {
-            // Errors that are fatal unless this is a sampling scan
-            fatal = !is_sampling_scan();
-        } else if (ec == errc::common::feature_not_available || ec == errc::common::invalid_argument ||
-                   ec == errc::common::temporary_failure) {
-            // Errors that are always fatal
-            fatal = true;
-        } else {
-            // Unexpected error - always fatal
-            CB_LOG_DEBUG("received unexpected error {} from stream for vbucket {} during range scan continue ({})",
+    agent_.range_scan_create(
+      vbucket_id_, create_options_, [self = shared_from_this()](auto res, auto ec) {
+        if (ec) {
+          if (ec == errc::key_value::document_not_found) {
+            // Benign error
+            CB_LOG_TRACE("ignoring vbucket_id {} because no documents exist for it",
+                         self->vbucket_id_);
+            self->state_ = failed{ ec, false };
+            self->stream_manager_->stream_failed(
+              self->node_id_, self->vbucket_id_, ec, self->error_is_fatal());
+          } else if (ec == errc::common::temporary_failure) {
+            // Retryable error - server is overwhelmed, retry after reducing concurrency
+            CB_LOG_DEBUG("received busy status during scan from vbucket with ID {} - reducing "
+                         "concurrency & retrying",
+                         self->vbucket_id_);
+            self->state_ = std::monostate{};
+            self->stream_manager_->stream_start_failed_awaiting_retry(self->node_id_,
+                                                                      self->vbucket_id_);
+          } else if (ec == errc::common::internal_server_failure ||
+                     ec == errc::common::collection_not_found) {
+            // Fatal errors
+            self->state_ = failed{ ec, true };
+            self->stream_manager_->stream_failed(
+              self->node_id_, self->vbucket_id_, ec, self->error_is_fatal());
+          } else {
+            // Unexpected errors
+            CB_LOG_DEBUG("received unexpected error {} from stream for vbucket {} during range "
+                         "scan continue ({})",
                          ec.value(),
-                         vbucket_id_,
+                         self->vbucket_id_,
                          ec.message());
-            fatal = true;
+            self->state_ = failed{ ec, true };
+            self->stream_manager_->stream_failed(
+              self->node_id_, self->vbucket_id_, ec, self->error_is_fatal());
+          }
+          return;
         }
 
-        state_ = failed{ ec, fatal };
-        stream_manager_->stream_failed(node_id_, vbucket_id_, ec, fatal);
-    }
+        self->state_ = running{ std::move(res.scan_uuid) };
 
-    void complete()
-    {
-        if (is_failed() || is_completed()) {
-            return;
-        }
+        return self->resume();
+      });
+  }
 
-        stream_manager_->stream_completed(node_id_, vbucket_id_);
-        state_ = completed{};
-    }
+  void should_cancel()
+  {
+    should_cancel_ = true;
+  }
 
-    void cancel()
-    {
-        auto scan_uuid = uuid();
-        if (scan_uuid.empty()) {
-            // The stream is not currently running
-            return;
-        }
+  [[nodiscard]] auto node_id() const -> std::int16_t
+  {
+    return node_id_;
+  }
 
-        asio::post(asio::bind_executor(io_, [self = shared_from_this(), scan_uuid]() mutable {
-            self->agent_.range_scan_cancel(scan_uuid, self->vbucket_id_, {}, [](auto /* res */, auto /* ec */) {});
-        }));
+private:
+  void fail(std::error_code ec)
+  {
+    if (is_failed()) {
+      return;
     }
 
-    [[nodiscard]] auto is_running() const -> bool
-    {
-        return std::holds_alternative<running>(state_);
+    bool fatal;
+    if (ec == errc::key_value::document_not_found || ec == errc::common::authentication_failure ||
+        ec == errc::common::collection_not_found || ec == errc::common::request_canceled) {
+      // Errors that are fatal unless this is a sampling scan
+      fatal = !is_sampling_scan();
+    } else if (ec == errc::common::feature_not_available || ec == errc::common::invalid_argument ||
+               ec == errc::common::temporary_failure) {
+      // Errors that are always fatal
+      fatal = true;
+    } else {
+      // Unexpected error - always fatal
+      CB_LOG_DEBUG(
+        "received unexpected error {} from stream for vbucket {} during range scan continue ({})",
+        ec.value(),
+        vbucket_id_,
+        ec.message());
+      fatal = true;
     }
 
-    [[nodiscard]] auto is_failed() const -> bool
-    {
-        return std::holds_alternative<failed>(state_);
-    }
+    state_ = failed{ ec, fatal };
+    stream_manager_->stream_failed(node_id_, vbucket_id_, ec, fatal);
+  }
 
-    [[nodiscard]] auto is_completed() const -> bool
-    {
-        return std::holds_alternative<completed>(state_);
+  void complete()
+  {
+    if (is_failed() || is_completed()) {
+      return;
     }
 
-    void resume()
-    {
-        if (!is_running()) {
-            return;
-        }
-        if (should_cancel_) {
-            cancel();
-            return;
-        }
+    stream_manager_->stream_completed(node_id_, vbucket_id_);
+    state_ = completed{};
+  }
 
-        asio::post(asio::bind_executor(io_, [self = shared_from_this()]() mutable {
-            self->agent_.range_scan_continue(
-              self->uuid(),
-              self->vbucket_id_,
-              self->continue_options_,
-              [self](auto item) {
-                  // The scan has already been cancelled, no need to send items
-                  if (self->should_cancel_) {
-                      return;
-                  }
-                  self->last_seen_key_ = item.key;
-                  self->stream_manager_->stream_received_item(std::move(item));
-              },
-              [self](auto res, auto ec) {
-                  if (ec) {
-                      return self->fail(ec);
-                  }
-                  if (res.complete) {
-                      return self->complete();
-                  }
-                  if (res.more) {
-                      return self->resume();
-                  }
-              });
-        }));
+  void cancel()
+  {
+    auto scan_uuid = uuid();
+    if (scan_uuid.empty()) {
+      // The stream is not currently running
+      return;
     }
 
-    [[nodiscard]] auto uuid() const -> std::vector<std::byte>
-    {
-        try {
-            return std::get<running>(state_).uuid;
-        } catch (std::bad_variant_access&) {
-            return {};
-        }
+    asio::post(asio::bind_executor(io_, [self = shared_from_this(), scan_uuid]() mutable {
+      self->agent_.range_scan_cancel(
+        scan_uuid, self->vbucket_id_, {}, [](auto /* res */, auto /* ec */) {
+        });
+    }));
+  }
+
+  [[nodiscard]] auto is_running() const -> bool
+  {
+    return std::holds_alternative<running>(state_);
+  }
+
+  [[nodiscard]] auto is_failed() const -> bool
+  {
+    return std::holds_alternative<failed>(state_);
+  }
+
+  [[nodiscard]] auto is_completed() const -> bool
+  {
+    return std::holds_alternative<completed>(state_);
+  }
+
+  void resume()
+  {
+    if (!is_running()) {
+      return;
     }
-
-    [[nodiscard]] auto error() const -> std::error_code
-    {
-        if (is_failed()) {
-            return std::get<failed>(state_).ec;
-        }
-        return {};
+    if (should_cancel_) {
+      cancel();
+      return;
     }
 
-    [[nodiscard]] auto error_is_fatal() const -> bool
-    {
-        if (is_failed()) {
-            return std::get<failed>(state_).fatal;
-        }
-        return {};
+    asio::post(asio::bind_executor(io_, [self = shared_from_this()]() mutable {
+      self->agent_.range_scan_continue(
+        self->uuid(),
+        self->vbucket_id_,
+        self->continue_options_,
+        [self](auto item) {
+          // The scan has already been cancelled, no need to send items
+          if (self->should_cancel_) {
+            return;
+          }
+          self->last_seen_key_ = item.key;
+          self->stream_manager_->stream_received_item(std::move(item));
+        },
+        [self](auto res, auto ec) {
+          if (ec) {
+            return self->fail(ec);
+          }
+          if (res.complete) {
+            return self->complete();
+          }
+          if (res.more) {
+            return self->resume();
+          }
+        });
+    }));
+  }
+
+  [[nodiscard]] auto uuid() const -> std::vector<std::byte>
+  {
+    try {
+      return std::get<running>(state_).uuid;
+    } catch (std::bad_variant_access&) {
+      return {};
     }
+  }
 
-    [[nodiscard]] auto is_sampling_scan() const -> bool
-    {
-        return std::holds_alternative<sampling_scan>(create_options_.scan_type);
+  [[nodiscard]] auto error() const -> std::error_code
+  {
+    if (is_failed()) {
+      return std::get<failed>(state_).ec;
     }
+    return {};
+  }
 
-    agent agent_;
-    asio::io_context& io_;
-    std::uint16_t vbucket_id_;
-    std::int16_t node_id_;
-    range_scan_create_options create_options_;
-    range_scan_continue_options continue_options_;
-    std::shared_ptr<scan_stream_manager> stream_manager_;
-    std::string last_seen_key_{};
-    std::variant<std::monostate, failed, running, completed> state_{};
-    bool should_cancel_{ false };
-    std::optional<std::chrono::time_point<std::chrono::steady_clock>> first_attempt_timestamp_{};
+  [[nodiscard]] auto error_is_fatal() const -> bool
+  {
+    if (is_failed()) {
+      return std::get<failed>(state_).fatal;
+    }
+    return {};
+  }
+
+  [[nodiscard]] auto is_sampling_scan() const -> bool
+  {
+    return std::holds_alternative<sampling_scan>(create_options_.scan_type);
+  }
+
+  agent agent_;
+  asio::io_context& io_;
+  std::uint16_t vbucket_id_;
+  std::int16_t node_id_;
+  range_scan_create_options create_options_;
+  range_scan_continue_options continue_options_;
+  std::shared_ptr<scan_stream_manager> stream_manager_;
+  std::string last_seen_key_{};
+  std::variant<std::monostate, failed, running, completed> state_{};
+  std::atomic<bool> should_cancel_{ false };
+  std::optional<std::chrono::time_point<std::chrono::steady_clock>> first_attempt_timestamp_{};
 };
 
 class range_scan_orchestrator_impl
@@ -321,263 +339,290 @@ class range_scan_orchestrator_impl
   , public range_scan_item_iterator
   , public scan_stream_manager
 {
-  public:
-    range_scan_orchestrator_impl(asio::io_context& io,
-                                 agent kv_provider,
-                                 topology::configuration::vbucket_map vbucket_map,
-                                 std::string scope_name,
-                                 std::string collection_name,
-                                 std::variant<std::monostate, range_scan, prefix_scan, sampling_scan> scan_type,
-                                 range_scan_orchestrator_options options)
-      : io_{ io }
-      , agent_{ std::move(kv_provider) }
-      , vbucket_map_{ std::move(vbucket_map) }
-      , scope_name_{ std::move(scope_name) }
-      , collection_name_{ std::move(collection_name) }
-      , load_balancer_{ vbucket_map_ }
-      , items_{ io, 1024 }
-      , scan_type_{ std::move(scan_type) }
-      , options_{ std::move(options) }
-      , vbucket_to_snapshot_requirements_{ mutation_state_to_snapshot_requirements(options_.consistent_with) }
-      , concurrency_{ options_.concurrency }
-    {
-
-        if (std::holds_alternative<sampling_scan>(scan_type_)) {
-            auto s = std::get<sampling_scan>(scan_type);
-            item_limit_ = s.limit;
-
-            // Set the seed of the load balancer to ensure that if the sampling scan is run multiple times the vbuckets
-            // are scanned in the same order when concurrency is 1. This guarantees that the items returned will be the
-            // same. We cannot guarantee this when concurrency is greater than 1, as the order of the vbucket scans
-            // depends on how long each scan takes and what the load on a node is at any given time.
-            if (s.seed.has_value()) {
-                load_balancer_.seed(s.seed.value());
-            }
-        }
+public:
+  range_scan_orchestrator_impl(
+    asio::io_context& io,
+    agent kv_provider,
+    topology::configuration::vbucket_map vbucket_map,
+    std::string scope_name,
+    std::string collection_name,
+    std::variant<std::monostate, range_scan, prefix_scan, sampling_scan> scan_type,
+    range_scan_orchestrator_options options)
+    : io_{ io }
+    , agent_{ std::move(kv_provider) }
+    , vbucket_map_{ std::move(vbucket_map) }
+    , scope_name_{ std::move(scope_name) }
+    , collection_name_{ std::move(collection_name) }
+    , load_balancer_{ vbucket_map_ }
+    , items_{ io, 1024 }
+    , scan_type_{ std::move(scan_type) }
+    , options_{ std::move(options) }
+    , vbucket_to_snapshot_requirements_{ mutation_state_to_snapshot_requirements(
+        options_.consistent_with) }
+    , concurrency_{ options_.concurrency }
+  {
+
+    if (std::holds_alternative<sampling_scan>(scan_type_)) {
+      auto s = std::get<sampling_scan>(scan_type);
+      item_limit_ = s.limit;
+
+      // Set the seed of the load balancer to ensure that if the sampling scan is run multiple times
+      // the vbuckets are scanned in the same order when concurrency is 1. This guarantees that the
+      // items returned will be the same. We cannot guarantee this when concurrency is greater than
+      // 1, as the order of the vbucket scans depends on how long each scan takes and what the load
+      // on a node is at any given time.
+      if (s.seed.has_value()) {
+        load_balancer_.seed(s.seed.value());
+      }
     }
+  }
 
-    void scan(scan_callback&& cb)
-    {
-        if (item_limit_ == 0 || concurrency_ <= 0) {
-            return cb(errc::common::invalid_argument, {});
-        }
-
-        get_collection_id_options const get_cid_options{ options_.retry_strategy, options_.timeout, options_.parent_span };
-        agent_.get_collection_id(
-          scope_name_,
-          collection_name_,
-          get_cid_options,
-          [self = shared_from_this(), cb = std::move(cb)](auto get_cid_res, auto ec) mutable {
-              if (ec) {
-                  return cb(ec, {});
-              }
-              self->collection_id_ = get_cid_res.collection_id;
-
-              auto batch_time_limit = std::chrono::duration_cast<std::chrono::milliseconds>(0.9 * self->options_.timeout);
-              range_scan_continue_options const continue_options{
-                  self->options_.batch_item_limit, self->options_.batch_byte_limit, batch_time_limit,
-                  self->options_.timeout,          self->options_.retry_strategy,
-              };
-
-              for (std::uint16_t vbucket = 0; vbucket < gsl::narrow_cast<std::uint16_t>(self->vbucket_map_.size()); ++vbucket) {
-                  const range_scan_create_options create_options{
-                      self->scope_name_,       {},
-                      self->scan_type_,        self->options_.timeout,
-                      self->collection_id_,    self->vbucket_to_snapshot_requirements_[vbucket],
-                      self->options_.ids_only, self->options_.retry_strategy,
-                  };
-
-                  // Get the active node for the vbucket (values in vbucket map are the active node id followed by the ids of the replicas)
-                  auto node_id = self->vbucket_map_[vbucket][0];
-
-                  auto stream = std::make_shared<range_scan_stream>(self->io_,
-                                                                    self->agent_,
-                                                                    vbucket,
-                                                                    node_id,
-                                                                    create_options,
-                                                                    continue_options,
-                                                                    std::static_pointer_cast<scan_stream_manager>(self));
-                  self->streams_[vbucket] = stream;
-              }
-              self->start_streams(self->concurrency_);
-              return cb({}, scan_result(self));
-          });
+  void scan(scan_callback&& cb)
+  {
+    if (item_limit_ == 0 || concurrency_ <= 0) {
+      return cb(errc::common::invalid_argument, {});
     }
 
-    void cancel() override
-    {
-        cancelled_ = true;
-        for (const auto& [vbucket_id, stream] : streams_) {
-            stream->should_cancel();
+    get_collection_id_options const get_cid_options{ options_.retry_strategy,
+                                                     options_.timeout,
+                                                     options_.parent_span };
+    agent_.get_collection_id(
+      scope_name_,
+      collection_name_,
+      get_cid_options,
+      [self = shared_from_this(), cb = std::move(cb)](auto get_cid_res, auto ec) mutable {
+        if (ec) {
+          return cb(ec, {});
         }
-    }
-
-    bool is_cancelled() override
-    {
-        return cancelled_;
-    }
-
-    auto next() -> std::future<tl::expected<range_scan_item, std::error_code>> override
-    {
-        auto barrier = std::make_shared<std::promise<tl::expected<range_scan_item, std::error_code>>>();
-        next([barrier](range_scan_item item, std::error_code ec) mutable {
-            if (ec) {
-                barrier->set_value(tl::unexpected{ ec });
-            } else {
-                barrier->set_value(std::move(item));
-            }
-        });
-        return barrier->get_future();
-    }
-
-    void next(utils::movable_function<void(range_scan_item, std::error_code)> callback) override
-    {
-        if (item_limit_ == 0 || item_limit_-- == 0) {
-            callback({}, errc::key_value::range_scan_completed);
-            cancel();
-        } else {
-            next_item(std::move(callback));
+        self->collection_id_ = get_cid_res.collection_id;
+
+        auto batch_time_limit =
+          std::chrono::duration_cast<std::chrono::milliseconds>(0.9 * self->options_.timeout);
+        range_scan_continue_options const continue_options{
+          self->options_.batch_item_limit, self->options_.batch_byte_limit, batch_time_limit,
+          self->options_.timeout,          self->options_.retry_strategy,
+        };
+
+        for (std::uint16_t vbucket = 0;
+             vbucket < gsl::narrow_cast<std::uint16_t>(self->vbucket_map_.size());
+             ++vbucket) {
+          const range_scan_create_options create_options{
+            self->scope_name_,       self->collection_name_,
+            self->scan_type_,        self->options_.timeout,
+            self->collection_id_,    self->vbucket_to_snapshot_requirements_[vbucket],
+            self->options_.ids_only, self->options_.retry_strategy,
+          };
+
+          // Get the active node for the vbucket (values in vbucket map are the active node id
+          // followed by the ids of the replicas)
+          auto node_id = self->vbucket_map_[vbucket][0];
+
+          auto stream = std::make_shared<range_scan_stream>(
+            self->io_,
+            self->agent_,
+            vbucket,
+            node_id,
+            create_options,
+            continue_options,
+            std::static_pointer_cast<scan_stream_manager>(self));
+          self->streams_[vbucket] = stream;
         }
+        self->start_streams(self->concurrency_);
+        return cb({}, scan_result(self));
+      });
+  }
+
+  void cancel() override
+  {
+    cancelled_ = true;
+    for (const auto& [vbucket_id, stream] : streams_) {
+      stream->should_cancel();
     }
-
-    template<typename Handler>
-    void next_item(Handler&& handler)
-    {
-        if (streams_.empty() || cancelled_) {
-            items_.cancel();
-            items_.close();
-            return handler({}, errc::key_value::range_scan_completed);
-        }
-        items_.async_receive([self = shared_from_this(), handler = std::forward<Handler>(handler)](
-                               std::error_code ec, std::variant<range_scan_item, scan_stream_end_signal> it) mutable {
-            if (ec) {
-                return handler({}, ec);
-            }
-
-            if (std::holds_alternative<range_scan_item>(it)) {
-                handler(std::get<range_scan_item>(it), {});
-            } else {
-                auto signal = std::get<scan_stream_end_signal>(it);
-                if (signal.error.has_value()) {
-                    // Fatal error
-                    handler({}, signal.error.value());
-                } else {
-                    // Empty signal means that stream has completed
-                    {
-                        std::lock_guard<std::mutex> const lock{ self->stream_map_mutex_ };
-                        self->streams_.erase(signal.vbucket_id);
-                    }
-                    return asio::post(asio::bind_executor(self->io_, [self, handler = std::forward<Handler>(handler)]() mutable {
-                        self->next_item(std::forward<Handler>(handler));
-                    }));
-                }
-            }
-        });
+  }
+
+  bool is_cancelled() override
+  {
+    return cancelled_;
+  }
+
+  auto next() -> std::future<tl::expected<range_scan_item, std::error_code>> override
+  {
+    auto barrier = std::make_shared<std::promise<tl::expected<range_scan_item, std::error_code>>>();
+    next([barrier](range_scan_item item, std::error_code ec) mutable {
+      if (ec) {
+        barrier->set_value(tl::unexpected{ ec });
+      } else {
+        barrier->set_value(std::move(item));
+      }
+    });
+    return barrier->get_future();
+  }
+
+  void next(utils::movable_function<void(range_scan_item, std::error_code)> callback) override
+  {
+    if (item_limit_ == 0 || item_limit_-- == 0) {
+      callback({}, errc::key_value::range_scan_completed);
+      cancel();
+    } else {
+      next_item(std::move(callback));
     }
-
-    void start_streams(std::uint16_t stream_count)
-    {
-        if (cancelled_) {
-            CB_LOG_TRACE("scan has been cancelled, do not start another stream");
-            return;
+  }
+
+  template<typename Handler>
+  void next_item(Handler&& handler)
+  {
+    if (streams_.empty() || cancelled_) {
+      items_.cancel();
+      items_.close();
+      return handler({}, errc::key_value::range_scan_completed);
+    }
+    items_.async_receive(
+      [self = shared_from_this(), handler = std::forward<Handler>(handler)](
+        std::error_code ec, std::variant<range_scan_item, scan_stream_end_signal> it) mutable {
+        if (ec) {
+          return handler({}, ec);
         }
 
-        std::uint16_t counter{ 0 };
-        while (counter < stream_count) {
-            auto vbucket_id = load_balancer_.select_vbucket();
-            if (!vbucket_id.has_value()) {
-                CB_LOG_TRACE("no more scans, all vbuckets have been scanned");
-                return;
-            }
-
-            auto v = vbucket_id.value();
-            std::shared_ptr<range_scan_stream> stream{};
+        if (std::holds_alternative<range_scan_item>(it)) {
+          handler(std::get<range_scan_item>(it), {});
+        } else {
+          auto signal = std::get<scan_stream_end_signal>(it);
+          if (signal.error.has_value()) {
+            // Fatal error
+            handler({}, signal.error.value());
+          } else {
+            // Empty signal means that stream has completed
             {
-                std::lock_guard<std::mutex> const lock{ stream_map_mutex_ };
-                stream = streams_.at(v);
+              std::lock_guard<std::mutex> const lock{ self->stream_map_mutex_ };
+              self->streams_.erase(signal.vbucket_id);
             }
-            CB_LOG_TRACE("scanning vbucket {} at node {}", vbucket_id.value(), stream->node_id());
-            active_stream_count_++;
-            counter++;
-            asio::post(asio::bind_executor(io_, [stream]() mutable { stream->start(); }));
+            return asio::post(asio::bind_executor(
+              self->io_, [self, handler = std::forward<Handler>(handler)]() mutable {
+                self->next_item(std::forward<Handler>(handler));
+              }));
+          }
         }
+      });
+  }
+
+  void start_streams(std::uint16_t stream_count)
+  {
+    if (cancelled_) {
+      CB_LOG_TRACE("scan has been cancelled, do not start another stream");
+      return;
     }
 
-    void stream_received_item(range_scan_item item) override
-    {
-        items_.async_send({}, std::move(item), [](std::error_code ec) {
-            if (ec && ec != asio::experimental::error::channel_closed && ec != asio::experimental::error::channel_cancelled) {
-                CB_LOG_WARNING("unexpected error while sending to scan item channel: {} ({})", ec.value(), ec.message());
-            }
-        });
+    std::uint16_t counter{ 0 };
+    while (counter < stream_count) {
+      auto vbucket_id = load_balancer_.select_vbucket();
+      if (!vbucket_id.has_value()) {
+        CB_LOG_TRACE("no more scans, all vbuckets have been scanned");
+        return;
+      }
+
+      auto v = vbucket_id.value();
+      std::shared_ptr<range_scan_stream> stream{};
+      {
+        std::lock_guard<std::mutex> const lock{ stream_map_mutex_ };
+        stream = streams_.at(v);
+      }
+      CB_LOG_TRACE("scanning vbucket {} at node {}", vbucket_id.value(), stream->node_id());
+      active_stream_count_++;
+      counter++;
+      asio::post(asio::bind_executor(io_, [stream]() mutable {
+        stream->start();
+      }));
     }
-
-    void stream_failed(std::int16_t node_id, std::uint16_t vbucket_id, std::error_code ec, bool fatal) override
-    {
-        if (!fatal) {
-            return stream_completed(node_id, vbucket_id);
-        }
-
-        load_balancer_.notify_stream_ended(node_id);
-        active_stream_count_--;
-        items_.async_send({}, scan_stream_end_signal{ vbucket_id, ec }, [](std::error_code ec) {
-            if (ec && ec != asio::experimental::error::channel_closed && ec != asio::experimental::error::channel_cancelled) {
-                CB_LOG_WARNING("unexpected error while sending to scan item channel: {} ({})", ec.value(), ec.message());
-            }
-        });
-        return cancel();
+  }
+
+  void stream_received_item(range_scan_item item) override
+  {
+    items_.async_send({}, std::move(item), [](std::error_code ec) {
+      if (ec && ec != asio::experimental::error::channel_closed &&
+          ec != asio::experimental::error::channel_cancelled) {
+        CB_LOG_WARNING(
+          "unexpected error while sending to scan item channel: {} ({})", ec.value(), ec.message());
+      }
+    });
+  }
+
+  void stream_failed(std::int16_t node_id,
+                     std::uint16_t vbucket_id,
+                     std::error_code ec,
+                     bool fatal) override
+  {
+    if (!fatal) {
+      return stream_completed(node_id, vbucket_id);
     }
 
-    void stream_completed(std::int16_t node_id, std::uint16_t vbucket_id) override
-    {
-        load_balancer_.notify_stream_ended(node_id);
-        active_stream_count_--;
-        items_.async_send({}, scan_stream_end_signal{ vbucket_id }, [](std::error_code ec) {
-            if (ec && ec != asio::experimental::error::channel_closed && ec != asio::experimental::error::channel_cancelled) {
-                CB_LOG_WARNING("unexpected error while sending to scan item channel: {} ({})", ec.value(), ec.message());
-            }
-        });
-        return start_streams(1);
-    }
+    load_balancer_.notify_stream_ended(node_id);
+    active_stream_count_--;
+    items_.async_send({}, scan_stream_end_signal{ vbucket_id, ec }, [](std::error_code ec) {
+      if (ec && ec != asio::experimental::error::channel_closed &&
+          ec != asio::experimental::error::channel_cancelled) {
+        CB_LOG_WARNING(
+          "unexpected error while sending to scan item channel: {} ({})", ec.value(), ec.message());
+      }
+    });
+    return cancel();
+  }
+
+  void stream_completed(std::int16_t node_id, std::uint16_t vbucket_id) override
+  {
+    load_balancer_.notify_stream_ended(node_id);
+    active_stream_count_--;
+    items_.async_send({}, scan_stream_end_signal{ vbucket_id }, [](std::error_code ec) {
+      if (ec && ec != asio::experimental::error::channel_closed &&
+          ec != asio::experimental::error::channel_cancelled) {
+        CB_LOG_WARNING(
+          "unexpected error while sending to scan item channel: {} ({})", ec.value(), ec.message());
+      }
+    });
+    return start_streams(1);
+  }
 
-    void stream_start_failed_awaiting_retry(std::int16_t node_id, std::uint16_t vbucket_id) override
-    {
-        load_balancer_.notify_stream_ended(node_id);
-        active_stream_count_--;
+  void stream_start_failed_awaiting_retry(std::int16_t node_id, std::uint16_t vbucket_id) override
+  {
+    load_balancer_.notify_stream_ended(node_id);
+    active_stream_count_--;
 
-        load_balancer_.enqueue_vbucket(node_id, vbucket_id);
-        if (active_stream_count_ == 0) {
-            return start_streams(1);
-        }
+    load_balancer_.enqueue_vbucket(node_id, vbucket_id);
+    if (active_stream_count_ == 0) {
+      return start_streams(1);
     }
-
-  private:
-    asio::io_context& io_;
-    agent agent_;
-    topology::configuration::vbucket_map vbucket_map_;
-    std::string scope_name_;
-    std::string collection_name_;
-    range_scan_load_balancer load_balancer_;
-    asio::experimental::concurrent_channel<void(std::error_code, std::variant<range_scan_item, scan_stream_end_signal>)> items_;
-    std::uint32_t collection_id_{ 0 };
-    std::variant<std::monostate, range_scan, prefix_scan, sampling_scan> scan_type_;
-    range_scan_orchestrator_options options_;
-    std::map<std::size_t, std::optional<range_snapshot_requirements>> vbucket_to_snapshot_requirements_;
-    std::map<std::uint16_t, std::shared_ptr<range_scan_stream>> streams_{};
-    std::mutex stream_map_mutex_{};
-    std::atomic_uint16_t active_stream_count_{ 0 };
-    std::uint16_t concurrency_{ 1 };
-    std::size_t item_limit_{ std::numeric_limits<std::size_t>::max() };
-    bool cancelled_{ false };
+  }
+
+private:
+  asio::io_context& io_;
+  agent agent_;
+  topology::configuration::vbucket_map vbucket_map_;
+  std::string scope_name_;
+  std::string collection_name_;
+  range_scan_load_balancer load_balancer_;
+  asio::experimental::concurrent_channel<
+    void(std::error_code, std::variant<range_scan_item, scan_stream_end_signal>)>
+    items_;
+  std::uint32_t collection_id_{ 0 };
+  std::variant<std::monostate, range_scan, prefix_scan, sampling_scan> scan_type_;
+  range_scan_orchestrator_options options_;
+  std::map<std::size_t, std::optional<range_snapshot_requirements>>
+    vbucket_to_snapshot_requirements_;
+  std::map<std::uint16_t, std::shared_ptr<range_scan_stream>> streams_{};
+  std::mutex stream_map_mutex_{};
+  std::atomic_uint16_t active_stream_count_{ 0 };
+  std::uint16_t concurrency_{ 1 };
+  std::size_t item_limit_{ std::numeric_limits<std::size_t>::max() };
+  std::atomic<bool> cancelled_{ false };
 };
 
-range_scan_orchestrator::range_scan_orchestrator(asio::io_context& io,
-                                                 agent kv_provider,
-                                                 topology::configuration::vbucket_map vbucket_map,
-                                                 std::string scope_name,
-                                                 std::string collection_name,
-                                                 std::variant<std::monostate, range_scan, prefix_scan, sampling_scan> scan_type,
-                                                 range_scan_orchestrator_options options)
+range_scan_orchestrator::range_scan_orchestrator(
+  asio::io_context& io,
+  agent kv_provider,
+  topology::configuration::vbucket_map vbucket_map,
+  std::string scope_name,
+  std::string collection_name,
+  std::variant<std::monostate, range_scan, prefix_scan, sampling_scan> scan_type,
+  range_scan_orchestrator_options options)
   : impl_{ std::make_shared<range_scan_orchestrator_impl>(io,
                                                           std::move(kv_provider),
                                                           std::move(vbucket_map),
@@ -591,20 +636,20 @@ range_scan_orchestrator::range_scan_orchestrator(asio::io_context& io,
 auto
 range_scan_orchestrator::scan() -> tl::expected<scan_result, std::error_code>
 {
-    auto barrier = std::make_shared<std::promise<tl::expected<scan_result, std::error_code>>>();
-    auto f = barrier->get_future();
-    scan([barrier](auto ec, auto res) mutable {
-        if (ec) {
-            return barrier->set_value(tl::unexpected(ec));
-        }
-        barrier->set_value(res);
-    });
-    return f.get();
+  auto barrier = std::make_shared<std::promise<tl::expected<scan_result, std::error_code>>>();
+  auto f = barrier->get_future();
+  scan([barrier](auto ec, auto res) mutable {
+    if (ec) {
+      return barrier->set_value(tl::unexpected(ec));
+    }
+    barrier->set_value(res);
+  });
+  return f.get();
 }
 
 void
 range_scan_orchestrator::scan(couchbase::core::scan_callback&& cb)
 {
-    return impl_->scan(std::move(cb));
+  return impl_->scan(std::move(cb));
 }
 } // namespace couchbase::core