Implementing graceful shutdown (#105)

* Initial commit for graceful shutdown

* fmt

* Add .vscode to gitignore

* Updates shutdown logic to use channels

* fmt

* fmt

* Adds shutdown timeout

* Fmt and updates tomls

* Updates readme

* fmt and updates log levels

* Update python tests to test shutdown

* merge changes

* Rename listener rx and update bash to be in line with master

* Update python test bash script ordering

* Adds error response message before shutdown

* Add details on shutdown event loop

* Fixes response length for error

* Adds handler for sigterm

* Uses ready for query function and fixes number of bytes

* fmt

tests/python/requirements.txt

@@ -1 +1,2 @@
 psycopg2==2.9.3
+psutil==5.9.1

tests/python/tests.py

@@ -1,22 +1,158 @@
+from typing import Tuple
 import psycopg2
+import psutil
+import os
+import signal
+import subprocess
+from threading import Thread
+import time
 
-def test_normal_db_access():
-    conn = psycopg2.connect("postgres://sharding_user:sharding_user@127.0.0.1:6432/sharded_db?application_name=testing_pgcat")
+SHUTDOWN_TIMEOUT = 5
+
+PGCAT_HOST = "127.0.0.1"
+PGCAT_PORT = "6432"
+
+
+def pgcat_start():
+    pg_cat_send_signal(signal.SIGTERM)
+    pgcat_start_command = "./target/debug/pgcat .circleci/pgcat.toml"
+    subprocess.Popen(pgcat_start_command.split())
+
+
+def pg_cat_send_signal(signal: signal.Signals):
+    for proc in psutil.process_iter(["pid", "name"]):
+        if "pgcat" == proc.name():
+            os.kill(proc.pid, signal)
+
+
+def connect_normal_db(
+    autocommit: bool = False,
+) -> Tuple[psycopg2.extensions.connection, psycopg2.extensions.cursor]:
+    conn = psycopg2.connect(
+        f"postgres://sharding_user:sharding_user@{PGCAT_HOST}:{PGCAT_PORT}/sharded_db?application_name=testing_pgcat"
+    )
+    conn.autocommit = autocommit
     cur = conn.cursor()
 
+    return (conn, cur)
+
+
+def cleanup_conn(conn: psycopg2.extensions.connection, cur: psycopg2.extensions.cursor):
+    cur.close()
+    conn.close()
+
+
+def test_normal_db_access():
+    conn, cur = connect_normal_db()
     cur.execute("SELECT 1")
     res = cur.fetchall()
     print(res)
+    cleanup_conn(conn, cur)
 
 
 def test_admin_db_access():
-    conn = psycopg2.connect("postgres://admin_user:admin_pass@127.0.0.1:6432/pgcat")
-    conn.autocommit = True # BEGIN/COMMIT is not supported by admin db
+    conn = psycopg2.connect(
+        f"postgres://admin_user:admin_pass@{PGCAT_HOST}:{PGCAT_PORT}/pgcat"
+    )
+    conn.autocommit = True  # BEGIN/COMMIT is not supported by admin db
     cur = conn.cursor()
 
     cur.execute("SHOW POOLS")
     res = cur.fetchall()
     print(res)
+    cleanup_conn(conn, cur)
+
+
+def test_shutdown_logic():
+
+    ##### NO ACTIVE QUERIES SIGINT HANDLING #####
+    # Start pgcat
+    server = Thread(target=pgcat_start)
+    server.start()
+
+    # Wait for server to fully start up
+    time.sleep(2)
+
+    # Create client connection and send query (not in transaction)
+    conn, cur = connect_normal_db(True)
+
+    cur.execute("BEGIN;")
+    cur.execute("SELECT 1;")
+    cur.execute("COMMIT;")
+
+    # Send sigint to pgcat
+    pg_cat_send_signal(signal.SIGINT)
+    time.sleep(1)
+
+    # Check that any new queries fail after sigint since server should close with no active transactions
+    try:
+        cur.execute("SELECT 1;")
+    except psycopg2.OperationalError as e:
+        pass
+    else:
+        # Fail if query execution succeeded
+        raise Exception("Server not closed after sigint")
+    cleanup_conn(conn, cur)
+
+    ##### HANDLE TRANSACTION WITH SIGINT #####
+    # Start pgcat
+    server = Thread(target=pgcat_start)
+    server.start()
+
+    # Wait for server to fully start up
+    time.sleep(2)
+
+    # Create client connection and begin transaction
+    conn, cur = connect_normal_db(True)
+
+    cur.execute("BEGIN;")
+    cur.execute("SELECT 1;")
+
+    # Send sigint to pgcat while still in transaction
+    pg_cat_send_signal(signal.SIGINT)
+    time.sleep(1)
+
+    # Check that any new queries succeed after sigint since server should still allow transaction to complete
+    try:
+        cur.execute("SELECT 1;")
+    except psycopg2.OperationalError as e:
+        # Fail if query fails since server closed
+        raise Exception("Server closed while in transaction", e.pgerror)
+
+    cleanup_conn(conn, cur)
+
+    ##### HANDLE SHUTDOWN TIMEOUT WITH SIGINT #####
+    # Start pgcat
+    server = Thread(target=pgcat_start)
+    server.start()
+
+    # Wait for server to fully start up
+    time.sleep(3)
+
+    # Create client connection and begin transaction, which should prevent server shutdown unless shutdown timeout is reached
+    conn, cur = connect_normal_db(True)
+
+    cur.execute("BEGIN;")
+    cur.execute("SELECT 1;")
+
+    # Send sigint to pgcat while still in transaction
+    pg_cat_send_signal(signal.SIGINT)
+
+    # pgcat shutdown timeout is set to SHUTDOWN_TIMEOUT seconds, so we sleep for SHUTDOWN_TIMEOUT + 1 seconds
+    time.sleep(SHUTDOWN_TIMEOUT + 1)
+
+    # Check that any new queries succeed after sigint since server should still allow transaction to complete
+    try:
+        cur.execute("SELECT 1;")
+    except psycopg2.OperationalError as e:
+        pass
+    else:
+        # Fail if query execution succeeded
+        raise Exception("Server not closed after sigint and expected timeout")
+
+    cleanup_conn(conn, cur)
+
 
 test_normal_db_access()
 test_admin_db_access()
+test_shutdown_logic()