Terminate background worker if postmaster dies and build with bindgen.

exts/rag_bge_small_en_v15/Cargo.toml

@@ -44,6 +44,8 @@ ort-sys = { path = "../../lib/ort-2.0.0-rc.4/ort-sys" }
 
 [build-dependencies]
 tonic-build = "0.12.3"
+bindgen = "0.65"
+cc      = "1.0"
 
 [dev-dependencies]
 pgrx-tests = "0.12.6"

exts/rag_bge_small_en_v15/bindgen_pmsignal.h

@@ -0,0 +1,16 @@
+#pragma once
+
+// get the C _bool_ type
+#include <stdbool.h>
+
+// only declare the function you need
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+extern bool PostmasterIsAliveInternal();
+
+#ifdef __cplusplus
+}
+#endif
+

exts/rag_bge_small_en_v15/build.rs

@@ -1,4 +1,41 @@
+use std::{env, path::PathBuf, process::Command};
+use bindgen::Builder;
+use cc::Build;
+
 fn main() -> Result<(), Box<dyn std::error::Error>> {
+    println!("cargo:rerun-if-changed=build.rs");
+    println!("cargo:rerun-if-changed=signature_check.c");
+    println!("cargo:rerun-if-env-changed=PGRX_PG_CONFIG");
+    println!("cargo:rerun-if-env-changed=PG_CONFIG");
+
+    let pg_config = env::var("PGRX_PG_CONFIG")
+        .or_else(|_| env::var("PG_CONFIG"))
+        .unwrap_or_else(|_| "pg_config".into());
+    let output = Command::new(&pg_config)
+        .arg("--includedir-server")
+        .output()
+        .expect("failed to run pg_config");
+    let pg_inc = String::from_utf8(output.stdout).unwrap().trim().to_string();
+
+    let bindings = bindgen::Builder::default()
+        .clang_arg(format!("-I."))
+        .header("bindgen_pmsignal.h")
+        .allowlist_function("PostmasterIsAliveInternal")
+        .generate()
+        .expect("bindgen failed");
+
+    let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
+    bindings
+        .write_to_file(out_dir.join("bindings.rs"))
+        .expect("couldn't write bindings");
+
+    cc::Build::new()
+        .include(&pg_inc)
+        .file("signature_check.c")
+        // you can silence warnings:
+        .warnings(false)
+        .compile("signature_check");
+
   tonic_build::compile_protos("proto/embeddings.proto")?;
   Ok(())
 }

exts/rag_bge_small_en_v15/signature_check.c

@@ -0,0 +1,9 @@
+// signature_check.c
+
+#include <stdbool.h>
+#include <c.h>
+#include <storage/pmsignal.h>
+
+// If the real signature ever changes, the assignment here will error:
+static bool (*_check_sig)() = &PostmasterIsAliveInternal;
+

exts/rag_bge_small_en_v15/src/lib.rs

@@ -1,5 +1,6 @@
 mod chunk;
 mod errors;
+
 mod embeddings {
     tonic::include_proto!("embeddings");
 }
@@ -23,6 +24,10 @@ use tokio::{
 use tokio_stream::wrappers::UnixListenerStream;
 use tonic::{transport::Server, Request, Response, Status};
 
+include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
+
+const _CHECK_SIG: unsafe extern "C" fn() -> bool = PostmasterIsAliveInternal;
+
 // macros
 
 mconst!(ext_name, "rag_bge_small_en_v15");
@@ -166,7 +171,7 @@ pub extern "C" fn background_main(arg: pg_sys::Datum) {
             Server::builder()
                 .add_service(EmbeddingGeneratorServer::new(embedder))
                 .serve_with_incoming_shutdown(uds_stream, async {
-                    while !BackgroundWorker::sigterm_received() {
+                    while !BackgroundWorker::sigterm_received() && unsafe { PostmasterIsAliveInternal() } {
                         sleep(Duration::from_millis(500)).await;
                     }
                 })

exts/rag_jina_reranker_v1_tiny_en/Cargo.toml

@@ -41,6 +41,8 @@ ort-sys = { path = "../../lib/ort-2.0.0-rc.4/ort-sys" }
 
 [build-dependencies]
 tonic-build = "0.12.3"
+bindgen = "0.65"
+cc      = "1.0"
 
 [dev-dependencies]
 pgrx-tests = "0.12.6"

exts/rag_jina_reranker_v1_tiny_en/bindgen_pmsignal.h

@@ -0,0 +1,16 @@
+#pragma once
+
+// get the C _bool_ type
+#include <stdbool.h>
+
+// only declare the function you need
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+extern bool PostmasterIsAliveInternal();
+
+#ifdef __cplusplus
+}
+#endif
+

exts/rag_jina_reranker_v1_tiny_en/build.rs

@@ -1,4 +1,41 @@
+use std::{env, path::PathBuf, process::Command};
+use bindgen::Builder;
+use cc::Build;
+
 fn main() -> Result<(), Box<dyn std::error::Error>> {
+    println!("cargo:rerun-if-changed=build.rs");
+    println!("cargo:rerun-if-changed=signature_check.c");
+    println!("cargo:rerun-if-env-changed=PGRX_PG_CONFIG");
+    println!("cargo:rerun-if-env-changed=PG_CONFIG");
+
+    let pg_config = env::var("PGRX_PG_CONFIG")
+        .or_else(|_| env::var("PG_CONFIG"))
+        .unwrap_or_else(|_| "pg_config".into());
+    let output = Command::new(&pg_config)
+        .arg("--includedir-server")
+        .output()
+        .expect("failed to run pg_config");
+    let pg_inc = String::from_utf8(output.stdout).unwrap().trim().to_string();
+
+    let bindings = bindgen::Builder::default()
+        .clang_arg(format!("-I."))
+        .header("bindgen_pmsignal.h")
+        .allowlist_function("PostmasterIsAliveInternal")
+        .generate()
+        .expect("bindgen failed");
+
+    let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
+    bindings
+        .write_to_file(out_dir.join("bindings.rs"))
+        .expect("couldn't write bindings");
+
+    cc::Build::new()
+        .include(&pg_inc)
+        .file("signature_check.c")
+        // you can silence warnings:
+        .warnings(false)
+        .compile("signature_check");
+
   tonic_build::compile_protos("proto/reranking.proto")?;
   Ok(())
 }

exts/rag_jina_reranker_v1_tiny_en/signature_check.c

@@ -0,0 +1,9 @@
+// signature_check.c
+
+#include <stdbool.h>
+#include <c.h>
+#include <storage/pmsignal.h>
+
+// If the real signature ever changes, the assignment here will error:
+static bool (*_check_sig)() = &PostmasterIsAliveInternal;
+

exts/rag_jina_reranker_v1_tiny_en/src/lib.rs

@@ -22,6 +22,10 @@ use tokio::{
 use tokio_stream::wrappers::UnixListenerStream;
 use tonic::{transport::Server, Request, Response, Status};
 
+include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
+
+const _CHECK_SIG: unsafe extern "C" fn() -> bool = PostmasterIsAliveInternal;
+
 // macros
 
 mconst!(ext_name, "rag_jina_reranker_v1_tiny_en");
@@ -170,7 +174,7 @@ pub extern "C" fn background_main(arg: pg_sys::Datum) {
             Server::builder()
                 .add_service(RerankerServer::new(reranker))
                 .serve_with_incoming_shutdown(uds_stream, async {
-                    while !BackgroundWorker::sigterm_received() {
+                    while !BackgroundWorker::sigterm_received() && unsafe { PostmasterIsAliveInternal() } {
                         sleep(Duration::from_millis(500)).await;
                     }
                 })