Cancel problems of query to pg_stat_statements

From f642c28124e1bd57764bd12296253b78fc9c083c Mon Sep 17 00:00:00 2001
From: rkhapov <r.khapov@ya.ru>
Date: Mon, 12 May 2025 15:29:17 +0500
Subject: [PATCH v2] pg_stat_statements.c: cancelable qtext_load_file

In the case of a large PGSS_TEXT_FILE, the work time of the qtext_load_file
function will be quite long, and the query to the pg_stat_statements table
will not be cancellable, as there is no CHECK_FOR_INTERRUPT in the function.

Also, the amount of bytes read can reach 1 GB, which leads to a slow read
system call that does not allow cancellation of the query. Testing the speed
of sequential read using fio with different block sizes shows that there is
no significant difference between 16 MB blocks and 1 GB blocks.

Therefore, this patch changes the maximum read value from 1 GB to 16 MB and
adds INTERRUPTS_PENDING_CONDITION() check in the read loop of qtext_load_file to make it cancellable.
For now, only statement execution is cancellable (fail_on_interrupt is true only for calls from pg_stat_statements_internal)

Signed-off-by: rkhapov <r.khapov@ya.ru>
---
 .../pg_stat_statements/pg_stat_statements.c   | 29 ++++++++++++++-----
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/contrib/pg_stat_statements/pg_stat_statements.c b/contrib/pg_stat_statements/pg_stat_statements.c
index 9778407cba3..587d49f4330 100644
--- a/contrib/pg_stat_statements/pg_stat_statements.c
+++ b/contrib/pg_stat_statements/pg_stat_statements.c
@@ -365,7 +365,7 @@ static pgssEntry *entry_alloc(pgssHashKey *key, Size query_offset, int query_len
 static void entry_dealloc(void);
 static bool qtext_store(const char *query, int query_len,
 						Size *query_offset, int *gc_count);
-static char *qtext_load_file(Size *buffer_size);
+static char *qtext_load_file(Size *buffer_size, bool fail_on_interrupts);
 static char *qtext_fetch(Size query_offset, int query_len,
 						 char *buffer, Size buffer_size);
 static bool need_gc_qtexts(void);
@@ -767,7 +767,7 @@ pgss_shmem_shutdown(int code, Datum arg)
 	if (fwrite(&num_entries, sizeof(int32), 1, file) != 1)
 		goto error;
 
-	qbuffer = qtext_load_file(&qbuffer_size);
+	qbuffer = qtext_load_file(&qbuffer_size, false /* fail_on_interrupts */ );
 	if (qbuffer == NULL)
 		goto error;
 
@@ -1767,7 +1767,7 @@ pg_stat_statements_internal(FunctionCallInfo fcinfo,
 
 		/* No point in loading file now if there are active writers */
 		if (n_writers == 0)
-			qbuffer = qtext_load_file(&qbuffer_size);
+			qbuffer = qtext_load_file(&qbuffer_size, true /* fail_on_interrupts */ );
 	}
 
 	/*
@@ -1800,7 +1800,7 @@ pg_stat_statements_internal(FunctionCallInfo fcinfo,
 			pgss->gc_count != gc_count)
 		{
 			free(qbuffer);
-			qbuffer = qtext_load_file(&qbuffer_size);
+			qbuffer = qtext_load_file(&qbuffer_size, true /* fail_on_interrupts */ );
 		}
 	}
 
@@ -1819,6 +1819,12 @@ pg_stat_statements_internal(FunctionCallInfo fcinfo,
 		memset(values, 0, sizeof(values));
 		memset(nulls, 0, sizeof(nulls));
 
+		/* Can't process interrupts here - pgss-lock is acquired */
+		if (INTERRUPTS_PENDING_CONDITION())
+		{
+			break;
+		}
+
 		values[i++] = ObjectIdGetDatum(entry->key.userid);
 		values[i++] = ObjectIdGetDatum(entry->key.dbid);
 		if (api_version >= PGSS_V1_9)
@@ -2014,6 +2020,8 @@ pg_stat_statements_internal(FunctionCallInfo fcinfo,
 
 	LWLockRelease(pgss->lock);
 
+	CHECK_FOR_INTERRUPTS();
+
 	free(qbuffer);
 }
 
@@ -2312,7 +2320,7 @@ error:
  * the caller is responsible for verifying that the result is sane.
  */
 static char *
-qtext_load_file(Size *buffer_size)
+qtext_load_file(Size *buffer_size, bool fail_on_interrupts)
 {
 	char	   *buf;
 	int			fd;
@@ -2365,7 +2373,14 @@ qtext_load_file(Size *buffer_size)
 	nread = 0;
 	while (nread < stat.st_size)
 	{
-		int			toread = Min(1024 * 1024 * 1024, stat.st_size - nread);
+		int			toread = Min(32 * 1024 * 1024, stat.st_size - nread);
+
+		if (fail_on_interrupts && INTERRUPTS_PENDING_CONDITION())
+		{
+			free(buf);
+			CloseTransientFile(fd);
+			return NULL;
+		}
 
 		/*
 		 * If we get a short read and errno doesn't get set, the reason is
@@ -2502,7 +2517,7 @@ gc_qtexts(void)
 	 * file is only going to get bigger; hoping for a future non-OOM result is
 	 * risky and can easily lead to complete denial of service.
 	 */
-	qbuffer = qtext_load_file(&qbuffer_size);
+	qbuffer = qtext_load_file(&qbuffer_size, false /* fail_on_interrupts */ );
 	if (qbuffer == NULL)
 		goto gc_fail;
 
-- 
2.39.5 (Apple Git-154)