fix listener on follower if region changed

components/br-stream/src/endpoint.rs

@@ -227,7 +227,6 @@ where
             self.router.clone(),
             self.regions.clone(),
             self.range_router.clone(),
-            self.store_id,
         )
     }
 
@@ -402,13 +401,16 @@ where
                 self.resolvers.as_ref().remove(&region.id);
             }
             ObserveOp::RefreshResolver { region } => {
-                self.observer.subs.deregister_region(region.id);
+                let canceled = self.observer.subs.deregister_region(region.id);
                 self.resolvers.as_ref().remove(&region.id);
-                if let Err(e) = self.observe_over(&region) {
-                    e.report(format!(
-                        "register region {} to raftstore when refreshing",
-                        region.get_id()
-                    ));
+
+                if canceled {
+                    if let Err(e) = self.observe_over(&region) {
+                        e.report(format!(
+                            "register region {} to raftstore when refreshing",
+                            region.get_id()
+                        ));
+                    }
                 }
             }
         }

components/br-stream/src/event_loader.rs

@@ -118,7 +118,6 @@ pub struct InitialDataLoader<E, R, RT> {
     // Note: maybe we can make it an abstract thing like `EventSink` with
     //       method `async (KvEvent) -> Result<()>`?
     sink: Router,
-    store_id: u64,
 
     _engine: PhantomData<E>,
 }
