Repository: tikv/tikv (language: rust)

Issue:
log-backup: add a RPC endpoint to force flush tasks
## Development Task
As title, this would be useful for testing and optimizing RPO when rolling restarting.


Discussion / hints:

Fix patch:
diff --git a/components/backup-stream/src/checkpoint_manager.rs b/components/backup-stream/src/checkpoint_manager.rs
index 602b8b43bec..63f7a472bf7 100644
--- a/components/backup-stream/src/checkpoint_manager.rs
+++ b/components/backup-stream/src/checkpoint_manager.rs
@@ -896,8 +896,8 @@ pub mod tests {
         let r = flush_observer.after(&task, rts).await;
         assert_eq!(r.is_ok(), true);
 
-        let serivce_id = format!("backup-stream-{}-{}", task, store_id);
-        let r = pd_cli.get_service_safe_point(serivce_id).unwrap();
+        let service_id = format!("backup-stream-{}-{}", task, store_id);
+        let r = pd_cli.get_service_safe_point(service_id).unwrap();
         assert_eq!(r.into_inner(), rts - 1);
     }
 }
diff --git a/components/backup-stream/src/endpoint.rs b/components/backup-stream/src/endpoint.rs
index 4dd35e926b0..af8f3a99de8 100644
--- a/components/backup-stream/src/endpoint.rs
+++ b/components/backup-stream/src/endpoint.rs
@@ -5,6 +5,7 @@ use std::{
     collections::HashSet,
     fmt,
     marker::PhantomData,
+    mem::ManuallyDrop,
     sync::{Arc, Mutex},
     time::Duration,
 };
@@ -64,7 +65,7 @@ use crate::{
     metadata::{store::MetaStore, MetadataClient, MetadataEvent, StreamTask},
     metrics::{self, TaskStatus},
     observer::BackupStreamObserver,
-    router::{self, ApplyEvents, FlushContext, Router, TaskSelector},
+    router::{self, ApplyEvents, FlushContext, Router, TaskSelector, TaskSelectorRef},
     subscription_manager::{RegionSubscriptionManager, ResolvedRegions},
     subscription_track::{Ref, RefMut, ResolveResult, SubscriptionTracer},
     try_send,
@@ -92,7 +93,7 @@ pub struct Endpoint<S, R, E: KvEngine, PDC> {
     // Note: some of the fields are public so test cases are able to access them.
     pub range_router: Router,
     observer: BackupStreamObserver,
-    pool: Runtime,
+    pool: ManuallyDrop<Runtime>,
     region_operator: Sender<ObserveOp>,
     failover_time: Option<Instant>,
     config: BackupStreamConfig,
@@ -106,6 +107,13 @@ pub struct Endpoint<S, R, E: KvEngine, PDC> {
     pub initial_scan_semaphore: Arc<Semaphore>,
 }
 
+impl<S, R, E: KvEngine, PDC> Drop for Endpoint<S, R, E, PDC> {
+    fn drop(&mut self) {
+        // SAFETY: won't access thread pool after dropping.
+        unsafe { ManuallyDrop::take(&mut self.pool).shutdown_background() }
+    }
+}
+
 impl<S, R, E, PDC> Endpoint<S, R, E, PDC>
 where
     R: RegionInfoProvider + 'static + Clone,
@@ -190,7 +198,7 @@ where
             range_router,
             scheduler,
             observer,
-            pool,
+            pool: ManuallyDrop::new(pool),
             store_id,
             regions: accessor,
             engine: PhantomData,
@@ -857,20 +865,36 @@ where
         }
     }
 
