From 6e26588d17d212e703915507a4913a592b3082d5 Mon Sep 17 00:00:00 2001
From: Dmitry Rodionov <dmitry@neon.tech>
Date: Tue, 7 Jun 2022 17:10:39 +0300
Subject: [PATCH] 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 | 14 +++++---------
 pageserver/src/lib.rs              |  2 --
 pageserver/src/page_service.rs     |  4 ++++
 3 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/libs/utils/src/postgres_backend.rs b/libs/utils/src/postgres_backend.rs
index ff71423122..79dca96fcf 100644
--- a/libs/utils/src/postgres_backend.rs
+++ b/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);
-}
diff --git a/pageserver/src/lib.rs b/pageserver/src/lib.rs
index fdce0e5c5f..a68c277114 100644
--- a/pageserver/src/lib.rs
+++ b/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
diff --git a/pageserver/src/page_service.rs b/pageserver/src/page_service.rs
index f6a088d4b5..30f0d241d6 100644
--- a/pageserver/src/page_service.rs
+++ b/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,