Protect referenced PLpgSQL_function structure until fextra is used

Author: okbob

src/cursors_leaks.c

@@ -33,7 +33,7 @@ int			plpgsql_check_cursors_leaks_level = WARNING;
 typedef struct CursorTrace
 {
 	int			stmtid;
-	int			rec_level;
+	PLpgSQL_execstate *estate;
 	char	   *curname;
 } CursorTrace;
 
@@ -195,16 +195,7 @@ func_end(PLpgSQL_execstate *estate,
 		 * Iterate over traced cursors. Remove slots for tracing immediately,
 		 * when traced cursor is closed already.
 		 */
-#if PG_VERSION_NUM >= 180000
-
-		if (ct->curname && ct->rec_level == func->cfunc.use_count)
-
-#else
-
-		if (ct->curname && ct->rec_level == func->use_count)
-
-#endif
-
+		if (ct->curname && ct->estate == estate)
 		{
 			if (SPI_cursor_find(ct->curname))
 			{
@@ -286,16 +277,7 @@ stmt_end(PLpgSQL_execstate *estate, PLpgSQL_stmt *stmt, plch_fextra *fextra)
 
 				if (SPI_cursor_find(ct->curname))
 				{
-#if PG_VERSION_NUM >= 180000
-
-					if (estate->func->cfunc.use_count == 1 && !plpgsql_check_cursors_leaks_strict)
-
-#else
-
-					if (estate->func->use_count == 1 && !plpgsql_check_cursors_leaks_strict)
-
-#endif
-
+					if (!plpgsql_check_cursors_leaks_strict)
 					{
 						char	   *context;
 
@@ -362,16 +344,7 @@ stmt_end(PLpgSQL_execstate *estate, PLpgSQL_stmt *stmt, plch_fextra *fextra)
 
 			ct->stmtid = stmt->stmtid;
 
-#if PG_VERSION_NUM >= 180000
-
-			ct->rec_level = estate->func->cfunc.use_count;
-
-#else
-
-			ct->rec_level = estate->func->use_count;
-
-#endif
-
+			ct->estate = estate;
 			ct->curname = pstrdup(curname);
 
 			MemoryContextSwitchTo(oldcxt);

src/fextra_cache.c

@@ -321,6 +321,17 @@ plch_get_fextra(PLpgSQL_function *func)
 	}
 
 	fextra->use_count++;
+	fextra->func = func;
+
+#if PG_VERSION_NUM >= 180000
+
+	func->cfunc.use_count++;
+
+#else
+
+	func->use_count++;
+
+#endif
 
 	return fextra;
 }
@@ -330,7 +341,21 @@ plch_release_fextra(plch_fextra *fextra)
 {
 	Assert(fextra->use_count > 0);
 
+	/* until now, referenced PLpgSQL_function should be still valid */
+	Assert(fextra->hk.fn_oid == fextra->func->fn_oid);
+
 	fextra->use_count--;
+
+#if PG_VERSION_NUM >= 180000
+
+	fextra->func->cfunc.use_count--;
+
+#else
+
+	fextra->func->use_count--;
+
+#endif
+
 }
 
 

src/pldbgapi3.c

@@ -44,7 +44,6 @@ typedef struct plpgsql_plugin_info
 	/* for assertations */
 	Oid			fn_oid;
 	PLpgSQL_execstate *estate;
-	int			use_count;
 
 	plch_fextra *fextra;
 
@@ -160,9 +159,12 @@ plugin_info_reset(void *arg)
 {
 	plpgsql_plugin_info *plugin_info = (plpgsql_plugin_info*) arg;
 	MemoryContext exec_mcxt = CurrentMemoryContext;
-	int			stmts_stack_size = plugin_info->stmts_stack_size;
+	int			stmts_stack_size;
 	int			i;
 
+	/* The memory should not be corrupted! */
+	Assert(plugin_info->magic == PLUGIN_INFO_MAGIC);
+
 	/*
 	 * PostgreSQL 19 can remove this callback. But we need to support
 	 * previous releases, so when fextra is already released, then
@@ -171,6 +173,7 @@ plugin_info_reset(void *arg)
 	if (!plugin_info->fextra)
 		return;
 
+	stmts_stack_size = plugin_info->stmts_stack_size;
 	plugin_info->stmts_stack_size = 0;
 
 	PG_TRY();
@@ -221,17 +224,6 @@ func_setup(PLpgSQL_execstate *estate, PLpgSQL_function *func)
 	plugin_info->fn_oid = func->fn_oid;
 	plugin_info->estate = estate;
 
-
-#if PG_VERSION_NUM >= 180000
-
-	plugin_info->use_count = func->cfunc.use_count;
-
-#else
-
-	plugin_info->use_count = func->use_count;
-
-#endif
-
 	for (i = 0; i < nplugins; i++)
 	{
 		if (plugins[i]->is_active(estate, func))
@@ -333,16 +325,6 @@ func_beg(PLpgSQL_execstate *estate, PLpgSQL_function *func)
 	Assert(plugin_info->estate == estate);
 	Assert(plugin_info->fn_oid == func->fn_oid);
 
-#if PG_VERSION_NUM >= 180000
-
-	Assert(plugin_info->use_count == func->cfunc.use_count);
-
-#else
-
-	Assert(plugin_info->use_count == func->use_count);
-
-#endif
-
 	PG_TRY();
 	{
 		if (plugin_info->fextra)
@@ -397,16 +379,6 @@ func_end(PLpgSQL_execstate *estate, PLpgSQL_function *func)
 	Assert(plugin_info->estate == estate);
 	Assert(plugin_info->fn_oid == func->fn_oid);
 
-#if PG_VERSION_NUM >= 180000
-
-	Assert(plugin_info->use_count == func->cfunc.use_count);
-
-#else
-
-	Assert(plugin_info->use_count == func->use_count);
-
-#endif
-
 	PG_TRY();
 	{
 		if (plugin_info->fextra)
@@ -481,16 +453,6 @@ stmt_beg(PLpgSQL_execstate *estate, PLpgSQL_stmt *stmt)
 	Assert(plugin_info->estate == estate);
 	Assert(plugin_info->fn_oid == estate->func->fn_oid);
 
-#if PG_VERSION_NUM >= 180000
-
-	Assert(plugin_info->use_count == estate->func->cfunc.use_count);
-
-#else
-
-	Assert(plugin_info->use_count == estate->func->use_count);
-
-#endif
-
 	if (plugin_info->fextra)
 	{
 		if (estate->cur_error)
@@ -572,16 +534,6 @@ stmt_end(PLpgSQL_execstate *estate, PLpgSQL_stmt *stmt)
 	Assert(plugin_info->estate == estate);
 	Assert(plugin_info->fn_oid == estate->func->fn_oid);
 
-#if PG_VERSION_NUM >= 180000
-
-	Assert(plugin_info->use_count == estate->func->cfunc.use_count);
-
-#else
-
-	Assert(plugin_info->use_count == estate->func->use_count);
-
-#endif
-
 	if (plugin_info->fextra)
 	{
 		Assert(plugin_info->stmts_stack_size > 0);

src/plpgsql_check.h

@@ -437,6 +437,7 @@ typedef struct
 {
 	plch_fextra_hk hk;
 	uint32		hashValue;
+	PLpgSQL_function *func;
 
 	Oid			fn_oid;
 	char	   *fn_signature;