-    pub fn on_force_flush(&self, task: String) {
+    pub fn on_force_flush(&self, task: TaskSelectorRef<'_>, sender: Sender<FlushResult>) {
         self.pool.block_on(async move {
-            let handler_res = self.range_router.get_task_handler(&task);
-            // This should only happen in testing, it would be to unwrap...
-            let _ = handler_res.unwrap().set_flushing_status_cas(false, true);
-            let mts = self.prepare_min_ts().await;
-            let sched = self.scheduler.clone();
-            self.region_op(ObserveOp::ResolveRegions {
-                callback: Box::new(move |res| {
-                    try_send!(sched, Task::ExecFlush(task, res));
-                }),
-                min_ts: mts,
-            })
-            .await;
+            info!("Triggering force flush."; "selector" => ?task);
+            let handlers = self.range_router.select_task_handler(task);
+            for hnd in handlers {
+                let mts = self.prepare_min_ts().await;
+                let sched = self.scheduler.clone();
+                let sender = sender.clone();
+                match hnd.set_flushing_status_cas(false, true) {
+                    Ok(_) => {
+                        self.region_op(ObserveOp::ResolveRegions {
+                            callback: Box::new(move |res| {
+                                try_send!(
+                                    sched,
+                                    Task::ExecFlush(hnd.task.info.name.to_owned(), res, sender)
+                                );
+                            }),
+                            min_ts: mts,
+                        })
+                        .await;
+                    }
+                    Err(_) => {
+                        let res = FlushResult {
+                            task: hnd.task.info.name.to_owned(),
+                            error: Some(Error::Other(box_err!("task is flushing"))),
+                        };
+                        let _ = sender.send(res).await;
+                    }
+                }
+            }
         });
     }
 
@@ -879,9 +903,10 @@ where
             let mts = self.prepare_min_ts().await;
             let sched = self.scheduler.clone();
             info!("min_ts prepared for flushing"; "min_ts" => %mts);
+            let (tx, _) = tokio::sync::mpsc::channel(1);
             self.region_op(ObserveOp::ResolveRegions {
                 callback: Box::new(move |res| {
-                    try_send!(sched, Task::ExecFlush(task, res));
+                    try_send!(sched, Task::ExecFlush(task, res, tx));
                 }),
                 min_ts: mts,
             })
@@ -889,14 +914,20 @@ where
         })
     }
 
-    fn on_exec_flush(&mut self, task: String, resolved: ResolvedRegions) {
+    fn on_exec_flush(&mut self, task: String, resolved: ResolvedRegions, cb: Sender<FlushResult>) {
         self.checkpoint_mgr.freeze();
-        self.pool
-            .spawn(root!("flush"; self.do_flush(task, resolved).map(|r| {
-                if let Err(err) = r {
-                    err.report("during updating flush status")
-                }
-            })));
+        let fut = self.do_flush(task.clone(), resolved);
+        self.pool.spawn(root!("flush"; async move {
+            let res = fut.await;
+            if let Err(ref err) = &res {
+                err.report("during updating flush status")
+            }
+            // If nobody waits us, it is no need to construct the result.
+            if !cb.is_closed() {
+                let flush_res = FlushResult { task, error: res.err() };
+                let _ = cb.send(flush_res).await;
+            }
+        }));
     }
 
     fn update_global_checkpoint(&self, task: String) -> future![()] {
@@ -1024,7 +1055,7 @@ where
             Task::BatchEvent(events) => self.do_backup(events),
             Task::Flush(task) => self.on_flush(task),
             Task::ModifyObserve(op) => self.on_modify_observe(op),
-            Task::ForceFlush(task) => self.on_force_flush(task),
+            Task::ForceFlush(sel, cb) => self.on_force_flush(sel.reference(), cb),
             Task::FatalError(task, err) => self.on_fatal_error(task, err),
             Task::ChangeConfig(cfg) => {
                 self.on_update_change_config(cfg);
@@ -1041,7 +1072,7 @@ where
                 }
             }
             Task::MarkFailover(t) => self.failover_time = Some(t),
-            Task::ExecFlush(task, min_ts) => self.on_exec_flush(task, min_ts),
+            Task::ExecFlush(task, min_ts, cb) => self.on_exec_flush(task, min_ts, cb),
             Task::RegionCheckpointsOp(s) => self.handle_region_checkpoints_op(s),
             Task::UpdateGlobalCheckpoint(task) => self.on_update_global_checkpoint(task),
         }
@@ -1239,6 +1270,12 @@ impl fmt::Debug for RegionCheckpointOperation {
     }
 }
 
