From 47402694bb91700d4028d91fa4b769782c4af771 Mon Sep 17 00:00:00 2001
From: limbooverlambda <schakra1@gmail.com>
Date: Mon, 8 Jul 2024 13:12:50 -0700
Subject: [PATCH] address PR feedback

Signed-off-by: limbooverlambda <schakra1@gmail.com>
---
 src/raw/client.rs | 70 ++---------------------------------------------
 1 file changed, 2 insertions(+), 68 deletions(-)

diff --git a/src/raw/client.rs b/src/raw/client.rs
index 5afb213b..a466079e 100644
--- a/src/raw/client.rs
+++ b/src/raw/client.rs
@@ -774,7 +774,7 @@ impl<PdC: PdClient> Client<PdC> {
             let scan_args = ScanInnerArgs {
                 start_key: current_key.clone(),
                 range: range.clone(),
-                limit,
+                limit: current_limit,
                 key_only,
                 reverse,
                 backoff: backoff.clone(),
@@ -789,12 +789,7 @@ impl<PdC: PdClient> Client<PdC> {
                 current_limit -= kvs.len() as u32;
                 result.append(&mut kvs);
             }
-            if end_key
-                .as_ref()
-                .map(|ek| ek <= next_key.as_ref())
-                .unwrap_or(false)
-                || next_key.is_empty()
-            {
+            if end_key.clone().is_some_and(|ek| ek <= next_key) {
                 break;
             } else {
                 current_key = next_key;
@@ -928,13 +923,11 @@ struct ScanInnerArgs {
 #[cfg(test)]
 mod tests {
     use std::any::Any;
-    use std::sync::atomic::{AtomicBool, Ordering};
     use std::sync::Arc;
 
     use super::*;
     use crate::mock::MockKvClient;
     use crate::mock::MockPdClient;
-    use crate::proto::errorpb::EpochNotMatch;
     use crate::proto::kvrpcpb;
     use crate::Result;
 
@@ -1016,63 +1009,4 @@ mod tests {
         );
         Ok(())
     }
-
-    #[tokio::test]
-    async fn test_raw_scan_retryer() -> Result<()> {
-        let epoch_not_match_error = EpochNotMatch {
-            current_regions: vec![],
-        };
-        let region_error = errorpb::Error {
-            epoch_not_match: Some(epoch_not_match_error),
-            ..Default::default()
-        };
-        let flag = Arc::new(AtomicBool::new(true));
-        let tikv_flag = flag.clone();
-        let error_handler_flag = flag.clone();
-        let mock_tikv_client = MockKvClient::with_dispatch_hook(move |req: &dyn Any| {
-            if req.downcast_ref::<RawScanRequest>().is_some() {
-                let v = tikv_flag.clone().load(Ordering::Relaxed);
-                let resp = if v {
-                    RawScanResponse {
-                        region_error: Some(region_error.clone()),
-                        ..Default::default()
-                    }
-                } else {
-                    RawScanResponse {
-                        ..Default::default()
-                    }
-                };
-                Ok(Box::new(resp) as Box<dyn Any>)
-            } else {
-                unreachable!()
-            }
-        });
-        let mock_pd_client = MockPdClient::new(mock_tikv_client);
-        let client = Client {
-            rpc: Arc::new(mock_pd_client),
-            cf: Some(ColumnFamily::Default),
-            backoff: DEFAULT_REGION_BACKOFF,
-            atomic: false,
-            keyspace: Keyspace::Enable { keyspace_id: 0 },
-        };
-
-        let scan_args = ScanInnerArgs {
-            start_key: "k1".to_string().into(),
-            range: BoundRange::from("k1".to_owned().."k2".to_owned()),
-            limit: 10,
-            key_only: false,
-            reverse: false,
-            backoff: Backoff::no_backoff(),
-        };
-        let error_handler = |_, _, _| {
-            error_handler_flag.clone().store(false, Ordering::Relaxed);
-            Box::pin(async move {
-                let res: Result<bool> = Ok(true);
-                res
-            })
-        } as _;
-        client.retryable_scan(scan_args, error_handler).await?;
-        assert!(!error_handler_flag.load(Ordering::Relaxed));
-        Ok(())
-    }
 }