From 79b68f1018e80b22671ed04d04301ec7dd1829bb Mon Sep 17 00:00:00 2001
From: usamoi <usamoi@outlook.com>
Date: Tue, 2 Jan 2024 20:16:10 +0800
Subject: [PATCH] chore: friendly OOM message

Signed-off-by: usamoi <usamoi@outlook.com>
---
 src/bgworker/mod.rs | 13 +++++++++++++
 src/lib.rs          |  1 +
 2 files changed, 14 insertions(+)

diff --git a/src/bgworker/mod.rs b/src/bgworker/mod.rs
index 5a8d84b52..200761202 100644
--- a/src/bgworker/mod.rs
+++ b/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");
diff --git a/src/lib.rs b/src/lib.rs
index 0f3faa181..b8d172d6f 100644
--- a/src/lib.rs
+++ b/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;