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:
Tim Visée
2023-08-30 10:13:17 +02:00
committed by GitHub
parent 8275b42f84
commit c8d2d33d5d
3 changed files with 208 additions and 62 deletions

View File

@@ -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());
}

View File

@@ -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");
}
}

View File

@@ -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,
}
}
}