On 28.01.2025 02:00, Sami Imseih wrote:
Hi,

I have a few comments about the patch.

1/
tests are missing for pg_stat_statements.track=all and with a sample
rate < 1 applied.
Can we add tests for this case?


I think I’ve come up with a good idea for implementing this. We can create a new file, sampling.sql, for regression tests, where we move written tests in select.out and check test with setting "pg_stat_statements.track = 'all'" and "pg_stat_statements.sample_rate = 0.5". Additionally, we can create two files in the 'expected' directory: sampling.out and sampling_1.out where we can verify either all nested statements are tracked or none of them.



2/

This declaration:

+static bool current_query_sampled = false;

should be moved near the declaration of nesting_level,
the other local variable.


+1




3/
I do not really like the calls to update_current_query_sampled()
the way there are, and I think the pgss_enabled should encapsulate
the new sample rate check as well.

My thoughts are:

1/ change the name of the function from update_current_query_sampled
to current_query_sampled

Also, the static bool should be called is_query_sampled or something like that.

The function should also allow you to skip the sample rate check, such as if
called from pgss_post_parse_analyze.

We can also remove the IsParallelWorker() check in this case, since that
is done in pgss_enabled.

something like this is what I am thinking:

static bool
current_query_sampled(bool skip)
{
        if (skip)
                return true;

         if (nesting_level == 0)
         {
                is_query_sampled = pgss_sample_rate != 0.0 &&
                                                (pgss_sample_rate == 1.0 ||

pg_prng_double(&pg_global_prng_state) <= pgss_sample_rate);
         }

        return is_query_sampled;
}

#define pgss_enabled(level, skip_sampling_check) \
         (!IsParallelWorker() && \
         (pgss_track == PGSS_TRACK_ALL || \
        (pgss_track == PGSS_TRACK_TOP && (level) == 0)) && \
        current_query_sampled(skip_sampling_check))

What do you think?


I agree with all of these changes; they make the code more readable. Adding comments to the current_query_sampled() function will make it even clearer.



4/ Now we have pg_stat_statements.track = "off" which is effectively
the same thing as pg_stat_statements.sample_rate = "0". Does this need
to be called out in docs?


+1, because we have the same thing in log_statement_sample_rate.


All the changes mentioned above are included in the v13 patch. Since the patch status is 'Ready for Committer,' I believe it is now better for upstream inclusion, with improved details in tests and documentation. Do you have any further suggestions?

--
Best regards,
Ilia Evdokimov,
Tantor Labs LLC.
From 149688fe3bccce25f8d01dfe8c01444555dce03c Mon Sep 17 00:00:00 2001
From: Ilia Evdokimov <ilya.evdoki...@tantorlabs.com>
Date: Tue, 28 Jan 2025 05:31:20 +0300
Subject: [PATCH] Allow setting sample rate for pg_stat_statements

New configuration parameter pg_stat_statements.sample_rate makes it
possible to track just a fraction of the queries meeting the configured
threshold, to reduce the amount of tracking.
---
 contrib/pg_stat_statements/Makefile           |  2 +-
 .../pg_stat_statements/expected/sampling.out  | 99 +++++++++++++++++++
 .../expected/sampling_1.out                   | 97 ++++++++++++++++++
 .../pg_stat_statements/pg_stat_statements.c   | 59 +++++++++--
 contrib/pg_stat_statements/sql/sampling.sql   | 34 +++++++
 doc/src/sgml/pgstatstatements.sgml            | 19 ++++
 6 files changed, 302 insertions(+), 8 deletions(-)
 create mode 100644 contrib/pg_stat_statements/expected/sampling.out
 create mode 100644 contrib/pg_stat_statements/expected/sampling_1.out
 create mode 100644 contrib/pg_stat_statements/sql/sampling.sql

diff --git a/contrib/pg_stat_statements/Makefile b/contrib/pg_stat_statements/Makefile
index 241c02587b..b70bdfaf2d 100644
--- a/contrib/pg_stat_statements/Makefile
+++ b/contrib/pg_stat_statements/Makefile
@@ -20,7 +20,7 @@ LDFLAGS_SL += $(filter -lm, $(LIBS))
 REGRESS_OPTS = --temp-config $(top_srcdir)/contrib/pg_stat_statements/pg_stat_statements.conf
 REGRESS = select dml cursors utility level_tracking planning \
 	user_activity wal entry_timestamp privileges extended \
