Allow to customize shutdown condition in PostgresBackend

Use it in PageServerHandler to check per thread shutdown condition
from thread_mgr which takes into account tenants and timelines

libs/utils/src/postgres_backend.rs

@@ -13,13 +13,10 @@ use std::fmt;
 use std::io::{self, Write};
 use std::net::{Shutdown, SocketAddr, TcpStream};
 use std::str::FromStr;
-use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
 use tracing::*;
 
-static PGBACKEND_SHUTDOWN_REQUESTED: AtomicBool = AtomicBool::new(false);
-
 pub trait Handler {
     /// Handle single query.
     /// postgres_backend will issue ReadyForQuery after calling this (this
@@ -45,6 +42,10 @@ pub trait Handler {
     fn check_auth_jwt(&mut self, _pgb: &mut PostgresBackend, _jwt_response: &[u8]) -> Result<()> {
         bail!("JWT auth failed")
     }
+
+    fn is_shutdown_requested(&self) -> bool {
+        false
+    }
 }
 
 /// PostgresBackend protocol state.
@@ -274,7 +275,7 @@ impl PostgresBackend {
 
         let mut unnamed_query_string = Bytes::new();
 
-        while !PGBACKEND_SHUTDOWN_REQUESTED.load(Ordering::Relaxed) {
+        while !handler.is_shutdown_requested() {
             match self.read_message() {
                 Ok(message) => {
                     if let Some(msg) = message {
@@ -493,8 +494,3 @@ impl PostgresBackend {
         Ok(ProcessMsgResult::Continue)
     }
 }
-
-// Set the flag to inform connections to cancel
-pub fn set_pgbackend_shutdown_requested() {
-    PGBACKEND_SHUTDOWN_REQUESTED.swap(true, Ordering::Relaxed);
-}

pageserver/src/lib.rs

@@ -24,7 +24,6 @@ pub mod walredo;
 
 use lazy_static::lazy_static;
 use tracing::info;
-use utils::postgres_backend;
 
 use crate::thread_mgr::ThreadKind;
 use metrics::{register_int_gauge_vec, IntGaugeVec};
@@ -73,7 +72,6 @@ pub fn shutdown_pageserver(exit_code: i32) {
     thread_mgr::shutdown_threads(Some(ThreadKind::LibpqEndpointListener), None, None);
 
     // Shut down any page service threads.
-    postgres_backend::set_pgbackend_shutdown_requested();
     thread_mgr::shutdown_threads(Some(ThreadKind::PageRequestHandler), None, None);
 
     // Shut down all the tenants. This flushes everything to disk and kills

pageserver/src/page_service.rs

@@ -676,6 +676,10 @@ impl postgres_backend::Handler for PageServerHandler {
         Ok(())
     }
 
+    fn is_shutdown_requested(&self) -> bool {
+        thread_mgr::is_shutdown_requested()
+    }
+
     fn process_query(
         &mut self,
         pgb: &mut PostgresBackend,