refactor: move the background task handler onto the parent IngesterGuard
* follow the pattern of the periodic wal rotation * do NOT follow the pattern of the wal.flusher_taskpull/24376/head
parent
b4b89699cd
commit
b961bc79c4
|
@ -99,6 +99,11 @@ pub struct IngesterGuard<T> {
|
|||
/// Aborted on drop.
|
||||
rotation_task: tokio::task::JoinHandle<()>,
|
||||
|
||||
/// The handle of the periodic disk protection task.
|
||||
///
|
||||
/// Aborted on drop.
|
||||
disk_protection_task: tokio::task::JoinHandle<()>,
|
||||
|
||||
/// The task handle executing the graceful shutdown once triggered.
|
||||
graceful_shutdown_handler: tokio::task::JoinHandle<()>,
|
||||
shutdown_complete: Shared<oneshot::Receiver<()>>,
|
||||
|
@ -125,6 +130,7 @@ where
|
|||
impl<T> Drop for IngesterGuard<T> {
|
||||
fn drop(&mut self) {
|
||||
self.rotation_task.abort();
|
||||
self.disk_protection_task.abort();
|
||||
self.graceful_shutdown_handler.abort();
|
||||
}
|
||||
}
|
||||
|
@ -332,7 +338,7 @@ where
|
|||
.map_err(InitError::WalInit)?;
|
||||
// Initialize the disk proetction after the WAL directory is initialized
|
||||
let disk_protection = InstrumentedDiskProtection::new(wal_directory, &metrics);
|
||||
disk_protection.start().await;
|
||||
let disk_protection_task = disk_protection.start().await;
|
||||
|
||||
// Replay the WAL log files, if any.
|
||||
let max_sequence_number =
|
||||
|
@ -411,6 +417,7 @@ where
|
|||
persist_handle,
|
||||
),
|
||||
rotation_task,
|
||||
disk_protection_task,
|
||||
graceful_shutdown_handler: shutdown_task,
|
||||
shutdown_complete: shutdown_rx.shared(),
|
||||
})
|
||||
|
|
|
@ -1,7 +1,6 @@
|
|||
use std::{borrow::Cow, path::PathBuf, sync::Arc, time::Duration};
|
||||
use std::{borrow::Cow, path::PathBuf, time::Duration};
|
||||
|
||||
use metric::{Attributes, U64Gauge};
|
||||
use parking_lot::Mutex;
|
||||
use sysinfo::{DiskExt, System, SystemExt};
|
||||
use tokio::{self, task::JoinHandle};
|
||||
|
||||
|
@ -66,8 +65,6 @@ impl DiskProtectionMetrics {
|
|||
pub struct InstrumentedDiskProtection {
|
||||
/// The metrics that are reported to the registry.
|
||||
metrics: DiskProtectionMetrics,
|
||||
/// The handle to terminate the background task.
|
||||
background_task: Mutex<Option<JoinHandle<()>>>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for InstrumentedDiskProtection {
|
||||
|
@ -81,27 +78,12 @@ impl InstrumentedDiskProtection {
|
|||
pub fn new(directory_to_track: PathBuf, registry: &metric::Registry) -> Self {
|
||||
let metrics = DiskProtectionMetrics::new(directory_to_track, registry);
|
||||
|
||||
Self {
|
||||
metrics,
|
||||
background_task: Default::default(),
|
||||
}
|
||||
Self { metrics }
|
||||
}
|
||||
|
||||
/// Start the [`InstrumentedDiskProtection`] background task.
|
||||
pub async fn start(self) {
|
||||
let rc_self = Arc::new(self);
|
||||
let rc_self_clone = Arc::clone(&rc_self);
|
||||
|
||||
*rc_self.background_task.lock() = Some(tokio::task::spawn(async move {
|
||||
rc_self_clone.background_task().await
|
||||
}));
|
||||
}
|
||||
|
||||
/// Stop the [`InstrumentedDiskProtection`] background task.
|
||||
pub fn stop(&mut self) {
|
||||
if let Some(t) = self.background_task.lock().take() {
|
||||
t.abort()
|
||||
}
|
||||
pub async fn start(self) -> JoinHandle<()> {
|
||||
tokio::task::spawn(async move { self.background_task().await })
|
||||
}
|
||||
|
||||
/// The background task that periodically performs the disk protection check.
|
||||
|
@ -118,15 +100,10 @@ impl InstrumentedDiskProtection {
|
|||
}
|
||||
}
|
||||
|
||||
impl Drop for InstrumentedDiskProtection {
|
||||
fn drop(&mut self) {
|
||||
// future-proof, such that stop does not need to be explicitly called.
|
||||
self.stop();
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use metric::Metric;
|
||||
|
||||
use super::*;
|
||||
|
@ -135,18 +112,20 @@ mod tests {
|
|||
async fn test_metrics() {
|
||||
let registry = Arc::new(metric::Registry::new());
|
||||
|
||||
struct MockAnyStruct;
|
||||
struct MockAnyStruct {
|
||||
abort_handle: JoinHandle<()>,
|
||||
}
|
||||
|
||||
impl MockAnyStruct {
|
||||
pub(crate) async fn new(registry: &metric::Registry) -> Self {
|
||||
let disk_protection = InstrumentedDiskProtection::new(PathBuf::from("/"), registry);
|
||||
disk_protection.start().await;
|
||||
let abort_handle = disk_protection.start().await;
|
||||
|
||||
Self
|
||||
Self { abort_handle }
|
||||
}
|
||||
}
|
||||
|
||||
let _mock = MockAnyStruct::new(®istry).await;
|
||||
let mock = MockAnyStruct::new(®istry).await;
|
||||
|
||||
tokio::time::sleep(2 * Duration::from_secs(2)).await;
|
||||
|
||||
|
@ -158,5 +137,6 @@ mod tests {
|
|||
.fetch();
|
||||
|
||||
assert!(recorded_metric > 0_u64);
|
||||
mock.abort_handle.abort();
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue