chore: friendly OOM message

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

Author: usamoi

src/bgworker/mod.rs

@@ -1,15 +1,21 @@
 pub mod normal;
 pub mod upgrade;
 
+pub struct OomError {
+    pub layout: std::alloc::Layout,
+}
+
 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();
 }
 
@@ -33,6 +39,10 @@ pub fn main() {
         builder.init();
     }
     std::panic::set_hook(Box::new(|info| {
+        if let Some(oom) = info.payload().downcast_ref::<OomError>() {
+            log::error!("Out of memory. Layout: {:?}.", oom.layout);
+            return;
+        }
         let backtrace;
         #[cfg(not(debug_assertions))]
         {
@@ -44,6 +54,9 @@ pub fn main() {
         }
         log::error!("Panickied. Info: {:?}. Backtrace: {}.", info, backtrace);
     }));
+    std::alloc::set_alloc_error_hook(|layout| {
+        std::panic::panic_any(OomError { layout });
+    });
     use service::worker::Worker;
     use std::path::Path;
     let path = Path::new("pg_vectors");

src/lib.rs

@@ -3,6 +3,7 @@
 //! Provides an easy-to-use extension for vector similarity search.
 #![feature(offset_of)]
 #![feature(never_type)]
+#![feature(alloc_error_hook)]
 
 mod bgworker;
 mod datatype;