Codecov/fluctuation

xline/src/server/mod.rs

@@ -17,6 +17,8 @@ mod watch_server;
 /// Xline server
 mod xline_server;
 
-pub use self::command::{Command, KeyRange};
 pub(crate) use self::maintenance::MAINTENANCE_SNAPSHOT_CHUNK_SIZE;
-pub use self::xline_server::XlineServer;
+pub use self::{
+    command::{Command, KeyRange},
+    xline_server::XlineServer,
+};

xline/src/server/watch_server.rs

@@ -83,6 +83,8 @@ where
             header_gen,
         );
         let mut ticker = tokio::time::interval(watch_progress_notify_interval);
+        let stop_listener = stop_notify.listen();
+        tokio::pin!(stop_listener);
         loop {
             tokio::select! {
                 req = req_rx.next() => {
@@ -111,7 +113,9 @@ where
                 _ = ticker.tick() => {
                     watch_handle.handle_tick_progress().await;
                 }
-                _ = stop_notify.listen() => {
+                // To ensure that each iteration invokes the same `stop_listener` and keeps
+                // events losing due to the cancellation of `stop_listener` at bay.
+                _ = &mut stop_listener => {
                     break;
                 }
             }
@@ -666,4 +670,49 @@ mod test {
         handle.await.unwrap();
         Ok(())
     }
+
+    #[tokio::test]
+    async fn watch_task_should_be_terminated_when_response_tx_is_closed(
+    ) -> Result<(), Box<dyn std::error::Error>> {
+        let (req_tx, req_rx) = mpsc::channel(CHANNEL_SIZE);
+        let (res_tx, res_rx) = mpsc::channel(CHANNEL_SIZE);
+        let req_stream: ReceiverStream<Result<WatchRequest, tonic::Status>> =
+            ReceiverStream::new(req_rx);
+        let header_gen = Arc::new(HeaderGenerator::new(0, 0));
+        let mut mock_watcher = MockKvWatcherOps::new();
+        let _ = mock_watcher.expect_watch().times(1).return_const(());
+        let _ = mock_watcher.expect_cancel().times(1).return_const(());
+        let watcher = Arc::new(mock_watcher);
+        let next_id = Arc::new(WatchIdGenerator::new(1));
+        let handle = tokio::spawn(WatchServer::<DB>::task(
+            next_id,
+            Arc::clone(&watcher),
+            res_tx,
+            req_stream,
+            header_gen,
+            Duration::from_millis(100),
+        ));
+
+        req_tx
+            .send(Ok(WatchRequest {
+                request_union: Some(RequestUnion::CreateRequest(WatchCreateRequest {
+                    key: "foo".into(),
+                    progress_notify: true,
+                    watch_id: 1,
+                    ..Default::default()
+                })),
+            }))
+            .await?;
+
+        drop(res_rx);
+
+        req_tx
+            .send(Ok(WatchRequest {
+                request_union: Some(RequestUnion::ProgressRequest(WatchProgressRequest {})),
+            }))
+            .await?;
+
+        assert!(timeout(Duration::from_secs(10), handle).await.is_ok());
+        Ok(())
+    }
 }

xline/src/storage/kvwatcher.rs

@@ -576,10 +576,20 @@ mod test {
     async fn watch_should_not_lost_events() {
         let (store, db, kv_watcher) = init_empty_store();
         let mut map = BTreeMap::new();
+        let (event_tx, mut event_rx) = mpsc::channel(128);
+        let stop_notify = Arc::new(event_listener::Event::new());
+        kv_watcher.watch(
+            123,
+            KeyRange::new_one_key("foo"),
+            10,
+            vec![],
+            stop_notify,
+            event_tx,
+        );
+        sleep(Duration::from_micros(50)).await;
         let handle = tokio::spawn({
             let store = Arc::clone(&store);
             async move {
-                // let mut revision = 2;
                 for i in 0..100_u8 {
                     put(
                         store.as_ref(),
@@ -592,17 +602,6 @@ mod test {
                 }
             }
         });
-        sleep(Duration::from_micros(500)).await;
-        let (event_tx, mut event_rx) = mpsc::channel(128);
-        let stop_notify = Arc::new(event_listener::Event::new());
-        kv_watcher.watch(
-            123,
-            KeyRange::new_one_key("foo"),
-            1,
-            vec![],
-            stop_notify,
-            event_tx,
-        );
 
         'outer: while let Some(event_batch) = timeout(Duration::from_secs(3), event_rx.recv())
             .await
@@ -618,8 +617,8 @@ mod test {
                 }
             }
         }
-
-        assert_eq!(map.len(), 100);
+        // The length of the map should be 1 + total_val_number - (start_rev - 1) = 1 + 100 - (10 - 1) = 92
+        assert_eq!(map.len(), 92);
         for (k, count) in map {
             assert_eq!(count, 1, "key {k} should be notified once");
         }