chore: friendly OOM message (#223)

Signed-off-by: usamoi <[email protected]>

Author: usamoi

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::<AllocErrorPanicPayload>() {
+            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");

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;