// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

use std::{fmt::Debug, marker::PhantomData, time::Duration};

use api_version::KvFormat;
use paste::paste;
use tikv_util::{cooperative_checkpoint, cooperative_safe_points_active, time::Instant};
use yatp::task::future::reschedule;

use super::{OwnedKvPair, OwnedKvPairEntry, Storage, range::*, ranges_iter::*};
use crate::error::StorageError;

const KEY_BUFFER_CAPACITY: usize = 64;
/// Batch executors are run in coroutines. This is the maximum time a coroutine
/// can run without being yielded when the legacy checker is in use.
pub const DEFAULT_RANGE_SCAN_MAX_TIME_SLICE_US: u64 = 3_000;
/// the number of scanned keys that should trigger a reschedule.
const CHECK_KEYS: usize = 32;

/// Caller-owned rescheduling policy for one range scan.
///
/// A request snapshots this once when it starts, so an online update does not
/// change the cadence partway through that request.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ScanRescheduleOptions {
    /// Number of completed point-get or scan-next operations between
    /// cooperative checkpoints. Zero disables these caller-owned checks.
    pub check_keys: usize,
    /// Wall time a coroutine may run between yields in the legacy checker.
    ///
    /// This bounds only the legacy path. When a cooperative safe-point context
    /// is active the read pool owns the slice instead, and this value is
    /// unused.
    pub max_time_slice: Duration,
}

impl Default for ScanRescheduleOptions {
    /// Disables the caller-owned key checks and keeps the legacy wall-time
    /// slice. Production callers set both fields from configuration.
    fn default() -> Self {
        Self {
            check_keys: 0,
            max_time_slice: Duration::from_micros(DEFAULT_RANGE_SCAN_MAX_TIME_SLICE_US),
        }
    }
}

/// A scanner that scans over multiple ranges. Each range can be a point range
/// containing only one row, or an interval range containing multiple rows.
pub struct RangesScanner<T, F> {
    storage: T,
    ranges_iter: RangesIterator,

    scan_backward_in_range: bool,
    is_key_only: bool,
    load_commit_ts: bool,

    scanned_rows_per_range: Vec<usize>,

    // The following fields are only used for calculating scanned range. Scanned range is only
    // useful in streaming mode, where the client need to know the underlying physical data range
    // of each response slice, so that partial retry can be non-overlapping.
    current_range: IntervalRange,
    working_range_begin_key: Vec<u8>,
    working_range_end_key: Vec<u8>,

    _phantom: PhantomData<F>,
    reschedule: ScanRescheduleOptions,
    keys_since_reschedule_check: usize,
    rescheduler: RescheduleChecker,
}

// TODO: maybe it's better to make it generic to avoid directly depending
// on yatp's rescheduler.
struct RescheduleChecker {
    prev_start: Instant,
    prev_key_count: usize,
    max_time_slice: Duration,
}

impl RescheduleChecker {
    fn new(max_time_slice: Duration) -> Self {
        Self {
            prev_start: Instant::now_coarse(),
            prev_key_count: 0,
            max_time_slice,
        }
    }

    #[inline(always)]
    async fn check_reschedule(&mut self, force_check: bool) {
        self.prev_key_count += 1;
        if (force_check || self.prev_key_count.is_multiple_of(CHECK_KEYS))
            && self.prev_start.saturating_elapsed() > self.max_time_slice
        {
            if tokio::runtime::Handle::try_current().is_ok() {
                tokio::task::yield_now().await;
            } else {
                reschedule().await;
            }
            self.prev_start = Instant::now_coarse();
            self.prev_key_count = 0;
        }
    }
}

pub struct RangesScannerOptions<T> {
    pub storage: T,
    pub ranges: Vec<Range>,
    pub scan_backward_in_range: bool, // TODO: This can be const generics
    pub is_key_only: bool,            // TODO: This can be const generics
    pub load_commit_ts: bool,
    /// Rescheduling policy this scan uses, snapshotted by the caller.
    pub reschedule: ScanRescheduleOptions,
}

pub struct IndexedKvPair {
    key: Vec<u8>,
    value: Vec<u8>,
    keyspace: Option<u32>,
    offset: usize,
    commit_ts: Option<u64>,
}

impl IndexedKvPair {
    pub fn from_kv_pair<F: KvFormat>((k, v): OwnedKvPair) -> Result<Self, StorageError> {
        Self::from_kv_pair_entry::<F>(OwnedKvPairEntry {
            key: k,
            value: v,
            commit_ts: None,
        })
    }

