Fix concurrent function stats drop race

src/backend/commands/functioncmds.c

@@ -67,6 +67,7 @@
 #include "utils/acl.h"
 #include "utils/builtins.h"
 #include "utils/guc.h"
+#include "utils/injection_point.h"
 #include "utils/lsyscache.h"
 #include "utils/rel.h"
 #include "utils/snapmgr.h"
@@ -2315,6 +2316,7 @@ ExecuteCallStmt(CallStmt *stmt, ParamListInfo params, bool atomic, DestReceiver
 		PopActiveSnapshot();
 
 	/* Here we actually call the procedure */
+	INJECTION_POINT("function-call-before-pgstat-init", NULL);
 	pgstat_init_function_usage(fcinfo, &fcusage);
 	retval = FunctionCallInvoke(fcinfo);
 	pgstat_end_function_usage(&fcusage, true);
@@ -2428,4 +2430,4 @@ CallStmtResultDesc(CallStmt *stmt)
 	}
 
 	return tupdesc;
-}
+}

src/backend/utils/activity/pgstat_function.c

@@ -18,6 +18,7 @@
 #include "postgres.h"
 
 #include "fmgr.h"
+#include "utils/injection_point.h"
 #include "utils/inval.h"
 #include "utils/pgstat_internal.h"
 #include "utils/syscache.h"
@@ -112,8 +113,10 @@ pgstat_init_function_usage(FunctionCallInfo fcinfo,
 		AcceptInvalidationMessages();
 		if (!SearchSysCacheExists1(PROCOID, ObjectIdGetDatum(fcinfo->flinfo->fn_oid)))
 		{
-			pgstat_drop_entry(PGSTAT_KIND_FUNCTION, MyDatabaseId,
-							  fcinfo->flinfo->fn_oid);
+			INJECTION_POINT("function-call-before-dropped-stats-drop", NULL);
+			if (!pgstat_drop_entry(PGSTAT_KIND_FUNCTION, MyDatabaseId,
+								   fcinfo->flinfo->fn_oid))
+				pgstat_request_entry_refs_gc();
 			ereport(ERROR, errcode(ERRCODE_UNDEFINED_FUNCTION),
 					errmsg("function call to dropped function"));
 		}
@@ -246,4 +249,4 @@ pgstat_fetch_stat_funcentry(Oid func_id)
 {
 	return (PgStat_StatFuncEntry *)
 		pgstat_fetch_entry(PGSTAT_KIND_FUNCTION, MyDatabaseId, func_id, NULL);
-}
+}

src/backend/utils/activity/pgstat_shmem.c

@@ -918,13 +918,18 @@ pgstat_drop_entry_internal(PgStatShared_HashEntry *shent,
 	 * backends to release their references.
 	 */
 	if (shent->dropped)
-		elog(ERROR,
-			 "trying to drop stats entry already dropped: kind=%s dboid=%u objid=%" PRIu64 " refcount=%u generation=%u",
-			 pgstat_get_kind_info(shent->key.kind)->name,
-			 shent->key.dboid,
-			 shent->key.objid,
-			 pg_atomic_read_u32(&shent->refcount),
-			 pg_atomic_read_u32(&shent->generation));
+	{
+		/*
+		 * Stats drops can race with non-transactional drops of already-stale
+		 * entries, for example when function stats are created for a function
+		 * whose catalog tuple was concurrently removed.  The first drop already
+		 * consumed the valid-entry refcount; treat later drops as requests to
+		 * wait for remaining local references to go away.
+		 */
+		if (!hstat)
+			dshash_release_lock(pgStatLocal.shared_hash, shent);
+		return false;
+	}
 	shent->dropped = true;
 
 	/* release refcount marking entry as not dropped */
@@ -1188,4 +1193,4 @@ pgstat_setup_memcxt(void)
 			AllocSetContextCreate(TopMemoryContext,
 								  "PgStat Shared Ref Hash",
 								  ALLOCSET_SMALL_SIZES);
-}
+}

src/backend/utils/activity/pgstat_xact.c

@@ -14,6 +14,7 @@
 
 #include "access/xact.h"
 #include "pgstat.h"
+#include "utils/injection_point.h"
 #include "utils/memutils.h"
 #include "utils/pgstat_internal.h"
 
@@ -85,6 +86,8 @@ AtEOXact_PgStat_DroppedStats(PgStat_SubXactStatus *xact_state, bool isCommit)
 			 * Transaction that dropped an object committed. Drop the stats
 			 * too.
 			 */
+			if (it->kind == PGSTAT_KIND_FUNCTION)
+				INJECTION_POINT("pgstat-before-drop-function-stats", NULL);
 			if (!pgstat_drop_entry(it->kind, it->dboid, objid))
 				not_freed_count++;
 		}
@@ -384,4 +387,4 @@ void
 pgstat_drop_transactional(PgStat_Kind kind, Oid dboid, uint64 objid)
 {
 	create_drop_transactional_internal(kind, dboid, objid, /* create */ false);
-}
+}

src/test/modules/injection_points/Makefile

@@ -19,7 +19,8 @@ ISOLATION = basic \
 	    repack_temporal_multirange \
 	    repack_toast \
 	    syscache-update-pruned \
-	    heap_lock_update
+	    heap_lock_update \
+	    function-stats-drop
 
 # some isolation tests require wal_level=replica
 ISOLATION_OPTS = --temp-config $(top_srcdir)/src/test/modules/injection_points/extra.conf
@@ -50,4 +51,4 @@ check:
 	@echo "injection points are disabled in this build"
 endif
 
-endif
+endif

src/test/modules/injection_points/expected/function-stats-drop.out