@@ -129,12 +128,11 @@ where
     R: RegionInfoProvider + Clone + 'static,
     RT: RaftStoreRouter<E>,
 {
-    pub fn new(router: RT, regions: R, sink: Router, store_id: u64) -> Self {
+    pub fn new(router: RT, regions: R, sink: Router) -> Self {
         Self {
             router,
             regions,
             sink,
-            store_id,
             _engine: PhantomData,
         }
     }

components/br-stream/src/observer.rs

@@ -68,9 +68,16 @@ impl BackupStreamObserver {
 
     /// Test whether a region should be observed by the observer.
     fn should_register_region(&self, region: &Region) -> bool {
+        // If the end key is empty, it actually meant infinity.
+        // However, this way is a little hacky, maybe we'd better make a
+        // `RangesBound<R>` version for `is_overlapping`.
+        let mut end_key = region.get_end_key();
+        if end_key.is_empty() {
+            end_key = &[0xffu8; 32];
+        }
         self.ranges
             .rl()
-            .is_overlapping((region.get_start_key(), region.get_end_key()))
+            .is_overlapping((region.get_start_key(), end_key))
     }
 }
 
@@ -89,20 +96,24 @@ impl SubscriptionTracer {
         }
     }
 
-    pub fn deregister_region(&self, region_id: u64) {
+    /// try to mark a region no longer be tracked by this observer.
+    /// returns whether success (it failed if the region hasn't been observed when calling this.)
+    pub fn deregister_region(&self, region_id: u64) -> bool {
         match self.0.remove(&region_id) {
             Some(o) => {
                 o.1.stop_observing();
                 info!("stop listen stream from store"; "observer" => ?o.1, "region_id"=> %region_id);
+                true
             }
             None => {
                 warn!("trying to deregister region not registered"; "region_id" => %region_id);
+                false
             }
         }
     }
 
     /// check whether the region_id should be observed by this observer.
-    pub fn should_observe(&self, region_id: u64) -> bool {
+    pub fn is_observing(&self, region_id: u64) -> bool {
         let mut exists = false;
 
         // The region traced, check it whether is still be observing,
@@ -147,7 +158,7 @@ impl<E: KvEngine> CmdObserver<E> for BackupStreamObserver {
                 !cb.is_empty()
                     && cb.level == ObserveLevel::All
                     // Once the observe has been canceled by outside things, we should be able to stop.
-                    && self.subs.should_observe(cb.region_id)
+                    && self.subs.is_observing(cb.region_id)
             })
             .cloned()
             .collect();
@@ -182,10 +193,22 @@ impl RegionChangeObserver for BackupStreamObserver {
         &self,
         ctx: &mut ObserverContext<'_>,
         event: RegionChangeEvent,
-        _role: StateRole,
+        role: StateRole,
     ) {
+        if !self.subs.is_observing(ctx.region().get_id()) {
+            return;
+        }
+        if role != StateRole::Leader {
+            try_send!(
+                self.scheduler,
+                Task::ModifyObserve(ObserveOp::Stop {
+                    region: ctx.region().clone(),
+                })
+            );
+            return;
+        }
         match event {
-            RegionChangeEvent::Destroy if self.subs.should_observe(ctx.region().get_id()) => {
+            RegionChangeEvent::Destroy => {
                 try_send!(
                     self.scheduler,
                     Task::ModifyObserve(ObserveOp::Stop {
@@ -229,6 +252,16 @@ mod tests {
 
     use super::BackupStreamObserver;
 
+    macro_rules! assert_let {
+        (let $p:pat = $e:expr; $cc:tt) => {
+            if let $p = $e {
+                $cc
+            } else {
+                panic!("{} doesn't matches {}", stringify!($e), stringify!($p))
+            }
+        };
+    }
+
     fn fake_region(id: u64, start: &[u8], end: &[u8]) -> Region {
         let mut r = Region::new();
         r.set_id(id);
@@ -255,9 +288,9 @@ mod tests {
         } else {
             panic!("unexpected message received: it is {}", task);
         }
-        assert!(o.subs.should_observe(42));
+        assert!(o.subs.is_observing(42));
         handle.stop_observing();
-        assert!(!o.subs.should_observe(42));
+        assert!(!o.subs.is_observing(42));
     }
 
     #[test]
@@ -274,11 +307,9 @@ mod tests {
         o.register_region(&r);
         let task = rx.recv_timeout(Duration::from_secs(0)).unwrap().unwrap();
         let handle = ObserveHandle::new();
-        if let Task::ModifyObserve(ObserveOp::Start { region }) = task {
+        assert_let!(let Task::ModifyObserve(ObserveOp::Start { region }) = task; {
             o.subs.register_region(region.get_id(), handle.clone())
-        } else {
-            panic!("unexpected message received: it is {}", task);
-        }
+        });
 
         // Test events with key in the range can be observed.
         let observe_info = CmdObserveInfo::from_handle(handle.clone(), ObserveHandle::new());
@@ -287,13 +318,11 @@ mod tests {
         let mut cmd_batches = vec![cb];
         o.on_flush_applied_cmd_batch(ObserveLevel::All, &mut cmd_batches, &mock_engine);
         let task = rx.recv_timeout(Duration::from_secs(0)).unwrap().unwrap();
-        if let Task::BatchEvent(batches) = task {
+        assert_let!(let Task::BatchEvent(batches) = task; {
             assert!(batches.len() == 1);
             assert!(batches[0].region_id == 42);
             assert!(batches[0].cdc_id == handle.id);
-        } else {
-            panic!("unexpected message received: it is {}", task);
-        }
+        });
 
         // Test event from other region should not be send.
         let observe_info = CmdObserveInfo::from_handle(ObserveHandle::new(), ObserveHandle::new());
@@ -311,21 +340,22 @@ mod tests {
         o.on_role_change(&mut ctx, StateRole::Leader);
         let task = rx.recv_timeout(Duration::from_millis(20));
         assert!(task.is_err(), "it is {:?}", task);
-        assert!(!o.subs.should_observe(43));
+        assert!(!o.subs.is_observing(43));
 
-        // Test region out of range won't be added to observe list.
+        // Test newly created region out of range won't be added to observe list.
         let mut ctx = ObserverContext::new(&r);
         o.on_region_changed(&mut ctx, RegionChangeEvent::Create, StateRole::Leader);
         let task = rx.recv_timeout(Duration::from_millis(20));
         assert!(task.is_err(), "it is {:?}", task);
-        assert!(!o.subs.should_observe(43));
+        assert!(!o.subs.is_observing(43));
 
         // Test give up subscripting when become follower.
         let r = fake_region(42, b"0008", b"0009");
         let mut ctx = ObserverContext::new(&r);
         o.on_role_change(&mut ctx, StateRole::Follower);
         let task = rx.recv_timeout(Duration::from_millis(20));
-        assert!(task.is_err(), "it is {:?}", task);
-        assert!(!o.subs.should_observe(42));
+        assert_let!(let Ok(Some(Task::ModifyObserve(ObserveOp::Stop { region, .. }))) = task; {
+            assert_eq!(region.id, 42);
+        });
     }
 }

components/br-stream/src/utils.rs

@@ -169,6 +169,18 @@ impl<K: Ord, V: Default> SegmentMap<K, V> {
     }
 }
 
+struct RangeToExclusiveRef<'a, T: ?Sized>(&'a T);
+
+impl<'a, T: ?Sized> RangeBounds<T> for RangeToExclusiveRef<'a, T> {
+    fn start_bound(&self) -> Bound<&T> {
+        Bound::Unbounded
+    }
+
+    fn end_bound(&self) -> Bound<&T> {
+        Bound::Excluded(self.0)
+    }
+}
+
 impl<K: Ord, V> SegmentMap<K, V> {
     /// Like `add`, but insert a value associated to the key.
     pub fn insert(&mut self, (start, end): (K, K), value: V) -> bool {
@@ -222,10 +234,32 @@ impl<K: Ord, V> SegmentMap<K, V> {
         K: Borrow<R>,
         R: Ord + ?Sized,
     {
-        self.get_interval_by_point(range.0).is_some()
-            || self
-                .get_interval_by_point(range.1)
-                .map_or(false, |rng| <K as Borrow<R>>::borrow(rng.0) != range.1)
+        // o: The Start Key.
+        // e: The End Key.
+        // +: The Boundary of Candidate Range.
+        // |------+-s----+----e----|
+        // Firstly, we check whether the start point is in some range.
+        // if true, it must be overlapping.
+        let overlap_with_start = self.get_interval_by_point(range.0).is_some();
+        // |--s----+-----+----e----|
+        // Otherwise, the possibility of being overlapping would be there are some sub range
+        // of the queried range...
+        // |--s----+----e----+-----|
+        // ...Or the end key is contained by some Range.
+        // For faster query, we merged the two cases together.
+        let covered_by_the_range = self
+             .0
+             // When querying possibility of overlapping by end key,
+             // we don't want the range [end key, ...) become a candidate.
+             // (which is impossible to overlapping with the range)
+             .range(RangeToExclusiveRef(range.1))
+             .next_back()
+             .filter(|(start, end)| {
+                 <K as Borrow<R>>::borrow(&end.range_end) > range.1
+                     || <K as Borrow<R>>::borrow(start) > range.0
+             })
+             .is_some();
+        overlap_with_start || covered_by_the_range
     }
 
     pub fn get_inner(&mut self) -> &mut BTreeMap<K, SegmentValue<K, V>> {
@@ -298,5 +332,6 @@ mod test {
         assert!(!tree.is_overlapping((&8, &42)));
         assert!(!tree.is_overlapping((&9, &10)));
         assert!(tree.is_overlapping((&2, &10)));
+        assert!(tree.is_overlapping((&0, &9999999)));
     }
 }