+#[derive(Debug)]
+pub struct FlushResult {
+    pub task: String,
+    pub error: Option<Error>,
+}
+
 pub enum Task {
     WatchTask(TaskOp),
     BatchEvent(Vec<CmdBatch>),
@@ -1246,19 +1283,16 @@ pub enum Task {
     /// Change the observe status of some region.
     ModifyObserve(ObserveOp),
     /// Convert status of some task into `flushing` and do flush then.
-    ForceFlush(String),
+    ForceFlush(TaskSelector, Sender<FlushResult>),
     /// FatalError pauses the task and set the error.
     FatalError(TaskSelector, Box<Error>),
-    /// Run the callback when see this message. Only for test usage.
-    /// NOTE: Those messages for testing are not guarded by `#[cfg(test)]` for
-    /// now, because the integration test would not enable test config when
-    /// compiling (why?)
+    /// Run the callback when see this message.
     Sync(
         // Run the closure if ...
         Box<dyn FnOnce() + Send>,
         // This returns `true`.
         // The argument should be `self`, but there are too many generic argument for `self`...
-        // So let the caller in test cases downcast this to the type they need manually...
+        // So let the caller downcast this to the type they need manually...
         Box<dyn FnMut(&mut dyn Any) -> bool + Send>,
     ),
     /// Mark the store as a failover store.
@@ -1271,7 +1305,7 @@ pub enum Task {
     Flush(String),
     /// Execute the flush with the calculated resolved result.
     /// This is an internal command only issued by the `Flush` task.
-    ExecFlush(String, ResolvedRegions),
+    ExecFlush(String, ResolvedRegions, Sender<FlushResult>),
     /// The command for getting region checkpoints.
     RegionCheckpointsOp(RegionCheckpointOperation),
     /// update global-checkpoint-ts to storage.
@@ -1377,7 +1411,7 @@ impl fmt::Debug for Task {
             Self::ChangeConfig(arg0) => f.debug_tuple("ChangeConfig").field(arg0).finish(),
             Self::Flush(arg0) => f.debug_tuple("Flush").field(arg0).finish(),
             Self::ModifyObserve(op) => f.debug_tuple("ModifyObserve").field(op).finish(),
-            Self::ForceFlush(arg0) => f.debug_tuple("ForceFlush").field(arg0).finish(),
+            Self::ForceFlush(sel, _) => f.debug_tuple("ForceFlush").field(sel).finish(),
             Self::FatalError(task, err) => {
                 f.debug_tuple("FatalError").field(task).field(err).finish()
             }
@@ -1386,7 +1420,7 @@ impl fmt::Debug for Task {
                 .debug_tuple("MarkFailover")
                 .field(&format_args!("{:?} ago", t.saturating_elapsed()))
                 .finish(),
-            Self::ExecFlush(arg0, arg1) => f
+            Self::ExecFlush(arg0, arg1, _) => f
                 .debug_tuple("ExecFlush")
                 .field(arg0)
                 .field(&arg1.global_checkpoint())
diff --git a/components/backup-stream/src/router.rs b/components/backup-stream/src/router.rs
index ce16776d9be..ed307871a46 100644
--- a/components/backup-stream/src/router.rs
+++ b/components/backup-stream/src/router.rs
@@ -619,6 +619,26 @@ impl RouterInner {
         r.get_value_by_point(key).cloned()
     }
 
+    pub fn select_task_handler(
+        &self,
+        selector: TaskSelectorRef<'_>,
+    ) -> impl Iterator<Item = Arc<StreamTaskHandler>> {
+        self.tasks
+            .iter()
+            .filter(|entry| {
+                let (name, info) = entry.pair();
+                selector.matches(
+                    name.as_str(),
+                    info.ranges
+                        .iter()
+                        .map(|(s, e)| (s.as_slice(), e.as_slice())),
+                )
+            })
+            .map(|entry| entry.value().clone())
+            .collect::<Vec<_>>()
+            .into_iter()
+    }
+
     pub fn select_task(&self, selector: TaskSelectorRef<'_>) -> Vec<String> {
         self.tasks
             .iter()
diff --git a/components/backup-stream/src/service.rs b/components/backup-stream/src/service.rs
index e639f44a731..3e1db6f1d0d 100644
--- a/components/backup-stream/src/service.rs
+++ b/components/backup-stream/src/service.rs
@@ -2,13 +2,15 @@
 
 use std::collections::HashSet;
 
-use grpcio::RpcContext;
+use futures::future::FutureExt;
+use grpcio::{RpcContext, RpcStatus, RpcStatusCode};
 use kvproto::{logbackuppb::*, metapb::Region};
-use tikv_util::{warn, worker::Scheduler};
+use tikv_util::{info, warn, worker::Scheduler};
 
 use crate::{
     checkpoint_manager::{GetCheckpointResult, RegionIdWithVersion},
     endpoint::{RegionCheckpointOperation, RegionSet},
+    router::TaskSelector,
     try_send, Task,
 };
 
@@ -40,6 +42,51 @@ impl From<RegionIdWithVersion> for RegionIdentity {
 }
 
 impl LogBackup for BackupStreamGrpcService {
+    fn flush_now(
+        &mut self,
+        ctx: grpcio::RpcContext<'_>,
+        _req: FlushNowRequest,
+        sink: grpcio::UnarySink<FlushNowResponse>,
+    ) {
+        info!("Client requests force flush."; "cli" => %ctx.peer());
+        let mut resp = FlushNowResponse::new();
+        let (tx, mut rx) = tokio::sync::mpsc::channel(1);
+        let task = Task::ForceFlush(TaskSelector::All, tx);
+        if let Err(err) = self.endpoint.schedule(task) {
+            ctx.spawn(
+                sink.fail(RpcStatus::with_message(
+                    RpcStatusCode::INTERNAL,
+                    format!(
+                        "failed to schedule the command, maybe busy or shutting down: {}",
+                        err
+                    ),
+                ))
+                .map(|res| {
+                    if let Err(err) = res {
+                        warn!("flush_now: failed to send an error response to client"; "err" => %err)
+                    }
+                }),
+            );
+            return;
+        };
+
+        ctx.spawn(async move {
+            while let Some(item) = rx.recv().await {
+                let mut res = FlushResult::new();
+                res.set_success(item.error.is_none());
+                if let Some(err) = item.error {
+                    res.set_error_message(err.to_string());
+                }
+                res.set_task_name(item.task);
+                resp.results.push(res);
+            }
+
+            if let Err(err) = sink.success(resp.clone()).await {
+                warn!("flush_now: failed to send success response to client"; "err" => %err, "resp" => ?resp);
+            }
+        })
+    }
+
     fn get_last_flush_ts_of_region(
         &mut self,
         _ctx: RpcContext<'_>,


Test patch:
diff --git a/components/backup-stream/tests/integration/mod.rs b/components/backup-stream/tests/integration/mod.rs
index ed161c68047..8b33a46558f 100644
--- a/components/backup-stream/tests/integration/mod.rs
+++ b/components/backup-stream/tests/integration/mod.rs
@@ -190,12 +190,10 @@ mod all {
             suite.write_records(0, 128, 1).await;
             let ts = suite.just_async_commit_prewrite(256, 1);
             suite.write_records(258, 128, 1).await;
-            suite.force_flush_files("test_async_commit");
-            std::thread::sleep(Duration::from_secs(4));
+            suite.force_flush_files_and_wait("test_async_commit").await;
             assert_eq!(suite.global_checkpoint(), 256);
             suite.just_commit_a_key(make_record_key(1, 256), TimeStamp::new(256), ts);
-            suite.force_flush_files("test_async_commit");
-            suite.wait_for_flush();
+            suite.force_flush_files_and_wait("test_async_commit").await;
             let cp = suite.global_checkpoint();
             assert!(cp > 256, "it is {:?}", cp);
         });
@@ -425,22 +423,22 @@ mod all {
         let leader = suite.cluster.leader_of_region(1).unwrap();
         suite.must_shuffle_leader(1);
         let round2 = run_async_test(suite.write_records(256, 128, 1));
+        let (tx, mut rx) = tokio::sync::mpsc::channel(1);
         suite
             .endpoints
             .get(&leader.store_id)
             .unwrap()
             .scheduler()
-            .schedule(Task::ForceFlush("r".to_owned()))
+            .schedule(Task::ForceFlush(TaskSelector::All, tx))
             .unwrap();
-        suite.sync();
-        std::thread::sleep(Duration::from_secs(2));
+        while rx.blocking_recv().is_some() {}
+
         suite.check_for_write_records(
             suite.flushed_files.path(),
             round1.iter().map(|x| x.as_slice()),
         );
         assert!(suite.global_checkpoint() > 256);
-        suite.force_flush_files("r");
-        suite.wait_for_flush();
+        run_async_test(suite.force_flush_files_and_wait("r"));
         assert!(suite.global_checkpoint() > 512);
         suite.check_for_write_records(
             suite.flushed_files.path(),
@@ -495,12 +493,11 @@ mod all {
 
     #[test]
     fn update_config() {
-        let suite = SuiteBuilder::new_named("network_partition")
-            .nodes(1)
-            .build();
+        let suite = SuiteBuilder::new_named("update_config").nodes(1).build();
         let mut basic_config = BackupStreamConfig::default();
         basic_config.initial_scan_concurrency = 4;
         suite.run(|| Task::ChangeConfig(basic_config.clone()));
+        suite.sync();
         suite.wait_with(|e| {
             assert_eq!(e.initial_scan_semaphore.available_permits(), 4,);
             true
@@ -513,4 +510,36 @@ mod all {
             true
         });
     }
+
+    #[test]
+    fn force_flush() {
+        let mut suite = SuiteBuilder::new_named("force_flush").nodes(1).build();
+        suite.must_register_task(1, "force_flush");
+        let recs = run_async_test(suite.write_records(0, 128, 1));
+        let mut strm = suite.flush_stream(true);
+        let tso = suite.tso();
+
+        suite.for_each_log_backup_cli(|_id, c| {
+            let res = c.flush_now(Default::default()).unwrap();
+            assert_eq!(res.results.len(), 1);
+            assert!(res.results[0].error_message.is_empty(), "{:?}", res);
+            assert!(res.results[0].success, "{:?}", res);
+        });
+
+        let Some((_, resp)) = run_async_test(strm.next()) else {
+            panic!("subscribe stream close early")
+        };
+        assert_eq!(resp.events.len(), 1, "{:?}", resp.events);
+        assert!(
+            resp.events[0].checkpoint > tso.into_inner(),
+            "{:?}, {}",
+            resp.events[0],
+            tso
+        );
+
+        suite.check_for_write_records(
+            suite.flushed_files.path(),
+            recs.iter().map(|v| v.as_slice()),
+        )
+    }
 }
diff --git a/components/backup-stream/tests/suite.rs b/components/backup-stream/tests/suite.rs
index e3254f16acc..a828ddee83c 100644
--- a/components/backup-stream/tests/suite.rs
+++ b/components/backup-stream/tests/suite.rs
@@ -595,10 +595,33 @@ impl Suite {
         ts
     }
 
+    pub fn for_each_log_backup_cli(&self, mut cb: impl FnMut(u64, &LogBackupClient)) {
+        for (k, v) in self.log_backup_cli.iter() {
+            cb(*k, v)
+        }
+    }
+
     pub fn force_flush_files(&self, task: &str) {
-        // TODO: use the callback to make the test more stable.
-        self.run(|| Task::ForceFlush(task.to_owned()));
-        self.sync();
+        // Force flush but not wait...
+        // Then the case may use `wait_flush`...
+        let _ = self.force_flush_files_and_wait(task);
+    }
+
+    pub fn force_flush_files_and_wait(&self, task: &str) -> impl Future<Output = ()> + '_ {
+        let (tx, mut rx) = tokio::sync::mpsc::channel(1);
+        self.run(|| Task::ForceFlush(TaskSelector::ByName(task.to_owned()), tx.clone()));
+        drop(tx);
+
+        async move {
+            while let Some(res) = tokio::time::timeout(Duration::from_secs(30), rx.recv())
+                .await
+                .expect("flush not finish after 30s")
+            {
+                if let Some(ref err) = res.error {
+                    panic!("failed to flush: {}", err)
+                }
+            }
+        }
     }
 
     pub fn run(&self, mut t: impl FnMut() -> Task) {