    pub fn from_kv_pair_entry<F: KvFormat>(entry: OwnedKvPairEntry) -> Result<Self, StorageError> {
        let (keyspace, user_key) =
            F::strip_keyspace(&entry.key).map_err(|e| StorageError(anyhow::Error::new(e)))?;
        let offset = entry.key.len() - user_key.len();

        Ok(IndexedKvPair {
            key: entry.key,
            value: entry.value,
            keyspace,
            offset,
            commit_ts: entry.commit_ts,
        })
    }

    pub fn key(&self) -> &[u8] {
        &self.key[self.offset..]
    }

    pub fn value(&self) -> &[u8] {
        &self.value
    }

    pub fn kv(&self) -> (&[u8], &[u8]) {
        (self.key(), self.value())
    }

    pub fn keyspace(&self) -> Option<u32> {
        self.keyspace
    }

    pub fn commit_ts(&self) -> Option<u64> {
        self.commit_ts
    }
}

impl PartialEq for IndexedKvPair {
    fn eq(&self, other: &IndexedKvPair) -> bool {
        self.kv() == other.kv()
    }
}

impl PartialEq<OwnedKvPair> for IndexedKvPair {
    fn eq(&self, other: &OwnedKvPair) -> bool {
        self.kv() == (&other.0, &other.1)
    }
}

impl Debug for IndexedKvPair {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("IndexedKvPair")
            .field("key", &log_wrappers::Value(self.key()))
            .field("value", &log_wrappers::Value(self.value()))
            .field("keyspace", &self.keyspace)
            .field("commit_ts", &self.commit_ts)
            .finish()
    }
}

/// These macros helps to use existed sync codes when IA is not enabled.
///
/// TODO: remove the macros and always use async methods.
///
/// `try_sync(storage, get, (key))` will generate:
///
/// ```text
/// if storage.is_snapshot_sync() {
///     storage.get(key)
/// } else {
///     storage.get_async(key).await
/// }
/// ```
macro_rules! try_sync {
    ($expr:expr, $fn:ident) => {{
        if $expr.is_snapshot_sync() {
            $expr.$fn()
        } else {
            paste! {$expr.[<$fn _async>]().await }
        }
    }};
    ($expr:expr, $fn:ident, ($($args:expr),*)) => {{
        if $expr.is_snapshot_sync() {
            $expr.$fn($($args),*)
        } else {
            paste! {$expr.[<$fn _async>]($($args),*).await }
        }
    }};
}

macro_rules! begin_scan {
    ($expr:expr, ($($args:expr),*)) => {{
        try_sync!($expr, begin_scan, ($($args),*))
    }};
}

macro_rules! scan_next_entry {
    ($expr:expr) => {{
        if $expr.is_snapshot_sync() && !cooperative_safe_points_active() {
            $expr.scan_next_entry()
        } else {
            $expr.scan_next_entry_async().await
        }
    }};
}

macro_rules! get_entry {
    ($expr:expr, ($($args:expr),*)) => {{
        try_sync!($expr, get_entry, ($($args),*))
    }};
}

impl<T: Storage, F: KvFormat> RangesScanner<T, F> {
    pub fn new(
        RangesScannerOptions {
            storage,
            ranges,
            scan_backward_in_range,
            is_key_only,
            load_commit_ts,
            reschedule,
        }: RangesScannerOptions<T>,
    ) -> RangesScanner<T, F> {
        let ranges_len = ranges.len();
        let ranges_iter = RangesIterator::new(ranges);
        RangesScanner {
            storage,
            ranges_iter,
            scan_backward_in_range,
            is_key_only,
            load_commit_ts,
            scanned_rows_per_range: Vec::with_capacity(ranges_len),
            current_range: IntervalRange {
                lower_inclusive: Vec::with_capacity(KEY_BUFFER_CAPACITY),
                upper_exclusive: Vec::with_capacity(KEY_BUFFER_CAPACITY),
            },
            working_range_begin_key: Vec::with_capacity(KEY_BUFFER_CAPACITY),
            working_range_end_key: Vec::with_capacity(KEY_BUFFER_CAPACITY),
            _phantom: PhantomData,
            reschedule,
            keys_since_reschedule_check: 0,
            rescheduler: RescheduleChecker::new(reschedule.max_time_slice),
        }
    }

    async fn check_read_budget_after_storage_op(&mut self) {
        if self.reschedule.check_keys == 0 {
            return;
        }
        self.keys_since_reschedule_check = self.keys_since_reschedule_check.saturating_add(1);
        if self.keys_since_reschedule_check < self.reschedule.check_keys {
            return;
        }
        self.keys_since_reschedule_check %= self.reschedule.check_keys;

        cooperative_checkpoint().await;
    }

    /// Fetches next row.
    // Note: This is not implemented over `Iterator` since it can fail.
    // TODO: Change to use reference to avoid allocation and copy.
    pub async fn next(&mut self) -> Result<Option<IndexedKvPair>, StorageError> {
        self.next_opt(true).await
    }

