Skip to content

Use WaitLatch to detect Postmaster death and exit. #6

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
+35 −5
Open

Use WaitLatch to detect Postmaster death and exit. #6

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
086ff32
Use WaitLatch to detect Postmaster death and exit.
bizwark Apr 21, 2025
b0d7b4b
Add SIGTERM case to WaitLatch, to also handle graceful shutdown.
bizwark Apr 21, 2025
4fe16cd
Remove unnecessary unsafe block that is causing compiler warning.
bizwark Apr 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 18 additions & 3 deletions exts/rag_bge_small_en_v15/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ use rayon::{ThreadPool, ThreadPoolBuilder};
use std::{fs, os::unix::fs::PermissionsExt, sync::OnceLock};
use tokio::{
net::UnixListener,
time::{sleep, Duration},
time::Duration,
};
use tokio_stream::wrappers::UnixListenerStream;
use tonic::{transport::Server, Request, Response, Status};
Expand Down Expand Up @@ -167,8 +167,23 @@ pub extern "C" fn background_main(arg: pg_sys::Datum) {
.add_service(EmbeddingGeneratorServer::new(embedder))
.serve_with_incoming_shutdown(uds_stream, async {
while !BackgroundWorker::sigterm_received() {
sleep(Duration::from_millis(500)).await;
}
unsafe {
let mask = pg_sys::WL_EXIT_ON_PM_DEATH | pg_sys::WL_LATCH_SET;
let my_latch = &mut (*pg_sys::MyProc).procLatch as *mut pg_sys::Latch;
let events = pg_sys::WaitLatch(
my_latch,
mask as i32,
0,
pg_sys::PG_WAIT_EXTENSION,
);
pg_sys::ResetLatch(my_latch);

if (events & pg_sys::WL_EXIT_ON_PM_DEATH as i32) != 0 {
// postmaster died
pg_sys::proc_exit(1);
}
}
Comment on lines +170 to +185
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think you can use pgrx's wait_latch function for this

}
})
.await
.expect_or_pg_err("Couldn't create server");
Expand Down
19 changes: 17 additions & 2 deletions exts/rag_jina_reranker_v1_tiny_en/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ use reranking::{
use std::{fs, os::unix::fs::PermissionsExt, sync::OnceLock};
use tokio::{
net::UnixListener,
time::{sleep, Duration},
time::Duration,
};
use tokio_stream::wrappers::UnixListenerStream;
use tonic::{transport::Server, Request, Response, Status};
Expand Down Expand Up @@ -171,7 +171,22 @@ pub extern "C" fn background_main(arg: pg_sys::Datum) {
.add_service(RerankerServer::new(reranker))
.serve_with_incoming_shutdown(uds_stream, async {
while !BackgroundWorker::sigterm_received() {
sleep(Duration::from_millis(500)).await;
unsafe {
let mask = pg_sys::WL_EXIT_ON_PM_DEATH | pg_sys::WL_LATCH_SET;
let my_latch = &mut (*pg_sys::MyProc).procLatch as *mut pg_sys::Latch;
let events = pg_sys::WaitLatch(
my_latch,
mask as i32,
0,
pg_sys::PG_WAIT_EXTENSION,
);
pg_sys::ResetLatch(my_latch);

if (events & pg_sys::WL_EXIT_ON_PM_DEATH as i32) != 0 {
// postmaster died
pg_sys::proc_exit(1);
}
}
}
})
.await
Expand Down