Merge branch 'master' into fix-update-regioncache

Author: sunxiaoguang

src/pd/client.rs

@@ -70,7 +70,7 @@ pub trait PdClient: Send + Sync + 'static {
     fn group_keys_by_region<K, K2>(
         self: Arc<Self>,
         keys: impl Iterator<Item = K> + Send + Sync + 'static,
-    ) -> BoxStream<'static, Result<(RegionId, Vec<K2>)>>
+    ) -> BoxStream<'static, Result<(RegionWithLeader, Vec<K2>)>>
     where
         K: AsRef<Key> + Into<K2> + Send + Sync + 'static,
         K2: Send + Sync + 'static,
@@ -81,15 +81,14 @@ pub trait PdClient: Send + Sync + 'static {
             async move {
                 if let Some(key) = keys.next() {
                     let region = this.region_for_key(key.as_ref()).await?;
-                    let id = region.id();
                     let mut grouped = vec![key.into()];
                     while let Some(key) = keys.peek() {
                         if !region.contains(key.as_ref()) {
                             break;
                         }
                         grouped.push(keys.next().unwrap().into());
                     }
-                    Ok(Some((keys, (id, grouped))))
+                    Ok(Some((keys, (region, grouped))))
                 } else {
                     Ok(None)
                 }
@@ -133,7 +132,7 @@ pub trait PdClient: Send + Sync + 'static {
     fn group_ranges_by_region(
         self: Arc<Self>,
         mut ranges: Vec<kvrpcpb::KeyRange>,
-    ) -> BoxStream<'static, Result<(RegionId, Vec<kvrpcpb::KeyRange>)>> {
+    ) -> BoxStream<'static, Result<(RegionWithLeader, Vec<kvrpcpb::KeyRange>)>> {
         ranges.reverse();
         stream_fn(Some(ranges), move |ranges| {
             let this = self.clone();
@@ -147,7 +146,6 @@ pub trait PdClient: Send + Sync + 'static {
                     let start_key: Key = range.start_key.clone().into();
                     let end_key: Key = range.end_key.clone().into();
                     let region = this.region_for_key(&start_key).await?;
-                    let id = region.id();
                     let region_start = region.start_key();
                     let region_end = region.end_key();
                     let mut grouped = vec![];
@@ -160,7 +158,7 @@ pub trait PdClient: Send + Sync + 'static {
                             start_key: region_end.into(),
                             end_key: end_key.into(),
                         });
-                        return Ok(Some((Some(ranges), (id, grouped))));
+                        return Ok(Some((Some(ranges), (region, grouped))));
                     }
                     grouped.push(range);
 
@@ -180,11 +178,11 @@ pub trait PdClient: Send + Sync + 'static {
                                 start_key: region_end.into(),
                                 end_key: end_key.into(),
                             });
-                            return Ok(Some((Some(ranges), (id, grouped))));
+                            return Ok(Some((Some(ranges), (region, grouped))));
                         }
                         grouped.push(range);
                     }
-                    Ok(Some((Some(ranges), (id, grouped))))
+                    Ok(Some((Some(ranges), (region, grouped))))
                 } else {
                     Ok(None)
                 }
@@ -452,7 +450,7 @@ pub mod test {
         let ranges3 = stream.next().unwrap().unwrap();
         let ranges4 = stream.next().unwrap().unwrap();
 
-        assert_eq!(ranges1.0, 1);
+        assert_eq!(ranges1.0.id(), 1);
         assert_eq!(
             ranges1.1,
             vec![
@@ -466,23 +464,23 @@ pub mod test {
                 }
             ]
         );
-        assert_eq!(ranges2.0, 2);
+        assert_eq!(ranges2.0.id(), 2);
         assert_eq!(
             ranges2.1,
             vec![kvrpcpb::KeyRange {
                 start_key: k_split.clone(),
                 end_key: k3
             }]
         );
-        assert_eq!(ranges3.0, 1);
+        assert_eq!(ranges3.0.id(), 1);
         assert_eq!(
             ranges3.1,
             vec![kvrpcpb::KeyRange {
                 start_key: k2,
                 end_key: k_split.clone()
             }]
         );
-        assert_eq!(ranges4.0, 2);
+        assert_eq!(ranges4.0.id(), 2);
         assert_eq!(
             ranges4.1,
             vec![kvrpcpb::KeyRange {

src/store.rs

@@ -39,10 +39,10 @@ where
     pd_client
         .clone()
         .group_keys_by_region(key_data)
-        .and_then(move |(region_id, key)| {
+        .and_then(move |(region, key)| {
             pd_client
                 .clone()
-                .store_for_id(region_id)
+                .map_region_to_store(region)
                 .map_ok(move |store| (key, store))
         })
         .boxed()
@@ -106,10 +106,10 @@ pub fn store_stream_for_ranges<PdC: PdClient>(
     pd_client
         .clone()
         .group_ranges_by_region(ranges)
-        .and_then(move |(region_id, range)| {
+        .and_then(move |(region, range)| {
             pd_client
                 .clone()
-                .store_for_id(region_id)
+                .map_region_to_store(region)
                 .map_ok(move |store| (range, store))
         })
         .boxed()