    /// Fetches next row.
    /// Note: `update_scanned_range` can control whether update the scanned
    /// range.
    pub async fn next_opt(
        &mut self,
        update_scanned_range: bool,
    ) -> Result<Option<IndexedKvPair>, StorageError> {
        loop {
            let mut force_check = true;
            let range = self.ranges_iter.next();
            let some_row = match range {
                IterStatus::NewRange(Range::Point(r)) => {
                    self.update_scanned_range_from_new_point(&r);
                    self.ranges_iter.notify_drained();
                    self.scanned_rows_per_range.push(0);
                    get_entry!(self.storage, (self.is_key_only, self.load_commit_ts, r))?
                }
                IterStatus::NewRange(Range::Interval(r)) => {
                    self.update_scanned_range_from_new_range(&r);
                    self.scanned_rows_per_range.push(0);
                    begin_scan!(
                        self.storage,
                        (
                            self.scan_backward_in_range,
                            self.is_key_only,
                            self.load_commit_ts,
                            r
                        )
                    )?;
                    scan_next_entry!(self.storage)?
                }
                IterStatus::Continue => {
                    force_check = false;
                    scan_next_entry!(self.storage)?
                }
                IterStatus::Drained => {
                    self.update_working_range_end_key();
                    return Ok(None); // drained
                }
            };
            if update_scanned_range {
                self.update_scanned_range_from_scanned_row(&some_row);
            }

            if cooperative_safe_points_active() {
                // The read-pool observer installs the new context only for
                // requests using the new fairness-compatible reschedule
                // policy. This layer owns the loop over point gets and
                // scan-next calls, including empty ranges, so it owns this
                // local cadence as well. Wall-time and fair-CPU helpers each
                // inspect their own context, so either capability can be
                // active independently.
                self.check_read_budget_after_storage_op().await;
            } else if some_row.is_some() {
                // Keep the legacy coroutine behavior for Tokio, remote
                // coprocessor, fairness-disabled, and other callers without a
                // poll-scoped context. In particular, a new range still
                // forces the old wall-time check.
                self.rescheduler.check_reschedule(force_check).await;
            }

            if let Some(row) = some_row {
                // Retrieved one row from point range or interval range.
                if let Some(r) = self.scanned_rows_per_range.last_mut() {
                    *r += 1;
                }

                let row = IndexedKvPair::from_kv_pair_entry::<F>(row)?;
                return Ok(Some(row));
            } else {
                // No more row in the range.
                self.ranges_iter.notify_drained();
            }
        }
    }

    /// Appends storage statistics collected so far to the given container and
    /// clears the collected statistics.
    pub fn collect_storage_stats(&mut self, dest: &mut T::Statistics) {
        self.storage.collect_statistics(dest)
    }

    /// Appends scanned rows of each range so far to the given container and
    /// clears the collected statistics.
    pub fn collect_scanned_rows_per_range(&mut self, dest: &mut Vec<usize>) {
        dest.append(&mut self.scanned_rows_per_range);
        self.scanned_rows_per_range.push(0);
    }

    pub fn peek_scanned_bytes_sum(&self) -> usize {
        self.storage.scanned_bytes()
    }

    /// Returns scanned range since last call.
    pub fn take_scanned_range(&mut self) -> IntervalRange {
        let mut range = IntervalRange::default();
        if !self.scan_backward_in_range {
            std::mem::swap(
                &mut range.lower_inclusive,
                &mut self.working_range_begin_key,
            );
            std::mem::swap(&mut range.upper_exclusive, &mut self.working_range_end_key);

            self.working_range_begin_key
                .extend_from_slice(&range.upper_exclusive);
        } else {
            std::mem::swap(&mut range.lower_inclusive, &mut self.working_range_end_key);
            std::mem::swap(
                &mut range.upper_exclusive,
                &mut self.working_range_begin_key,
            );

            self.working_range_begin_key
                .extend_from_slice(&range.lower_inclusive);
        }

        range
    }

    #[inline]
    pub fn can_be_cached(&self) -> bool {
        self.storage.met_uncacheable_data() == Some(false)
    }

    fn update_scanned_range_from_new_point(&mut self, point: &PointRange) {
        // Only update current_range for the first and the last range.
        if self.current_range.lower_inclusive.is_empty() || self.ranges_iter.is_drained() {
            self.current_range.lower_inclusive.clear();
            self.current_range.upper_exclusive.clear();
            self.current_range
                .lower_inclusive
                .extend_from_slice(&point.0);
            self.current_range
                .upper_exclusive
                .extend_from_slice(&point.0);
            self.current_range.upper_exclusive.push(0);
        }
        self.update_working_range_begin_key();
    }

