From 543c809ac2086d421fc48792a6f197ac35f06fa4 Mon Sep 17 00:00:00 2001
From: Zhao Junwang <zhjwpku@gmail.com>
Date: Fri, 22 Dec 2023 23:13:07 +0800
Subject: [PATCH v15 2/2] set transaction_timeout before next command

Signed-off-by: Zhao Junwang <zhjwpku@gmail.com>
---
 src/backend/tcop/postgres.c | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index a2611cf8e6..75f5b21dc0 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -2745,10 +2745,6 @@ start_xact_command(void)
 	{
 		StartTransactionCommand();
 
-		/* Schedule or reschedule transaction timeout */
-		if (TransactionTimeout > 0)
-			enable_timeout_after(TRANSACTION_TIMEOUT, TransactionTimeout);
-
 		xact_started = true;
 	}
 
@@ -4144,6 +4140,7 @@ PostgresMain(const char *dbname, const char *username)
 	volatile bool send_ready_for_query = true;
 	volatile bool idle_in_transaction_timeout_enabled = false;
 	volatile bool idle_session_timeout_enabled = false;
+	volatile bool transaction_timeout_enabled = false;
 
 	Assert(dbname != NULL);
 	Assert(username != NULL);
@@ -4357,6 +4354,7 @@ PostgresMain(const char *dbname, const char *username)
 		QueryCancelPending = false;
 		idle_in_transaction_timeout_enabled = false;
 		idle_session_timeout_enabled = false;
+		transaction_timeout_enabled = false;
 
 		/* Not reading from the client anymore. */
 		DoingCommandRead = false;
@@ -4527,6 +4525,13 @@ PostgresMain(const char *dbname, const char *username)
 					enable_timeout_after(IDLE_IN_TRANSACTION_SESSION_TIMEOUT,
 										 IdleInTransactionSessionTimeout);
 				}
+
+				/* Schedule or reschedule transaction timeout */
+				if (!transaction_timeout_enabled && TransactionTimeout > 0)
+				{
+					enable_timeout_after(TRANSACTION_TIMEOUT, TransactionTimeout);
+					transaction_timeout_enabled = true;
+				}
 			}
 			else
 			{
@@ -4580,8 +4585,10 @@ PostgresMain(const char *dbname, const char *username)
 										 IdleSessionTimeout);
 				}
 
-				if (get_timeout_active(TRANSACTION_TIMEOUT))
+				if (transaction_timeout_enabled) {
 					disable_timeout(TRANSACTION_TIMEOUT, false);
+					transaction_timeout_enabled = false;
+				}
 			}
 
 			/* Report any recently-changed GUC options */
-- 
2.41.0