-	parallel cleanup oldextversions
+	parallel sampling cleanup oldextversions
 # Disabled because these tests require "shared_preload_libraries=pg_stat_statements",
 # which typical installcheck users do not have (e.g. buildfarm clients).
 NO_INSTALLCHECK = 1
diff --git a/contrib/pg_stat_statements/expected/sampling.out b/contrib/pg_stat_statements/expected/sampling.out
new file mode 100644
index 0000000000..6e73a0a318
--- /dev/null
+++ b/contrib/pg_stat_statements/expected/sampling.out
@@ -0,0 +1,99 @@
+--
+-- sample statements
+--
+SET pg_stat_statements.sample_rate = 0.0;
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+SELECT 1 AS "int";
+ int 
+-----
+   1
+(1 row)
+
+SELECT query, calls FROM pg_stat_statements ORDER BY query COLLATE "C";
+ query | calls 
+-------+-------
+(0 rows)
+
+SET pg_stat_statements.sample_rate = 1.0;
+SELECT 1 AS "int";
+ int 
+-----
+   1
+(1 row)
+
+SELECT query, calls FROM pg_stat_statements ORDER BY query COLLATE "C";
+       query        | calls 
+--------------------+-------
+ SELECT $1 AS "int" |     1
+(1 row)
+
+-- sample_rate is (0, 1) and track = "all"
+SET pg_stat_statements.track = "all";
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+SET pg_stat_statements.sample_rate = 0.5;
+-- either all nested statements will be tracked (sample.out) or none will (sample_1.out).
+EXPLAIN (COSTS OFF) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+    ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | EXPLAIN (COSTS OFF) SELECT $1
+ f        |     1 | SELECT $1
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | SET pg_stat_statements.sample_rate = $1
+(4 rows)
+
+SET pg_stat_statements.track = "top";
+SET pg_stat_statements.sample_rate = 0.0;
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+SELECT 1 \parse stmt
+\bind_named stmt \g
+ ?column? 
+----------
+        1
+(1 row)
+
+SELECT query, calls FROM pg_stat_statements ORDER BY query COLLATE "C";
+ query | calls 
+-------+-------
+(0 rows)
+
+SET pg_stat_statements.sample_rate = 1.0;
+\bind_named stmt \g
+ ?column? 
+----------
+        1
+(1 row)
+
+SELECT query, calls FROM pg_stat_statements ORDER BY query COLLATE "C";
+   query   | calls 
+-----------+-------
+ SELECT $1 |     1
+(1 row)
+
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
diff --git a/contrib/pg_stat_statements/expected/sampling_1.out b/contrib/pg_stat_statements/expected/sampling_1.out
new file mode 100644
index 0000000000..3a390226fa
--- /dev/null
+++ b/contrib/pg_stat_statements/expected/sampling_1.out
@@ -0,0 +1,97 @@
+--
+-- sample statements
+--
+SET pg_stat_statements.sample_rate = 0.0;
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+SELECT 1 AS "int";
+ int 
+-----
+   1
+(1 row)
+
+SELECT query, calls FROM pg_stat_statements ORDER BY query COLLATE "C";
+ query | calls 
+-------+-------
+(0 rows)
+
+SET pg_stat_statements.sample_rate = 1.0;
+SELECT 1 AS "int";
+ int 
+-----
+   1
+(1 row)
+
+SELECT query, calls FROM pg_stat_statements ORDER BY query COLLATE "C";
+       query        | calls 
+--------------------+-------
+ SELECT $1 AS "int" |     1
+(1 row)
+
+-- sample_rate is (0, 1) and track = "all"
+SET pg_stat_statements.track = "all";
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+SET pg_stat_statements.sample_rate = 0.5;
+-- either all nested statements will be tracked (sample.out) or none will (sample_1.out).
+EXPLAIN (COSTS OFF) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+    ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | SET pg_stat_statements.sample_rate = $1
+(2 rows)
+
+SET pg_stat_statements.track = "top";
+SET pg_stat_statements.sample_rate = 0.0;
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+SELECT 1 \parse stmt
+\bind_named stmt \g
+ ?column? 
+----------
+        1
+(1 row)
+
+SELECT query, calls FROM pg_stat_statements ORDER BY query COLLATE "C";
+ query | calls 
+-------+-------
+(0 rows)
+
+SET pg_stat_statements.sample_rate = 1.0;
+\bind_named stmt \g
+ ?column? 
+----------
+        1
+(1 row)
+
+SELECT query, calls FROM pg_stat_statements ORDER BY query COLLATE "C";
+   query   | calls 
+-----------+-------
+ SELECT $1 |     1
+(1 row)
+
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
diff --git a/contrib/pg_stat_statements/pg_stat_statements.c b/contrib/pg_stat_statements/pg_stat_statements.c
index bebf8134eb..afd0537713 100644
--- a/contrib/pg_stat_statements/pg_stat_statements.c
+++ b/contrib/pg_stat_statements/pg_stat_statements.c
@@ -50,6 +50,7 @@
 #include "access/parallel.h"
 #include "catalog/pg_authid.h"
 #include "common/int.h"
