diff --git a/src/bgworker/mod.rs b/src/bgworker/mod.rs index 5a8d84b52..6c65b73ed 100644 --- a/src/bgworker/mod.rs +++ b/src/bgworker/mod.rs @@ -4,12 +4,14 @@ pub mod upgrade; pub unsafe fn init() { use pgrx::bgworkers::BackgroundWorkerBuilder; use pgrx::bgworkers::BgWorkerStartTime; + use std::time::Duration; BackgroundWorkerBuilder::new("vectors") .set_function("vectors_main") .set_library("vectors") .set_argument(None) .enable_shmem_access(None) .set_start_time(BgWorkerStartTime::PostmasterStart) + .set_restart_time(Some(Duration::from_secs(1))) .load(); } @@ -19,6 +21,9 @@ extern "C" fn vectors_main(_arg: pgrx::pg_sys::Datum) { } pub fn main() { + pub struct AllocErrorPanicPayload { + pub layout: std::alloc::Layout, + } { let mut builder = env_logger::builder(); builder.target(env_logger::Target::Stderr); @@ -33,6 +38,10 @@ pub fn main() { builder.init(); } std::panic::set_hook(Box::new(|info| { + if let Some(oom) = info.payload().downcast_ref::() { + log::error!("Out of memory. Layout: {:?}.", oom.layout); + return; + } let backtrace; #[cfg(not(debug_assertions))] { @@ -44,6 +53,9 @@ pub fn main() { } log::error!("Panickied. Info: {:?}. Backtrace: {}.", info, backtrace); })); + std::alloc::set_alloc_error_hook(|layout| { + std::panic::panic_any(AllocErrorPanicPayload { layout }); + }); use service::worker::Worker; use std::path::Path; let path = Path::new("pg_vectors"); diff --git a/src/lib.rs b/src/lib.rs index 14910973d..fc99a9488 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,6 +2,7 @@ //! //! Provides an easy-to-use extension for vector similarity search. #![feature(never_type)] +#![feature(alloc_error_hook)] mod bgworker; mod datatype;