    fn update_scanned_range_from_new_range(&mut self, range: &IntervalRange) {
        // Only update current_range for the first and the last range.
        if self.current_range.lower_inclusive.is_empty() || self.ranges_iter.is_drained() {
            self.current_range.lower_inclusive.clear();
            self.current_range.upper_exclusive.clear();
            self.current_range
                .lower_inclusive
                .extend_from_slice(&range.lower_inclusive);
            self.current_range
                .upper_exclusive
                .extend_from_slice(&range.upper_exclusive);
        }
        self.update_working_range_begin_key();
    }

    fn update_working_range_begin_key(&mut self) {
        if self.working_range_begin_key.is_empty() {
            if !self.scan_backward_in_range {
                self.working_range_begin_key
                    .extend(&self.current_range.lower_inclusive);
            } else {
                self.working_range_begin_key
                    .extend(&self.current_range.upper_exclusive);
            }
        }
    }

    fn update_working_range_end_key(&mut self) {
        self.working_range_end_key.clear();
        if !self.scan_backward_in_range {
            self.working_range_end_key
                .extend(&self.current_range.upper_exclusive);
        } else {
            self.working_range_end_key
                .extend(&self.current_range.lower_inclusive);
        }
    }

    fn update_scanned_range_from_scanned_row(&mut self, some_row: &Option<OwnedKvPairEntry>) {
        if let Some(OwnedKvPairEntry { key, .. }) = some_row {
            self.working_range_end_key.clear();
            self.working_range_end_key.extend(key);
            if !self.scan_backward_in_range {
                self.working_range_end_key.push(0);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::Instant as StdInstant;

    use api_version::ApiV1;
    use futures::executor::block_on;
    use tikv_util::task_reschedule::{DEFAULT_MAX_TIME_SLICE_US, PollRescheduleContextGuard};

    use super::*;
    use crate::storage::{IntervalRange, PointRange, Range, test_fixture::FixtureStorage};

    fn create_storage() -> FixtureStorage {
        let data: &[(&'static [u8], &'static [u8])] = &[
            (b"bar", b"2"),
            (b"bar_2", b"4"),
            (b"foo", b"1"),
            (b"foo_2", b"3"),
            (b"foo_3", b"5"),
        ];
        FixtureStorage::from(data)
    }

    #[test]
    fn test_reschedule_uses_local_context_cadence_or_legacy_fallback() {
        let ranges = vec![IntervalRange::from(("foo", "foo_3a")).into()];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage: create_storage(),
            ranges,
            scan_backward_in_range: false,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions {
                check_keys: 2,
                ..Default::default()
            },
        });

        {
            let _guard = PollRescheduleContextGuard::new(
                StdInstant::now(),
                Duration::from_micros(DEFAULT_MAX_TIME_SLICE_US),
            );
            block_on(scanner.next()).unwrap().unwrap();
            assert_eq!(scanner.keys_since_reschedule_check, 1);
            block_on(scanner.next()).unwrap().unwrap();
            assert_eq!(scanner.keys_since_reschedule_check, 0);
            assert_eq!(scanner.rescheduler.prev_key_count, 0);
        }
        block_on(scanner.next()).unwrap().unwrap();
        assert_eq!(scanner.rescheduler.prev_key_count, 1);
    }

    #[test]
    fn test_legacy_checker_uses_the_configured_slice() {
        let ranges = vec![IntervalRange::from(("foo", "foo_3a")).into()];
        let scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage: create_storage(),
            ranges,
            scan_backward_in_range: false,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions {
                max_time_slice: Duration::from_micros(1_500),
                ..Default::default()
            },
        });
        assert_eq!(
            scanner.rescheduler.max_time_slice,
            Duration::from_micros(1_500)
        );