@@ -0,0 +1,45 @@
+Parsed test spec with 4 sessions
+
+starting permutation: call_proc drop_proc wait_drop_stop wake_call_start wait_call_drop_stop fetch_stats wake_call_drop wake_drop
+step call_proc: CALL proc_test(); <waiting ...>
+step drop_proc: DROP PROCEDURE proc_test(); <waiting ...>
+step wait_drop_stop: 
+	SELECT wait_for_injection_point(
+		'function-stats-drop-drop',
+		'pgstat-before-drop-function-stats');
+
+wait_for_injection_point
+------------------------
+
+(1 row)
+
+step wake_call_start: 
+	SELECT FROM injection_points_wakeup('function-call-before-pgstat-init');
+
+step wait_call_drop_stop: 
+	SELECT wait_for_injection_point(
+		'function-stats-drop-call',
+		'function-call-before-dropped-stats-drop');
+
+wait_for_injection_point
+------------------------
+
+(1 row)
+
+step fetch_stats: 
+	SELECT pg_stat_get_function_calls(oid) FROM proc_test_oid;
+
+pg_stat_get_function_calls
+--------------------------
+                         0
+(1 row)
+
+step wake_call_drop: 
+	SELECT FROM injection_points_wakeup('function-call-before-dropped-stats-drop');
+
+step call_proc: <... completed>
+ERROR:  function call to dropped function
+step wake_drop: 
+	SELECT FROM injection_points_wakeup('pgstat-before-drop-function-stats');
+
+step drop_proc: <... completed>

src/test/modules/injection_points/meson.build

@@ -51,6 +51,7 @@ tests += {
       'repack_toast',
       'syscache-update-pruned',
       'heap_lock_update',
+      'function-stats-drop',
     ],
     'runningcheck': false, # see syscache-update-pruned
     # Some tests wait for all snapshots, so avoid parallel execution
@@ -60,4 +61,4 @@ tests += {
       '--temp-config', files('extra.conf'),
     ],
   },
-}
+}

src/test/modules/injection_points/specs/function-stats-drop.spec

@@ -0,0 +1,103 @@
+# Test concurrent function stats drops.
+#
+# A CALL can resolve a procedure OID, then observe invalidations after a
+# concurrent DROP PROCEDURE commits.  With track_functions enabled this creates
+# a function stats entry for the already-dropped OID and immediately drops it.
+# If another backend keeps a local reference to that stats entry, the DROP
+# transaction's post-commit stats cleanup used to ERROR while in COMMIT state.
+
+setup
+{
+	CREATE EXTENSION injection_points;
+	CREATE FUNCTION wait_for_injection_point(appname text, event text)
+	RETURNS void LANGUAGE plpgsql AS $$
+	DECLARE
+		stop_at timestamptz := clock_timestamp() + interval '10 seconds';
+	BEGIN
+		LOOP
+			IF EXISTS (
+				SELECT FROM pg_stat_activity
+				WHERE application_name = appname
+				  AND wait_event_type = 'InjectionPoint'
+				  AND wait_event = event
+			) THEN
+				RETURN;
+			END IF;
+
+			IF clock_timestamp() > stop_at THEN
+				RAISE EXCEPTION 'timed out waiting for % at injection point %',
+					appname, event;
+			END IF;
+
+			PERFORM pg_sleep(0.01);
+		END LOOP;
+	END
+	$$;
+	CREATE PROCEDURE proc_test() LANGUAGE plpgsql AS $$ BEGIN END $$;
+	CREATE TABLE proc_test_oid AS
+		SELECT 'proc_test()'::regprocedure::oid AS oid;
+}
+
+teardown
+{
+	DROP TABLE proc_test_oid;
+	DROP FUNCTION wait_for_injection_point(text, text);
+	DROP EXTENSION injection_points;
+}
+
+session call_s
+setup
+{
+	SET application_name = 'function-stats-drop-call';
+	SET track_functions = 'all';
+	SELECT FROM injection_points_set_local();
+	SELECT FROM injection_points_attach('function-call-before-pgstat-init', 'wait');
+	SELECT FROM injection_points_attach('function-call-before-dropped-stats-drop', 'wait');
+}
+step call_proc	{ CALL proc_test(); }
+
+session drop_s
+setup
+{
+	SET application_name = 'function-stats-drop-drop';
+	SELECT FROM injection_points_set_local();
+	SELECT FROM injection_points_attach('pgstat-before-drop-function-stats', 'wait');
+}
+step drop_proc	{ DROP PROCEDURE proc_test(); }
+
+# Hold a local reference to the transient function stats entry.
+session stats_s
+step fetch_stats	{
+	SELECT pg_stat_get_function_calls(oid) FROM proc_test_oid;
+}
+
+session ctl_s
+step wait_drop_stop	{
+	SELECT wait_for_injection_point(
+		'function-stats-drop-drop',
+		'pgstat-before-drop-function-stats');
+}
+step wake_call_start	{
+	SELECT FROM injection_points_wakeup('function-call-before-pgstat-init');
+}
+step wait_call_drop_stop	{
+	SELECT wait_for_injection_point(
+		'function-stats-drop-call',
+		'function-call-before-dropped-stats-drop');
+}
+step wake_call_drop	{
+	SELECT FROM injection_points_wakeup('function-call-before-dropped-stats-drop');
+}
+step wake_drop	{
+	SELECT FROM injection_points_wakeup('pgstat-before-drop-function-stats');
+}
+
+permutation
+	call_proc
+	drop_proc
+	wait_drop_stop
+	wake_call_start
+	wait_call_drop_stop
+	fetch_stats
+	wake_call_drop
+	wake_drop