Optimized pgstat_init_function_usage/pg_newlocale_from_collation (#188)

* Optimize pgstat_init_function_usage/pg_newlocale_from_collation function in creating index Signed-off-by: Parikshit Sarode <[email protected]> Task: BABEL-4300
babelfish-for-postgresql · Sep 26, 2023 · c7572cd · c7572cd
1 parent b327d84
commit c7572cd
Show file tree

Hide file tree

Showing 5 changed files with 51 additions and 14 deletions.
diff --git a/src/backend/utils/activity/pgstat_function.c b/src/backend/utils/activity/pgstat_function.c
@@ -16,7 +16,7 @@
  */
 
 #include "postgres.h"
-
+#include "pgstat.h"
 #include "access/htup_details.h"
 #include "catalog/pg_proc.h"
 #include "fmgr.h"
@@ -80,18 +80,6 @@ pgstat_init_function_usage(FunctionCallInfo fcinfo,
 	PgStat_BackendFunctionEntry *pending;
 	bool		created_entry;
 
-	if (pre_function_call_hook && IsTransactionState())
-	{
-		HeapTuple proctup = SearchSysCache1(PROCOID, ObjectIdGetDatum(fcinfo->flinfo->fn_oid));
-		if (HeapTupleIsValid(proctup))
-		{
-			Form_pg_proc proc = (Form_pg_proc) GETSTRUCT(proctup);
-			(*pre_function_call_hook)(NameStr(proc->proname));
-		}
-
-		ReleaseSysCache(proctup);
-	}
-
 	if (pgstat_track_functions <= fcinfo->flinfo->fn_stats)
 	{
 		/* stats not wanted */

diff --git a/src/backend/utils/adt/pg_locale.c b/src/backend/utils/adt/pg_locale.c
@@ -102,6 +102,8 @@ char	   *localized_full_months[12 + 1];
 /* is the databases's LC_CTYPE the C locale? */
 bool		database_ctype_is_c = false;
 
+PGDLLIMPORT collation_cache_entry_hook_type collation_cache_entry_hook = NULL;
+
 /* indicates whether locale information cache is valid */
 static bool CurrentLocaleConvValid = false;
 static bool CurrentLCTimeValid = false;
@@ -1521,6 +1523,7 @@ pg_locale_t
 pg_newlocale_from_collation(Oid collid)
 {
 	collation_cache_entry *cache_entry;
+	pg_locale_t 		prev_local = NULL;
 
 	/* Callers must pass a valid OID */
 	Assert(OidIsValid(collid));
@@ -1534,6 +1537,16 @@ pg_newlocale_from_collation(Oid collid)
 	}
 
 	cache_entry = lookup_collation_cache(collid, false);
+
+	/* Call hook to get pervious cached value and return if not null*/
+	if(collation_cache_entry_hook)
+	{
+		prev_local =  (pg_locale_t)(*collation_cache_entry_hook)(collid,NULL);
+		if(prev_local)
+		{
+			return prev_local;
+		}
+	}
 
 	if (cache_entry->locale == 0)
 	{
@@ -1676,6 +1689,12 @@ pg_newlocale_from_collation(Oid collid)
 		cache_entry->locale = resultp;
 	}
 
+/* Call hook to save the cached value */
+	if(collation_cache_entry_hook)
+	{
+		(*collation_cache_entry_hook)(cache_entry->collid, &cache_entry->locale);
+	}	
+
 	return cache_entry->locale;
 }
 

diff --git a/src/backend/utils/fmgr/fmgr.c b/src/backend/utils/fmgr/fmgr.c
@@ -44,6 +44,7 @@ PGDLLIMPORT needs_fmgr_hook_type needs_fmgr_hook = NULL;
 PGDLLIMPORT fmgr_hook_type fmgr_hook = NULL;
 PGDLLIMPORT non_tsql_proc_entry_hook_type non_tsql_proc_entry_hook = NULL;
 PGDLLIMPORT get_func_language_oids_hook_type get_func_language_oids_hook = NULL;
+PGDLLIMPORT pgstat_function_wrapper_hook_type pgstat_function_wrapper_hook = NULL;
 
 /*
  * Hashtable for fast lookup of external C functions
@@ -60,7 +61,6 @@ typedef struct
 
 static HTAB *CFuncHash = NULL;
 
-
 static void fmgr_info_cxt_security(Oid functionId, FmgrInfo *finfo, MemoryContext mcxt,
 								   bool ignore_security);
 static void fmgr_info_C_lang(Oid functionId, FmgrInfo *finfo, HeapTuple procedureTuple);
@@ -706,6 +706,7 @@ fmgr_security_definer(PG_FUNCTION_ARGS)
 	int			sys_func_count = 0;
 	int			non_tsql_proc_count = 0;
 	void	   *newextra = NULL;
+	char 	   *cacheTupleProcname = NULL;
 
 	if (!fcinfo->flinfo->fn_extra)
 	{
@@ -729,6 +730,8 @@ fmgr_security_definer(PG_FUNCTION_ARGS)
 				 fcinfo->flinfo->fn_oid);
 		procedureStruct = (Form_pg_proc) GETSTRUCT(tuple);
 
+		cacheTupleProcname = pstrdup(procedureStruct->proname.data);
+
 		if (procedureStruct->prosecdef)
 			fcache->userid = procedureStruct->proowner;
 
@@ -832,8 +835,25 @@ fmgr_security_definer(PG_FUNCTION_ARGS)
 		fcinfo->flinfo = &fcache->flinfo;
 
 		/* See notes in fmgr_info_cxt_security */
+
+		/* 
+		 * The wrapper function hook is called if hook is not null 
+		 * and the dialect is TSQL then we call this func using hook 
+		 * otherwise we will fall back to pgstat_init_function_usage
+		*/
+
+		if(pgstat_function_wrapper_hook && sql_dialect == SQL_DIALECT_TSQL)
+		{
+			(*pgstat_function_wrapper_hook)(fcinfo, &fcusage, cacheTupleProcname);
+		}
+
 		pgstat_init_function_usage(fcinfo, &fcusage);
 
+		if(cacheTupleProcname)
+		{
+			pfree(cacheTupleProcname);
+		}
+
 		result = FunctionCallInvoke(fcinfo);
 
 		/*

diff --git a/src/include/pgstat.h b/src/include/pgstat.h
@@ -493,6 +493,8 @@ extern void pgstat_drop_function(Oid proid);
 struct FunctionCallInfoBaseData;
 extern void pgstat_init_function_usage(struct FunctionCallInfoBaseData *fcinfo,
 									   PgStat_FunctionCallUsage *fcu);
+extern void pgstat_init_function_usage_wrapper(struct FunctionCallInfoBaseData *fcinfo,
+									   PgStat_FunctionCallUsage *fcusageptr, char *proname);
 extern void pgstat_end_function_usage(PgStat_FunctionCallUsage *fcu,
 									  bool finalize);
 
@@ -720,4 +722,7 @@ extern PGDLLIMPORT guc_newval_hook_type guc_newval_hook;
 typedef bool (*tsql_has_pgstat_permissions_hook_type) (Oid role);
 extern PGDLLIMPORT tsql_has_pgstat_permissions_hook_type tsql_has_pgstat_permissions_hook;
 
+typedef void (*pgstat_function_wrapper_hook_type)(FunctionCallInfo, PgStat_FunctionCallUsage *, char *);
+extern PGDLLIMPORT pgstat_function_wrapper_hook_type pgstat_function_wrapper_hook;
+
 #endif							/* PGSTAT_H */
diff --git a/src/include/utils/pg_locale.h b/src/include/utils/pg_locale.h
@@ -105,6 +105,8 @@ struct pg_locale_struct
 
 typedef struct pg_locale_struct *pg_locale_t;
 
+extern pg_locale_t *collation_cache_entry_hook_function (Oid, pg_locale_t *);
+
 extern PGDLLIMPORT struct pg_locale_struct default_locale;
 
 extern void make_icu_collator(const char *iculocstr,
@@ -126,4 +128,7 @@ extern size_t wchar2char(char *to, const wchar_t *from, size_t tolen,
 extern size_t char2wchar(wchar_t *to, size_t tolen,
 						 const char *from, size_t fromlen, pg_locale_t locale);
 
+typedef pg_locale_t* (*collation_cache_entry_hook_type)(Oid, pg_locale_t *);
+extern  PGDLLIMPORT collation_cache_entry_hook_type collation_cache_entry_hook;
+
 #endif							/* _PG_LOCALE_ */