file watcher singleton thread pool (#18)

cp
00ee677
and
cd8486a

Signed-off-by: SpadeA <tangchenjie1210@gmail.com>

Author: SpadeA-Tang

Cargo.toml

@@ -70,7 +70,7 @@ tokenizer-api = { version = "0.2", path = "./tokenizer-api", package = "tantivy-
 sketches-ddsketch = { version = "0.2.1", features = ["use_serde"] }
 futures-util = { version = "0.3.28", optional = true }
 lazy_static = "1.5.0"
-tokio = { version = "1.29.1", features = ["rt-multi-thread"] }
+tokio = { version = "1.29.1", features = ["rt-multi-thread", "macros", "time"] }
 async-channel = "2.3.1"
 
 [target.'cfg(windows)'.dependencies]
@@ -88,7 +88,6 @@ futures = "0.3.21"
 paste = "1.0.11"
 more-asserts = "0.3.1"
 rand_distr = "0.4.3"
-tokio = { version = "1.29.1", features = ["rt-multi-thread", "macros"] }
 
 [target.'cfg(not(windows))'.dev-dependencies]
 criterion = "0.5"

src/directory/file_watcher.rs

@@ -1,13 +1,15 @@
 use std::io::BufRead;
 use std::path::Path;
 use std::sync::atomic::{AtomicUsize, Ordering};
-use std::sync::Arc;
+use std::sync::{Arc, RwLock};
 use std::time::Duration;
 use std::{fs, io, thread};
 
 use crc32fast::Hasher;
+use tokio::task::JoinHandle;
 
 use crate::directory::{WatchCallback, WatchCallbackList, WatchHandle};
+use crate::indexer::TOKIO_FILE_WATCHER_WORKER_RUNTIME;
 
 const POLLING_INTERVAL: Duration = Duration::from_millis(if cfg!(test) { 1 } else { 500 });
 
@@ -16,6 +18,8 @@ pub struct FileWatcher {
     path: Arc<Path>,
     callbacks: Arc<WatchCallbackList>,
     state: Arc<AtomicUsize>, // 0: new, 1: runnable, 2: terminated
+    watch_handle: RwLock<Option<JoinHandle<()>>>,
+    wakeup_channel: RwLock<Option<async_channel::Sender<()>>>,
 }
 
 impl FileWatcher {
@@ -24,6 +28,8 @@ impl FileWatcher {
             path: Arc::from(path),
             callbacks: Default::default(),
             state: Default::default(),
+            watch_handle: RwLock::new(None),
+            wakeup_channel: RwLock::new(None),
         }
     }
 
@@ -40,29 +46,34 @@ impl FileWatcher {
         let callbacks = self.callbacks.clone();
         let state = self.state.clone();
 
-        thread::Builder::new()
-            .name("thread-tantivy-meta-file-watcher".to_string())
-            .spawn(move || {
-                let mut current_checksum_opt = None;
-
-                while state.load(Ordering::SeqCst) == 1 {
-                    if let Ok(checksum) = FileWatcher::compute_checksum(&path) {
-                        let metafile_has_changed = current_checksum_opt
-                            .map(|current_checksum| current_checksum != checksum)
-                            .unwrap_or(true);
-                        if metafile_has_changed {
-                            info!("Meta file {:?} was modified", path);
-                            current_checksum_opt = Some(checksum);
-                            // We actually ignore callbacks failing here.
-                            // We just wait for the end of their execution.
-                            let _ = callbacks.broadcast().wait();
-                        }
+        let (tx, rx) = async_channel::bounded(5);
+        self.wakeup_channel.write().unwrap().replace(tx);
+        let task = async move {
+            let mut current_checksum_opt = None;
+            while state.load(Ordering::SeqCst) == 1 {
+                if let Ok(checksum) = FileWatcher::compute_checksum(&path) {
+                    let metafile_has_changed = current_checksum_opt
+                        .map(|current_checksum| current_checksum != checksum)
+                        .unwrap_or(true);
+                    if metafile_has_changed {
+                        info!("Meta file {:?} was modified", path);
+                        current_checksum_opt = Some(checksum);
+                        // We actually ignore callbacks failing here.
+                        // We just wait for the end of their execution.
+                        let _ = callbacks.broadcast().wait();
                     }
+                }
 
-                    thread::sleep(POLLING_INTERVAL);
+                tokio::select! {
+                    _ = tokio::time::sleep(POLLING_INTERVAL) => {},
+                    _ = rx.recv() => {
+                        // Early wake up from sleep
+                    }
                 }
-            })
-            .expect("Failed to spawn meta file watcher thread");
+            }
+        };
+        let watch_handle = TOKIO_FILE_WATCHER_WORKER_RUNTIME.spawn(task);
+        self.watch_handle.write().unwrap().replace(watch_handle);
     }
 
     pub fn watch(&self, callback: WatchCallback) -> WatchHandle {
@@ -91,6 +102,11 @@ impl FileWatcher {
 
     pub fn graceful_stop(&self) {
         self.state.store(2, Ordering::SeqCst);
+        if let Some(handle) = self.watch_handle.write().unwrap().take() {
+            let _ = self.wakeup_channel.write().unwrap().take();
+            handle.abort();
+            info!("Meta file watcher thread joined/aborted");
+        }
     }
 }
 

src/directory/watch_event_router.rs

@@ -1,5 +1,6 @@
 use std::sync::{Arc, RwLock, Weak};
 
+use crate::indexer::MISC_POOL;
 use crate::FutureResult;
 
 /// Cloneable wrapper for callbacks registered when watching files of a `Directory`.
@@ -86,20 +87,14 @@ impl WatchCallbackList {
             let _ = sender.send(Ok(()));
             return result;
         }
-        let spawn_res = std::thread::Builder::new()
-            .name("watch-callbacks".to_string())
-            .spawn(move || {
-                for callback in callbacks {
-                    callback.call();
-                }
-                let _ = sender.send(Ok(()));
-            });
-        if let Err(err) = spawn_res {
-            error!(
-                "Failed to spawn thread to call watch callbacks. Cause: {:?}",
-                err
-            );
-        }
+
+        MISC_POOL.spawn(move || {
+            for callback in callbacks {
+                callback.call();
+            }
+            let _ = sender.send(Ok(()));
+        });
+
         result
     }
 }

src/indexer/mod.rs

@@ -27,7 +27,9 @@ pub use self::index_writer::IndexWriter;
 pub use self::log_merge_policy::LogMergePolicy;
 pub use self::merge_operation::MergeOperation;
 pub use self::merge_policy::{MergeCandidate, MergePolicy, NoMergePolicy};
-pub use self::pool::{TOKIO_DOCSTORE_COMPRESS_RUNTIME, TOKIO_RUNTIME};
+pub use self::pool::{
+    MISC_POOL, TOKIO_DOCSTORE_COMPRESS_RUNTIME, TOKIO_FILE_WATCHER_WORKER_RUNTIME, TOKIO_RUNTIME,
+};
 pub use self::prepared_commit::PreparedCommit;
 pub use self::segment_entry::SegmentEntry;
 pub use self::segment_serializer::SegmentSerializer;

src/indexer/pool.rs

@@ -9,6 +9,8 @@ const MILVUS_TANTIVY_WRITER_THREAD_NUM: &str = "MILVUS_TANTIVY_WRITER_THREAD_NUM
 const MILVUS_TOKIO_THREAD_NUM: &str = "MILVUS_TOKIO_THREAD_NUM";
 const MILVUS_TOKIO_DOCSTORE_COMPRESS_THREAD_NUM: &str =
     "MILVUS_TANTIVY_DOCSTORE_COMPRESS_THREAD_NUM";
+const MILVUS_TOKIO_FILE_WATCHER_THREAD_NUM: &str = "MILVUS_TOKIO_FILE_WATCHER_THREAD_NUM";
+const MILVUS_MISC_POLL: &str = "MILVUS_MISC_POLL";
 
 lazy_static! {
     pub static ref TOKIO_RUNTIME: tokio::runtime::Runtime =
@@ -35,6 +37,18 @@ lazy_static! {
             .enable_all()
             .build()
             .expect("Failed to create tokio runtime");
+    pub static ref TOKIO_FILE_WATCHER_WORKER_RUNTIME: tokio::runtime::Runtime =
+        tokio::runtime::Builder::new_multi_thread()
+            .worker_threads(get_num_thread(MILVUS_TOKIO_FILE_WATCHER_THREAD_NUM))
+            .thread_name("tantivy-file-watcher")
+            .enable_all()
+            .build()
+            .expect("Failed to create tokio runtime");
+    pub static ref MISC_POOL: ThreadPool = ThreadPoolBuilder::new()
+        .num_threads(get_num_thread(MILVUS_MISC_POLL))
+        .thread_name(|sz| format!("tantivy-misc{}", sz))
+        .build()
+        .expect("Failed to create tantivy-misc thread pool");
 }
 
 fn default_num_thread() -> usize {