mirror of
https://github.com/qdrant/qdrant.git
synced 2026-08-07 02:20:55 -05:00
Catch optimizer panics (#2485)
* Add panic handler to stoppable task * Define panic handler for optimizer task, report panic to segment holder * Report panic payload message when optimizer panics * Add function to handle and clean finished optimizer tasks * Rewrite optimizer worker, periodically clean finished optimizer tasks * Minor improvements and codespell fixes * Add simple test to ensure panic callback is called * Explicitly handle elapsed timeout errors in optimizer handle handling
This commit is contained in:
@@ -479,6 +479,8 @@ impl<'s> SegmentHolder {
|
||||
}
|
||||
|
||||
pub fn report_optimizer_error<E: Into<CollectionError>>(&mut self, error: E) {
|
||||
// Save only the first error
|
||||
// If is more likely to be the real cause of all further problems
|
||||
if self.optimizer_errors.is_none() {
|
||||
self.optimizer_errors = Some(error.into());
|
||||
}
|
||||
|
||||
@@ -1,12 +1,16 @@
|
||||
use std::any::Any;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Arc, Weak};
|
||||
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
type PanicPayload = Box<dyn Any + Send + 'static>;
|
||||
|
||||
pub struct StoppableTaskHandle<T> {
|
||||
pub join_handle: JoinHandle<Option<T>>,
|
||||
started: Arc<AtomicBool>,
|
||||
stopped: Weak<AtomicBool>,
|
||||
panic_handler: Option<Box<dyn Fn(PanicPayload) + Sync + Send>>,
|
||||
}
|
||||
|
||||
impl<T> StoppableTaskHandle<T> {
|
||||
@@ -28,9 +32,48 @@ impl<T> StoppableTaskHandle<T> {
|
||||
self.ask_to_stop();
|
||||
self.is_started().then_some(self.join_handle)
|
||||
}
|
||||
|
||||
/// Join this stoppable task and handle any panics
|
||||
///
|
||||
/// Any panics are propagated through the configured panic handler. If no handler is
|
||||
/// configured, nothing happens.
|
||||
///
|
||||
/// To call this, the task must already be finished. Otherwise it panics in development, or
|
||||
/// blocks in release.
|
||||
pub async fn join_and_handle_panic(self) {
|
||||
debug_assert!(
|
||||
self.join_handle.is_finished(),
|
||||
"Task must be finished, we cannot block here on awaiting the join handle",
|
||||
);
|
||||
|
||||
match self.join_handle.await {
|
||||
Ok(_) => {}
|
||||
Err(err) if err.is_cancelled() => {}
|
||||
// Propagate panic
|
||||
Err(err) if err.is_panic() && self.panic_handler.is_some() => {
|
||||
log::trace!("Handling stoppable task panic through custom panic handler");
|
||||
let panic = err.into_panic();
|
||||
let panic_handler = self.panic_handler.unwrap();
|
||||
panic_handler(panic);
|
||||
}
|
||||
Err(err) if err.is_panic() => {
|
||||
log::debug!("Stoppable task panicked without panic handler");
|
||||
}
|
||||
// Log error on unknown error
|
||||
Err(err) => {
|
||||
log::error!("Stoppable task handle error for unknown reason: {err}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn spawn_stoppable<F, T>(f: F) -> StoppableTaskHandle<T>
|
||||
/// Spawn stoppable task `f`
|
||||
///
|
||||
/// An optional `panic_handler` may be given, eventually called if the task panicked.
|
||||
pub fn spawn_stoppable<F, T>(
|
||||
f: F,
|
||||
panic_handler: Option<Box<dyn Fn(PanicPayload) + Sync + Send>>,
|
||||
) -> StoppableTaskHandle<T>
|
||||
where
|
||||
F: FnOnce(&AtomicBool) -> T + Send + 'static,
|
||||
T: Send + 'static,
|
||||
@@ -57,11 +100,25 @@ where
|
||||
}),
|
||||
started: started_c,
|
||||
stopped: stopped_w,
|
||||
panic_handler,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a panic payload into a string
|
||||
///
|
||||
/// This converts `String` and `&str` panic payloads into a string.
|
||||
/// Other payload types are formatted as is, and may be non descriptive.
|
||||
pub(crate) fn panic_payload_into_string(payload: PanicPayload) -> String {
|
||||
payload
|
||||
.downcast::<&str>()
|
||||
.map(|msg| msg.to_string())
|
||||
.or_else(|payload| payload.downcast::<String>().map(|msg| msg.to_string()))
|
||||
.unwrap_or_else(|payload| format!("{payload:?}"))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Mutex;
|
||||
use std::thread;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
@@ -91,7 +148,7 @@ mod tests {
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn test_task_stop() {
|
||||
let handle = spawn_stoppable(counting_task);
|
||||
let handle = spawn_stoppable(counting_task, None);
|
||||
|
||||
// Signal task to stop after ~20 steps
|
||||
sleep(STEP * 20).await;
|
||||
@@ -117,7 +174,7 @@ mod tests {
|
||||
const TASKS: usize = 64;
|
||||
|
||||
let handles = (0..TASKS)
|
||||
.map(|_| spawn_stoppable(counting_task))
|
||||
.map(|_| spawn_stoppable(counting_task, None))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
// Signal tasks to stop after ~20 steps
|
||||
@@ -139,4 +196,31 @@ mod tests {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn test_task_panic() {
|
||||
let panic_payload = Arc::new(Mutex::new(String::new()));
|
||||
let handle = spawn_stoppable(
|
||||
|_| {
|
||||
thread::sleep(STEP * 50);
|
||||
panic!("stoppable task panicked");
|
||||
},
|
||||
Some(Box::new({
|
||||
let panic_payload = panic_payload.clone();
|
||||
move |payload| {
|
||||
*panic_payload.lock().unwrap() = panic_payload_into_string(payload);
|
||||
}
|
||||
})),
|
||||
);
|
||||
|
||||
sleep(STEP * 20).await;
|
||||
assert!(!handle.is_finished());
|
||||
sleep(STEP * 100).await;
|
||||
assert!(handle.is_finished());
|
||||
|
||||
// Join handle to call back panic
|
||||
handle.join_and_handle_panic().await;
|
||||
|
||||
assert_eq!(*panic_payload.lock().unwrap(), "stoppable task panicked");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -11,19 +11,27 @@ use tokio::runtime::Handle;
|
||||
use tokio::sync::mpsc::{self, Receiver, Sender};
|
||||
use tokio::sync::{oneshot, Mutex as TokioMutex};
|
||||
use tokio::task::JoinHandle;
|
||||
use tokio::time::Duration;
|
||||
use tokio::time::error::Elapsed;
|
||||
use tokio::time::{timeout, Duration};
|
||||
|
||||
use crate::collection_manager::collection_updater::CollectionUpdater;
|
||||
use crate::collection_manager::holders::segment_holder::LockedSegmentHolder;
|
||||
use crate::collection_manager::optimizers::segment_optimizer::SegmentOptimizer;
|
||||
use crate::collection_manager::optimizers::{Tracker, TrackerLog, TrackerStatus};
|
||||
use crate::common::stoppable_task::{spawn_stoppable, StoppableTaskHandle};
|
||||
use crate::common::stoppable_task::{
|
||||
panic_payload_into_string, spawn_stoppable, StoppableTaskHandle,
|
||||
};
|
||||
use crate::operations::shared_storage_config::SharedStorageConfig;
|
||||
use crate::operations::types::{CollectionError, CollectionResult};
|
||||
use crate::operations::CollectionUpdateOperations;
|
||||
use crate::shards::local_shard::LockedWal;
|
||||
use crate::wal::WalError;
|
||||
|
||||
/// Interval at which the optimizer worker cleans up old optimization handles
|
||||
///
|
||||
/// The longer the duration, the longer it takes for panicked tasks to be reported.
|
||||
const OPTIMIZER_CLEANUP_INTERVAL: Duration = Duration::from_secs(5);
|
||||
|
||||
pub type Optimizer = dyn SegmentOptimizer + Sync + Send;
|
||||
|
||||
/// Information, required to perform operation and notify regarding the result
|
||||
@@ -228,54 +236,73 @@ impl UpdateHandler {
|
||||
optimizer.check_condition(segments.clone(), &scheduled_segment_ids);
|
||||
if nonoptimal_segment_ids.is_empty() {
|
||||
break;
|
||||
} else {
|
||||
let optim = optimizer.clone();
|
||||
let optimizers_log = optimizers_log.clone();
|
||||
let segs = segments.clone();
|
||||
let nsi = nonoptimal_segment_ids.clone();
|
||||
for sid in &nsi {
|
||||
scheduled_segment_ids.insert(*sid);
|
||||
}
|
||||
let callback = callback.clone();
|
||||
|
||||
handles.push(spawn_stoppable(move |stopped| {
|
||||
// Track optimizer status
|
||||
let tracker = Tracker::start(optim.as_ref().name(), nsi.clone());
|
||||
let tracker_handle = tracker.handle();
|
||||
optimizers_log.lock().register(tracker);
|
||||
|
||||
// Optimize and handle result
|
||||
match optim.as_ref().optimize(segs.clone(), nsi, stopped) {
|
||||
Ok(result) => {
|
||||
tracker_handle.update(TrackerStatus::Done);
|
||||
callback(result); // Perform some actions when optimization if finished
|
||||
result
|
||||
}
|
||||
Err(error) => match error {
|
||||
CollectionError::Cancelled { description } => {
|
||||
log::debug!("Optimization cancelled - {}", description);
|
||||
tracker_handle.update(TrackerStatus::Cancelled(description));
|
||||
false
|
||||
}
|
||||
_ => {
|
||||
// Save only the first error
|
||||
// If is more likely to be the real cause of all further problems
|
||||
segs.write().report_optimizer_error(error.clone());
|
||||
|
||||
// Error of the optimization can not be handled by API user
|
||||
// It is only possible to fix after full restart,
|
||||
// so the best available action here is to stop whole
|
||||
// optimization thread and log the error
|
||||
log::error!("Optimization error: {}", error);
|
||||
|
||||
tracker_handle.update(TrackerStatus::Error(error.to_string()));
|
||||
|
||||
panic!("Optimization error: {error}");
|
||||
}
|
||||
},
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
let optimizer = optimizer.clone();
|
||||
let optimizers_log = optimizers_log.clone();
|
||||
let segments = segments.clone();
|
||||
let nsi = nonoptimal_segment_ids.clone();
|
||||
scheduled_segment_ids.extend(&nsi);
|
||||
let callback = callback.clone();
|
||||
|
||||
let handle = spawn_stoppable(
|
||||
// Stoppable task
|
||||
{
|
||||
let segments = segments.clone();
|
||||
move |stopped| {
|
||||
// Track optimizer status
|
||||
let tracker = Tracker::start(optimizer.as_ref().name(), nsi.clone());
|
||||
let tracker_handle = tracker.handle();
|
||||
optimizers_log.lock().register(tracker);
|
||||
|
||||
// Optimize and handle result
|
||||
match optimizer.as_ref().optimize(segments.clone(), nsi, stopped) {
|
||||
// Perform some actions when optimization if finished
|
||||
Ok(result) => {
|
||||
tracker_handle.update(TrackerStatus::Done);
|
||||
callback(result);
|
||||
result
|
||||
}
|
||||
// Handle and report errors
|
||||
Err(error) => match error {
|
||||
CollectionError::Cancelled { description } => {
|
||||
log::debug!("Optimization cancelled - {}", description);
|
||||
tracker_handle
|
||||
.update(TrackerStatus::Cancelled(description));
|
||||
false
|
||||
}
|
||||
_ => {
|
||||
segments.write().report_optimizer_error(error.clone());
|
||||
|
||||
// Error of the optimization can not be handled by API user
|
||||
// It is only possible to fix after full restart,
|
||||
// so the best available action here is to stop whole
|
||||
// optimization thread and log the error
|
||||
log::error!("Optimization error: {}", error);
|
||||
|
||||
tracker_handle
|
||||
.update(TrackerStatus::Error(error.to_string()));
|
||||
|
||||
panic!("Optimization error: {error}");
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
},
|
||||
// Panic handler
|
||||
Some(Box::new(move |panic_payload| {
|
||||
let panic_msg = panic_payload_into_string(panic_payload);
|
||||
log::warn!(
|
||||
"Optimization task panicked, collection may be in unstable state: {panic_msg}"
|
||||
);
|
||||
segments
|
||||
.write()
|
||||
.report_optimizer_error(CollectionError::service_error(format!(
|
||||
"Optimization task panicked: {panic_msg}"
|
||||
)));
|
||||
})),
|
||||
);
|
||||
handles.push(handle);
|
||||
}
|
||||
}
|
||||
handles
|
||||
@@ -302,7 +329,33 @@ impl UpdateHandler {
|
||||
);
|
||||
let mut handles = optimization_handles.lock().await;
|
||||
handles.append(&mut new_handles);
|
||||
handles.retain(|h| !h.is_finished())
|
||||
}
|
||||
|
||||
/// Cleanup finalized optimization task handles
|
||||
///
|
||||
/// This finds and removes completed tasks from our list of optimization handles.
|
||||
/// It also propagates any panics (and unknown errors) so we properly handle them if desired.
|
||||
///
|
||||
/// It is essential to call this every once in a while for handling panics in time.
|
||||
async fn cleanup_optimization_handles(
|
||||
optimization_handles: Arc<TokioMutex<Vec<StoppableTaskHandle<bool>>>>,
|
||||
) {
|
||||
// Remove finished handles
|
||||
let finished_handles: Vec<_> = {
|
||||
let mut handles = optimization_handles.lock().await;
|
||||
(0..handles.len())
|
||||
.filter(|i| handles[*i].is_finished())
|
||||
.collect::<Vec<_>>()
|
||||
.into_iter()
|
||||
.rev()
|
||||
.map(|i| handles.remove(i))
|
||||
.collect()
|
||||
};
|
||||
|
||||
// Finalize all finished handles to propagate panics
|
||||
for handle in finished_handles {
|
||||
handle.join_and_handle_panic().await;
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
@@ -316,18 +369,27 @@ impl UpdateHandler {
|
||||
optimizers_log: Arc<Mutex<TrackerLog>>,
|
||||
max_handles: usize,
|
||||
) {
|
||||
while let Some(signal) = receiver.recv().await {
|
||||
match signal {
|
||||
OptimizerSignal::Nop | OptimizerSignal::Operation(_) => {
|
||||
loop {
|
||||
let receiver = timeout(OPTIMIZER_CLEANUP_INTERVAL, receiver.recv());
|
||||
let result = receiver.await;
|
||||
|
||||
// Always clean up on any signal
|
||||
Self::cleanup_optimization_handles(optimization_handles.clone()).await;
|
||||
|
||||
match result {
|
||||
// Channel closed or stop signal
|
||||
Ok(None | Some(OptimizerSignal::Stop)) => break,
|
||||
// Clean up interval
|
||||
Err(Elapsed { .. }) => continue,
|
||||
// Optimizer signal
|
||||
Ok(Some(signal @ (OptimizerSignal::Nop | OptimizerSignal::Operation(_)))) => {
|
||||
// If not forcing with Nop, wait on next signal if we have too many handles
|
||||
if signal != OptimizerSignal::Nop
|
||||
&& optimization_handles.lock().await.len() >= max_handles
|
||||
{
|
||||
let mut handles = optimization_handles.lock().await;
|
||||
handles.retain(|h| !h.is_finished());
|
||||
continue;
|
||||
}
|
||||
// We skip the check for number of optimization handles here
|
||||
// Because `Nop` usually means that we need to force the optimization
|
||||
|
||||
if Self::try_recover(segments.clone(), wal.clone())
|
||||
.await
|
||||
.is_err()
|
||||
@@ -343,8 +405,6 @@ impl UpdateHandler {
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
OptimizerSignal::Stop => break,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user