+#include "common/pg_prng.h"
 #include "executor/instrument.h"
 #include "funcapi.h"
 #include "jit/jit.h"
@@ -256,6 +257,9 @@ typedef struct pgssSharedState
 /* Current nesting depth of planner/ExecutorRun/ProcessUtility calls */
 static int	nesting_level = 0;
 
+/* Is the current top-level query to be sampled? */
+static bool is_query_sampled = false;
+
 /* Saved hook values */
 static shmem_request_hook_type prev_shmem_request_hook = NULL;
 static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
@@ -294,12 +298,14 @@ static bool pgss_track_utility = true;	/* whether to track utility commands */
 static bool pgss_track_planning = false;	/* whether to track planning
 											 * duration */
 static bool pgss_save = true;	/* whether to save stats across shutdown */
+static double pgss_sample_rate = 1.0; /* fraction of statements to track */
 
 
-#define pgss_enabled(level) \
+#define pgss_enabled(level, skip_sampling_check) \
 	(!IsParallelWorker() && \
 	(pgss_track == PGSS_TRACK_ALL || \
-	(pgss_track == PGSS_TRACK_TOP && (level) == 0)))
+	(pgss_track == PGSS_TRACK_TOP && (level) == 0)) && \
+	current_query_sampled(skip_sampling_check))
 
 #define record_gc_qtexts() \
 	do { \
@@ -373,6 +379,7 @@ static char *generate_normalized_query(JumbleState *jstate, const char *query,
 static void fill_in_constant_lengths(JumbleState *jstate, const char *query,
 									 int query_loc);
 static int	comp_location(const void *a, const void *b);
+static bool	current_query_sampled(bool skip);
 
 
 /*
@@ -414,6 +421,19 @@ _PG_init(void)
 							NULL,
 							NULL);
 
+	DefineCustomRealVariable("pg_stat_statements.sample_rate",
+							 "Fraction of queries to track.",
+							 NULL,
+							 &pgss_sample_rate,
+							 1.0,
+							 0.0,
+							 1.0,
+							 PGC_SUSET,
+							 0,
+							 NULL,
+							 NULL,
+							 NULL);
+
 	DefineCustomEnumVariable("pg_stat_statements.track",
 							 "Selects which statements are tracked by pg_stat_statements.",
 							 NULL,
@@ -836,7 +856,7 @@ pgss_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jstate)
 		prev_post_parse_analyze_hook(pstate, query, jstate);
 
 	/* Safety check... */
-	if (!pgss || !pgss_hash || !pgss_enabled(nesting_level))
+	if (!pgss || !pgss_hash || !pgss_enabled(nesting_level, true))
 		return;
 
 	/*
@@ -894,7 +914,7 @@ pgss_planner(Query *parse,
 	 * pgss_store needs it.  We also ignore query without queryid, as it would
 	 * be treated as a utility statement, which may not be the case.
 	 */
-	if (pgss_enabled(nesting_level)
+	if (pgss_enabled(nesting_level, false)
 		&& pgss_track_planning && query_string
 		&& parse->queryId != UINT64CONST(0))
 	{
@@ -999,7 +1019,8 @@ pgss_ExecutorStart(QueryDesc *queryDesc, int eflags)
 	 * counting of optimizable statements that are directly contained in
 	 * utility statements.
 	 */
-	if (pgss_enabled(nesting_level) && queryDesc->plannedstmt->queryId != UINT64CONST(0))
+	if (pgss_enabled(nesting_level, false) &&
+		queryDesc->plannedstmt->queryId != UINT64CONST(0))
 	{
 		/*
 		 * Set up to track total elapsed time in ExecutorRun.  Make sure the
@@ -1068,7 +1089,7 @@ pgss_ExecutorEnd(QueryDesc *queryDesc)
 	uint64		queryId = queryDesc->plannedstmt->queryId;
 
 	if (queryId != UINT64CONST(0) && queryDesc->totaltime &&
-		pgss_enabled(nesting_level))
+		pgss_enabled(nesting_level, false))
 	{
 		/*
 		 * Make sure stats accumulation is done.  (Note: it's okay if several
@@ -1111,7 +1132,7 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 	uint64		saved_queryId = pstmt->queryId;
 	int			saved_stmt_location = pstmt->stmt_location;
 	int			saved_stmt_len = pstmt->stmt_len;
-	bool		enabled = pgss_track_utility && pgss_enabled(nesting_level);
+	bool		enabled = pgss_track_utility && pgss_enabled(nesting_level, false);
 
 	/*
 	 * Force utility statements to get queryId zero.  We do this even in cases
@@ -3011,3 +3032,27 @@ comp_location(const void *a, const void *b)
 
 	return pg_cmp_s32(l, r);
 }
+
+/*
+ * Determine whether the current query should be sampled.
+ *
+ * In some cases, sampling may need to be skipped entirely, such as during
+ * query parsing for normalization.
+ *
+ * At the beginning of each top-level statement, decide whether we'll
+ * sample this statement. If nested-statement tracking is enabled,
+ * either all nested statements will be tracked or none will.
+ */
+static bool
+current_query_sampled(bool skip)
+{
+	if (skip)
+		return true;
+
+	if (nesting_level == 0)
+		is_query_sampled = pgss_sample_rate != 0.0 &&
+			(pgss_sample_rate == 1.0 ||
+			pg_prng_double(&pg_global_prng_state) <= pgss_sample_rate);
+
+	return is_query_sampled;
+}
\ No newline at end of file
diff --git a/contrib/pg_stat_statements/sql/sampling.sql b/contrib/pg_stat_statements/sql/sampling.sql
new file mode 100644
index 0000000000..4dabe1540c
--- /dev/null
+++ b/contrib/pg_stat_statements/sql/sampling.sql
@@ -0,0 +1,34 @@
+--
+-- sample statements
+--
+
+SET pg_stat_statements.sample_rate = 0.0;
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+SELECT 1 AS "int";
+SELECT query, calls FROM pg_stat_statements ORDER BY query COLLATE "C";
+SET pg_stat_statements.sample_rate = 1.0;
+SELECT 1 AS "int";
+SELECT query, calls FROM pg_stat_statements ORDER BY query COLLATE "C";
+
+
+-- sample_rate is (0, 1) and track = "all"
+SET pg_stat_statements.track = "all";
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+SET pg_stat_statements.sample_rate = 0.5;
+
+-- either all nested statements will be tracked (sample.out) or none will (sample_1.out).
+EXPLAIN (COSTS OFF) SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+    ORDER BY query COLLATE "C", toplevel;
+SET pg_stat_statements.track = "top";
+
+
+SET pg_stat_statements.sample_rate = 0.0;
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+SELECT 1 \parse stmt
+\bind_named stmt \g
+SELECT query, calls FROM pg_stat_statements ORDER BY query COLLATE "C";
+SET pg_stat_statements.sample_rate = 1.0;
+\bind_named stmt \g
+SELECT query, calls FROM pg_stat_statements ORDER BY query COLLATE "C";
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
\ No newline at end of file
diff --git a/doc/src/sgml/pgstatstatements.sgml b/doc/src/sgml/pgstatstatements.sgml
index 501b468e9a..cd3fdaeffe 100644
--- a/doc/src/sgml/pgstatstatements.sgml
+++ b/doc/src/sgml/pgstatstatements.sgml
@@ -936,6 +936,25 @@
     </listitem>
    </varlistentry>
 
+   <varlistentry>
+    <term>
+     <varname>pg_stat_statements.sample_rate</varname> (<type>real</type>)
+     <indexterm>
+      <primary><varname>pg_stat_statements.sample_rate</varname> configuration parameter</primary>
+     </indexterm>
+    </term>
+
+    <listitem>
+     <para>
+      <varname>pg_stat_statements.sample_rate</varname> causes pg_stat_statements to only
+      track a fraction of the statements in each session. The default is <literal>1</literal>,
+      meaning track all the queries. Setting this to <literal>0</literal> disables sampled statements
+      tracking, the same as setting <varname>pg_stat_statements.track</varname> to <literal>none</literal>.
+      In case of nested statements, either all will be tracked or none. Only superusers can change this setting.
+     </para>
+    </listitem>
+   </varlistentry>
+
    <varlistentry>
     <term>
      <varname>pg_stat_statements.save</varname> (<type>boolean</type>)
-- 
2.34.1

Reply via email to