        // The default keeps the slice this checker has always used.
        let ranges = vec![IntervalRange::from(("foo", "foo_3a")).into()];
        let scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage: create_storage(),
            ranges,
            scan_backward_in_range: false,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });
        assert_eq!(scanner.rescheduler.max_time_slice, Duration::from_millis(3));
    }

    #[test]
    fn test_next() {
        let storage = create_storage();

        // Currently we accept unordered ranges.
        let ranges: Vec<Range> = vec![
            IntervalRange::from(("foo", "foo_2a")).into(),
            PointRange::from("foo_2b").into(),
            PointRange::from("foo_3").into(),
            IntervalRange::from(("a", "c")).into(),
        ];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage: storage.clone(),
            ranges,
            scan_backward_in_range: false,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });
        assert_eq!(
            block_on(scanner.next()).unwrap().unwrap(),
            (b"foo".to_vec(), b"1".to_vec())
        );
        assert_eq!(
            block_on(scanner.next()).unwrap().unwrap(),
            (b"foo_2".to_vec(), b"3".to_vec())
        );
        assert_eq!(
            block_on(scanner.next()).unwrap().unwrap(),
            (b"foo_3".to_vec(), b"5".to_vec())
        );
        assert_eq!(
            block_on(scanner.next()).unwrap().unwrap(),
            (b"bar".to_vec(), b"2".to_vec())
        );
        assert_eq!(
            block_on(scanner.next()).unwrap().unwrap(),
            (b"bar_2".to_vec(), b"4".to_vec())
        );
        assert_eq!(block_on(scanner.next()).unwrap(), None);

        // Backward in range
        let ranges: Vec<Range> = vec![
            IntervalRange::from(("foo", "foo_2a")).into(),
            PointRange::from("foo_2b").into(),
            PointRange::from("foo_3").into(),
            IntervalRange::from(("a", "bar_2")).into(),
        ];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage: storage.clone(),
            ranges,
            scan_backward_in_range: true,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });
        assert_eq!(
            block_on(scanner.next()).unwrap().unwrap(),
            (b"foo_2".to_vec(), b"3".to_vec())
        );
        assert_eq!(
            block_on(scanner.next()).unwrap().unwrap(),
            (b"foo".to_vec(), b"1".to_vec())
        );
        assert_eq!(
            block_on(scanner.next()).unwrap().unwrap(),
            (b"foo_3".to_vec(), b"5".to_vec())
        );
        assert_eq!(
            block_on(scanner.next()).unwrap().unwrap(),
            (b"bar".to_vec(), b"2".to_vec())
        );
        assert_eq!(block_on(scanner.next()).unwrap(), None);

        // Key only
        let ranges: Vec<Range> = vec![
            IntervalRange::from(("bar", "foo_2a")).into(),
            PointRange::from("foo_3").into(),
            PointRange::from("bar_3").into(),
        ];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage,
            ranges,
            scan_backward_in_range: false,
            is_key_only: true,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });
        assert_eq!(
            block_on(scanner.next()).unwrap().unwrap(),
            (b"bar".to_vec(), Vec::new())
        );
        assert_eq!(
            block_on(scanner.next()).unwrap().unwrap(),
            (b"bar_2".to_vec(), Vec::new())
        );
        assert_eq!(
            block_on(scanner.next()).unwrap().unwrap(),
            (b"foo".to_vec(), Vec::new())
        );
        assert_eq!(
            block_on(scanner.next()).unwrap().unwrap(),
            (b"foo_2".to_vec(), Vec::new())
        );
        assert_eq!(
            block_on(scanner.next()).unwrap().unwrap(),
            (b"foo_3".to_vec(), Vec::new())
        );
        assert_eq!(block_on(scanner.next()).unwrap(), None);
    }

    #[test]
    fn test_scanned_rows() {
        let storage = create_storage();

        let ranges: Vec<Range> = vec![
            IntervalRange::from(("foo", "foo_2a")).into(),
            PointRange::from("foo_2b").into(),
            PointRange::from("foo_3").into(),
            IntervalRange::from(("a", "z")).into(),
        ];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage,
            ranges,
            scan_backward_in_range: false,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });
        let mut scanned_rows_per_range = Vec::new();

        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"foo");
        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"foo_2");
        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"foo_3");

        scanner.collect_scanned_rows_per_range(&mut scanned_rows_per_range);
        assert_eq!(scanned_rows_per_range, vec![2, 0, 1]);
        scanned_rows_per_range.clear();

        scanner.collect_scanned_rows_per_range(&mut scanned_rows_per_range);
        assert_eq!(scanned_rows_per_range, vec![0]);
        scanned_rows_per_range.clear();

        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"bar");
        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"bar_2");

        scanner.collect_scanned_rows_per_range(&mut scanned_rows_per_range);
        assert_eq!(scanned_rows_per_range, vec![0, 2]);
        scanned_rows_per_range.clear();

        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"foo");

        scanner.collect_scanned_rows_per_range(&mut scanned_rows_per_range);
        assert_eq!(scanned_rows_per_range, vec![1]);
        scanned_rows_per_range.clear();

        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"foo_2");
        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"foo_3");
        assert_eq!(block_on(scanner.next()).unwrap(), None);

        scanner.collect_scanned_rows_per_range(&mut scanned_rows_per_range);
        assert_eq!(scanned_rows_per_range, vec![2]);
        scanned_rows_per_range.clear();

        assert_eq!(block_on(scanner.next()).unwrap(), None);

        scanner.collect_scanned_rows_per_range(&mut scanned_rows_per_range);
        assert_eq!(scanned_rows_per_range, vec![0]);
        scanned_rows_per_range.clear();
    }

    #[test]
    fn test_scanned_bytes() {
        let storage = create_storage();
        let ranges: Vec<Range> = vec![
            IntervalRange::from(("foo", "foo_3a")).into(),
            PointRange::from("bar_2").into(),
            PointRange::from("bar_3").into(),
        ];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage,
            ranges,
            scan_backward_in_range: false,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });

        assert_eq!(scanner.peek_scanned_bytes_sum(), 0);

        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"foo");
        assert_eq!(scanner.peek_scanned_bytes_sum(), 4);

        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"foo_2");
        assert_eq!(scanner.peek_scanned_bytes_sum(), 10);

        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"foo_3");
        assert_eq!(scanner.peek_scanned_bytes_sum(), 16);

        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"bar_2");
        assert_eq!(scanner.peek_scanned_bytes_sum(), 22);

        assert_eq!(block_on(scanner.next()).unwrap(), None);
        assert_eq!(scanner.peek_scanned_bytes_sum(), 22);
    }

    #[test]
    fn test_scanned_range_forward() {
        let storage = create_storage();

        // No range
        let ranges = vec![];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage: storage.clone(),
            ranges,
            scan_backward_in_range: false,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"");
        assert_eq!(&r.upper_exclusive, b"");

        assert_eq!(block_on(scanner.next()).unwrap(), None);

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"");
        assert_eq!(&r.upper_exclusive, b"");

        // Empty interval range
        let ranges = vec![IntervalRange::from(("x", "xb")).into()];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage: storage.clone(),
            ranges,
            scan_backward_in_range: false,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });

        assert_eq!(block_on(scanner.next()).unwrap(), None);

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"x");
        assert_eq!(&r.upper_exclusive, b"xb");

        // Empty point range
        let ranges = vec![PointRange::from("x").into()];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage: storage.clone(),
            ranges,
            scan_backward_in_range: false,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });

        assert_eq!(block_on(scanner.next()).unwrap(), None);

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"x");
        assert_eq!(&r.upper_exclusive, b"x\0");

        // Filled interval range
        let ranges = vec![IntervalRange::from(("foo", "foo_8")).into()];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage: storage.clone(),
            ranges,
            scan_backward_in_range: false,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });

        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"foo");
        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"foo_2");

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"foo");
        assert_eq!(&r.upper_exclusive, b"foo_2\0");

        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"foo_3");

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"foo_2\0");
        assert_eq!(&r.upper_exclusive, b"foo_3\0");

        assert_eq!(block_on(scanner.next()).unwrap(), None);

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"foo_3\0");
        assert_eq!(&r.upper_exclusive, b"foo_8");

        // Multiple ranges
        // TODO: caller should not pass in unordered ranges otherwise scanned ranges
        // would be unsound.
        let ranges = vec![
            IntervalRange::from(("foo", "foo_3")).into(),
            IntervalRange::from(("foo_5", "foo_50")).into(),
            IntervalRange::from(("bar", "bar_")).into(),
            PointRange::from("bar_2").into(),
            PointRange::from("bar_3").into(),
            IntervalRange::from(("bar_4", "box")).into(),
        ];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage,
            ranges,
            scan_backward_in_range: false,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });

        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"foo");

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"foo");
        assert_eq!(&r.upper_exclusive, b"foo\0");

        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"foo_2");

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"foo\0");
        assert_eq!(&r.upper_exclusive, b"foo_2\0");

        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"bar");

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"foo_2\0");
        assert_eq!(&r.upper_exclusive, b"bar\0");

        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"bar_2");

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"bar\0");
        assert_eq!(&r.upper_exclusive, b"bar_2\0");

        assert_eq!(block_on(scanner.next()).unwrap(), None);

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"bar_2\0");
        assert_eq!(&r.upper_exclusive, b"box");
    }

    #[test]
    fn test_scanned_range_backward() {
        let storage = create_storage();

        // No range
        let ranges = vec![];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage: storage.clone(),
            ranges,
            scan_backward_in_range: true,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"");
        assert_eq!(&r.upper_exclusive, b"");

        assert_eq!(block_on(scanner.next()).unwrap(), None);

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"");
        assert_eq!(&r.upper_exclusive, b"");

        // Empty interval range
        let ranges = vec![IntervalRange::from(("x", "xb")).into()];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage: storage.clone(),
            ranges,
            scan_backward_in_range: true,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });

        assert_eq!(block_on(scanner.next()).unwrap(), None);

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"x");
        assert_eq!(&r.upper_exclusive, b"xb");

        // Empty point range
        let ranges = vec![PointRange::from("x").into()];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage: storage.clone(),
            ranges,
            scan_backward_in_range: true,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });

        assert_eq!(block_on(scanner.next()).unwrap(), None);

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"x");
        assert_eq!(&r.upper_exclusive, b"x\0");

        // Filled interval range
        let ranges = vec![IntervalRange::from(("foo", "foo_8")).into()];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage: storage.clone(),
            ranges,
            scan_backward_in_range: true,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });

        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"foo_3");
        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"foo_2");

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"foo_2");
        assert_eq!(&r.upper_exclusive, b"foo_8");

        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"foo");

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"foo");
        assert_eq!(&r.upper_exclusive, b"foo_2");

        assert_eq!(block_on(scanner.next()).unwrap(), None);

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"foo");
        assert_eq!(&r.upper_exclusive, b"foo");

        // Multiple ranges
        let ranges = vec![
            IntervalRange::from(("bar_4", "box")).into(),
            PointRange::from("bar_3").into(),
            PointRange::from("bar_2").into(),
            IntervalRange::from(("bar", "bar_")).into(),
            IntervalRange::from(("foo_5", "foo_50")).into(),
            IntervalRange::from(("foo", "foo_3")).into(),
        ];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage,
            ranges,
            scan_backward_in_range: true,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });

        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"bar_2");

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"bar_2");
        assert_eq!(&r.upper_exclusive, b"box");

        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"bar");

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"bar");
        assert_eq!(&r.upper_exclusive, b"bar_2");

        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"foo_2");
        assert_eq!(&block_on(scanner.next()).unwrap().unwrap().key(), b"foo");

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"foo");
        assert_eq!(&r.upper_exclusive, b"bar");

        assert_eq!(block_on(scanner.next()).unwrap(), None);

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"foo");
        assert_eq!(&r.upper_exclusive, b"foo");
    }

    #[test]
    fn test_scanned_range_forward2() {
        let storage = create_storage();
        // Filled interval range
        let ranges = vec![IntervalRange::from(("foo", "foo_8")).into()];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage: storage.clone(),
            ranges,
            scan_backward_in_range: false,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });

        // Only lower_inclusive is updated.
        assert_eq!(
            &block_on(scanner.next_opt(false)).unwrap().unwrap().key(),
            b"foo"
        );
        assert_eq!(&scanner.working_range_begin_key, b"foo");
        assert_eq!(&scanner.working_range_end_key, b"");

        // Upper_exclusive is updated.
        assert_eq!(
            &block_on(scanner.next_opt(true)).unwrap().unwrap().key(),
            b"foo_2"
        );
        assert_eq!(&scanner.working_range_begin_key, b"foo");
        assert_eq!(&scanner.working_range_end_key, b"foo_2\0");

        // Upper_exclusive is not updated.
        assert_eq!(
            &block_on(scanner.next_opt(false)).unwrap().unwrap().key(),
            b"foo_3"
        );
        assert_eq!(&scanner.working_range_begin_key, b"foo");
        assert_eq!(&scanner.working_range_end_key, b"foo_2\0");

        // Drained.
        assert_eq!(block_on(scanner.next_opt(false)).unwrap(), None);
        assert_eq!(&scanner.working_range_begin_key, b"foo");
        assert_eq!(&scanner.working_range_end_key, b"foo_8");

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"foo");
        assert_eq!(&r.upper_exclusive, b"foo_8");

        // Multiple ranges
        // TODO: caller should not pass in unordered ranges otherwise scanned ranges
        // would be unsound.
        let ranges = vec![
            IntervalRange::from(("foo", "foo_3")).into(),
            IntervalRange::from(("foo_5", "foo_50")).into(),
            IntervalRange::from(("bar", "bar_")).into(),
            PointRange::from("bar_2").into(),
            PointRange::from("bar_3").into(),
            IntervalRange::from(("bar_4", "box")).into(),
        ];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage,
            ranges,
            scan_backward_in_range: false,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });

        // Only lower_inclusive is updated.
        assert_eq!(
            &block_on(scanner.next_opt(false)).unwrap().unwrap().key(),
            b"foo"
        );
        assert_eq!(&scanner.working_range_begin_key, b"foo");
        assert_eq!(&scanner.working_range_end_key, b"");

        // Upper_exclusive is updated. Updated by scanned row.
        assert_eq!(
            &block_on(scanner.next_opt(true)).unwrap().unwrap().key(),
            b"foo_2"
        );
        assert_eq!(&scanner.working_range_begin_key, b"foo");
        assert_eq!(&scanner.working_range_end_key, b"foo_2\0");

        // Upper_exclusive is not updated.
        assert_eq!(
            &block_on(scanner.next_opt(false)).unwrap().unwrap().key(),
            b"bar"
        );
        assert_eq!(&scanner.working_range_begin_key, b"foo");
        assert_eq!(&scanner.working_range_end_key, b"foo_2\0");

        // Upper_exclusive is not updated.
        assert_eq!(
            &block_on(scanner.next_opt(false)).unwrap().unwrap().key(),
            b"bar_2"
        );
        assert_eq!(&scanner.working_range_begin_key, b"foo");
        assert_eq!(&scanner.working_range_end_key, b"foo_2\0");

        // Drain.
        assert_eq!(block_on(scanner.next_opt(false)).unwrap(), None);
        assert_eq!(&scanner.working_range_begin_key, b"foo");
        assert_eq!(&scanner.working_range_end_key, b"box");

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"foo");
        assert_eq!(&r.upper_exclusive, b"box");
    }

    #[test]
    fn test_scanned_range_backward2() {
        let storage = create_storage();
        // Filled interval range
        let ranges = vec![IntervalRange::from(("foo", "foo_8")).into()];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage: storage.clone(),
            ranges,
            scan_backward_in_range: true,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });

        // Only lower_inclusive is updated.
        assert_eq!(
            &block_on(scanner.next_opt(false)).unwrap().unwrap().key(),
            b"foo_3"
        );
        assert_eq!(&scanner.working_range_begin_key, b"foo_8");
        assert_eq!(&scanner.working_range_end_key, b"");

        // Upper_exclusive is updated.
        assert_eq!(
            &block_on(scanner.next_opt(true)).unwrap().unwrap().key(),
            b"foo_2"
        );
        assert_eq!(&scanner.working_range_begin_key, b"foo_8");
        assert_eq!(&scanner.working_range_end_key, b"foo_2");

        // Upper_exclusive is not updated.
        assert_eq!(
            &block_on(scanner.next_opt(false)).unwrap().unwrap().key(),
            b"foo"
        );
        assert_eq!(&scanner.working_range_begin_key, b"foo_8");
        assert_eq!(&scanner.working_range_end_key, b"foo_2");

        // Drained.
        assert_eq!(block_on(scanner.next_opt(false)).unwrap(), None);
        assert_eq!(&scanner.working_range_begin_key, b"foo_8");
        assert_eq!(&scanner.working_range_end_key, b"foo");

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"foo");
        assert_eq!(&r.upper_exclusive, b"foo_8");

        // Multiple ranges
        let ranges = vec![
            IntervalRange::from(("bar_4", "box")).into(),
            PointRange::from("bar_3").into(),
            PointRange::from("bar_2").into(),
            IntervalRange::from(("bar", "bar_")).into(),
            IntervalRange::from(("foo_5", "foo_50")).into(),
            IntervalRange::from(("foo", "foo_3")).into(),
        ];
        let mut scanner = RangesScanner::<_, ApiV1>::new(RangesScannerOptions {
            storage,
            ranges,
            scan_backward_in_range: true,
            is_key_only: false,
            load_commit_ts: false,
            reschedule: ScanRescheduleOptions::default(),
        });

        // Lower_inclusive is updated. Upper_exclusive is not update.
        assert_eq!(
            &block_on(scanner.next_opt(false)).unwrap().unwrap().key(),
            b"bar_2"
        );
        assert_eq!(&scanner.working_range_begin_key, b"box");
        assert_eq!(&scanner.working_range_end_key, b"");

        // Upper_exclusive is updated. Updated by scanned row.
        assert_eq!(
            &block_on(scanner.next_opt(true)).unwrap().unwrap().key(),
            b"bar"
        );
        assert_eq!(&scanner.working_range_begin_key, b"box");
        assert_eq!(&scanner.working_range_end_key, b"bar");

        // Upper_exclusive is not update.
        assert_eq!(
            &block_on(scanner.next_opt(false)).unwrap().unwrap().key(),
            b"foo_2"
        );
        assert_eq!(&scanner.working_range_begin_key, b"box");
        assert_eq!(&scanner.working_range_end_key, b"bar");

        // Upper_exclusive is not update.
        assert_eq!(
            &block_on(scanner.next_opt(false)).unwrap().unwrap().key(),
            b"foo"
        );
        assert_eq!(&scanner.working_range_begin_key, b"box");
        assert_eq!(&scanner.working_range_end_key, b"bar");

        // Drain.
        assert_eq!(block_on(scanner.next_opt(false)).unwrap(), None);
        assert_eq!(&scanner.working_range_begin_key, b"box");
        assert_eq!(&scanner.working_range_end_key, b"foo");

        let r = scanner.take_scanned_range();
        assert_eq!(&r.lower_inclusive, b"foo");
        assert_eq!(&r.upper_exclusive, b"box");
    }
}