Use ResourceOwner to track internally prepared statements

XL uses prepared statements internally to execute querys on remote datanodes.
But in case of transaction aborts, these statements are left open on the
datanodes causing issues later. We now track such prepared statements using
transaction resource owner and drop them in case of aborts.

5 files changed, 153 insertions, 5 deletions

diff --git a/src/backend/commands/prepare.c b/src/backend/commands/prepare.c
index 219608b571..8ce63780dd 100644
--- a/src/backend/commands/prepare.c
+++ b/src/backend/commands/prepare.c
@@ -184,7 +184,8 @@ PrepareQuery(PrepareStmt *stmt, const char *queryString)
 	 */
 	StorePreparedStatement(stmt->name,
 						   plansource,
-						   true);
+						   true,
+						   false);
 }
 
 /*
@@ -583,7 +584,8 @@ SetRemoteStatementName(Plan *plan, const char *stmt_name, int num_params,
 void
 StorePreparedStatement(const char *stmt_name,
 					   CachedPlanSource *plansource,
-					   bool from_sql)
+					   bool from_sql,
+					   bool use_resowner)
 {
 	PreparedStatement *entry;
 	TimestampTz cur_ts = GetCurrentStatementStartTimestamp();
@@ -610,9 +612,19 @@ StorePreparedStatement(const char *stmt_name,
 	entry->plansource = plansource;
 	entry->from_sql = from_sql;
 	entry->prepare_time = cur_ts;
+	entry->use_resowner = use_resowner;
 
 	/* Now it's safe to move the CachedPlanSource to permanent memory */
 	SaveCachedPlan(plansource);
+
+#ifdef XCP	
+	if (use_resowner)
+	{
+		ResourceOwnerEnlargePreparedStmts(CurTransactionResourceOwner);
+		ResourceOwnerRememberPreparedStmt(CurTransactionResourceOwner,
+				entry->stmt_name);
+	}
+#endif	
 }
 
 /*
@@ -722,6 +734,11 @@ DropPreparedStatement(const char *stmt_name, bool showError)
 
 		/* Now we can remove the hash table entry */
 		hash_search(prepared_queries, entry->stmt_name, HASH_REMOVE, NULL);
+#ifdef XCP
+		if (entry->use_resowner)
+			ResourceOwnerForgetPreparedStmt(CurTransactionResourceOwner,
+					entry->stmt_name);
+#endif		
 	}
 }
 
@@ -747,6 +764,12 @@ DropAllPreparedStatements(void)
 
 		/* Now we can remove the hash table entry */
 		hash_search(prepared_queries, entry->stmt_name, HASH_REMOVE, NULL);
+
+#ifdef XCP
+		if (entry->use_resowner)
+			ResourceOwnerForgetPreparedStmt(CurTransactionResourceOwner,
+					entry->stmt_name);
+#endif		
 	}
 }
 
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index 9bf91a3e33..465b03fe9b 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -1625,7 +1625,7 @@ exec_parse_message(const char *query_string,	/* string to execute */
 		/*
 		 * Store the query as a prepared statement.
 		 */
-		StorePreparedStatement(stmt_name, psrc, false);
+		StorePreparedStatement(stmt_name, psrc, false, false);
 	}
 	else
 	{
@@ -1768,7 +1768,7 @@ exec_plan_message(const char *query_string,	/* source of the query */
 	/*
 	 * Store the query as a prepared statement.  See above comments.
 	 */
-	StorePreparedStatement(stmt_name, psrc, false);
+	StorePreparedStatement(stmt_name, psrc, false, true);
 
 	SetRemoteSubplan(psrc, plan_string);
 
diff --git a/src/backend/utils/resowner/resowner.c b/src/backend/utils/resowner/resowner.c
index 7c771eb491..8743fff77c 100644
--- a/src/backend/utils/resowner/resowner.c
+++ b/src/backend/utils/resowner/resowner.c
@@ -76,6 +76,13 @@ typedef struct ResourceOwnerData
 	int			nfiles;			/* number of owned temporary files */
 	File	   *files;			/* dynamically allocated array */
 	int			maxfiles;		/* currently allocated array size */
+
+#ifdef XCP	
+	/* We have built-in support for remembering prepared statements */
+	int			nstmts;			/* number of remote statements */
+	char		*stmts;		/* dynamically allocated array */
+	int			maxstmts;		/* currently allocated array size */
+#endif	
 }	ResourceOwnerData;
 
 
@@ -110,6 +117,9 @@ static void PrintPlanCacheLeakWarning(CachedPlan *plan);
 static void PrintTupleDescLeakWarning(TupleDesc tupdesc);
 static void PrintSnapshotLeakWarning(Snapshot snapshot);
 static void PrintFileLeakWarning(File file);
+#ifdef XCP
+static void PrintPreparedStmtLeakWarning(char *stmt);
+#endif
 
 
 /*****************************************************************************
@@ -192,6 +202,10 @@ ResourceOwnerRelease(ResourceOwner owner,
 	CurrentResourceOwner = save;
 }
 
+#ifdef XCP
+#define CNAME_MAXLEN 32
+#endif
+
 static void
 ResourceOwnerReleaseInternal(ResourceOwner owner,
 							 ResourceReleasePhase phase,
@@ -332,6 +346,17 @@ ResourceOwnerReleaseInternal(ResourceOwner owner,
 			FileClose(owner->files[owner->nfiles - 1]);
 		}
 
+#ifdef XCP
+		/* Ditto for prepared statements */
+		while (owner->nstmts > 0)
+		{
+			char *stmt = owner->stmts + ((owner->nstmts - 1) * CNAME_MAXLEN);
+			if (isCommit)
+				PrintPreparedStmtLeakWarning(stmt);
+			DropPreparedStatement(stmt);
+		}
+#endif
+
 		/* Clean up index scans too */
 		ReleaseResources_hash();
 	}
@@ -1138,3 +1163,90 @@ PrintFileLeakWarning(File file)
 		 "temporary file leak: File %d still referenced",
 		 file);
 }
+
+#ifdef XCP
+/*
+ * Make sure there is room for at least one more entry in a ResourceOwner's
+ * prepared statements reference array.
+ *
+ * This is separate from actually inserting an entry because if we run out
+ * of memory, it's critical to do so *before* acquiring the resource.
+ */
+void
+ResourceOwnerEnlargePreparedStmts(ResourceOwner owner)
+{
+	int			newmax;
+
+	if (owner->nstmts < owner->maxstmts)
+		return;					/* nothing to do */
+
+	if (owner->stmts == NULL)
+	{
+		newmax = 16;
+		owner->stmts = (char *)
+			MemoryContextAlloc(TopMemoryContext, newmax * CNAME_MAXLEN);
+		owner->maxstmts = newmax;
+	}
+	else
+	{
+		newmax = owner->maxstmts * 2;
+		owner->stmts = (char *)
+			repalloc(owner->stmts, newmax * CNAME_MAXLEN);
+		owner->maxstmts = newmax;
+	}
+}
+
+/*
+ * Remember that a prepared statement is owned by a ResourceOwner
+ *
+ * Caller must have previously done ResourceOwnerEnlargePreparedStmts()
+ */
+void
+ResourceOwnerRememberPreparedStmt(ResourceOwner owner, char *stmt)
+{
+	Assert(owner->nstmts < owner->maxstmts);
+	strncpy(owner->stmts + (owner->nstmts * CNAME_MAXLEN), stmt, CNAME_MAXLEN);
+	owner->nstmts++;
+}
+
+/*
+ * Forget that a temporary file is owned by a ResourceOwner
+ */
+void
+ResourceOwnerForgetPreparedStmt(ResourceOwner owner, char *stmt)
+{
+	char	   *stmts = owner->stmts;
+	int			ns1 = owner->nstmts - 1;
+	int			i;
+
+	for (i = ns1; i >= 0; i--)
+	{
+		if (strncmp(stmts + (i * CNAME_MAXLEN), stmt, CNAME_MAXLEN) == 0)
+		{
+			while (i < ns1)
+			{
+				strncpy(stmts + (i * CNAME_MAXLEN),
+						stmts + ((i + 1) * CNAME_MAXLEN),
+						CNAME_MAXLEN);
+				i++;
+			}
+			owner->nstmts = ns1;
+			return;
+		}
+	}
+	elog(ERROR, "prepared statement %s is not owned by resource owner %s",
+		 stmt, owner->name);
+}
+
+
+/*
+ * Debugging subroutine
+ */
+static void
+PrintPreparedStmtLeakWarning(char *stmt)
+{
+	elog(WARNING,
+		 "prepared statement leak: Statement %s still referenced",
+		 stmt);
+}
+#endif
diff --git a/src/include/commands/prepare.h b/src/include/commands/prepare.h
index 12fa60fd4f..222734b7aa 100644
--- a/src/include/commands/prepare.h
+++ b/src/include/commands/prepare.h
@@ -29,6 +29,9 @@ typedef struct
 	char		stmt_name[NAMEDATALEN];
 	CachedPlanSource *plansource;		/* the actual cached plan */
 	bool		from_sql;		/* prepared via SQL, not FE/BE protocol? */
+#ifdef XCP	
+	bool		use_resowner;	/* does it use resowner for tracking? */
+#endif	
 	TimestampTz prepare_time;	/* the time when the stmt was prepared */
 } PreparedStatement;
 
@@ -55,7 +58,8 @@ extern void ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into,
 /* Low-level access to stored prepared statements */
 extern void StorePreparedStatement(const char *stmt_name,
 					   CachedPlanSource *plansource,
-					   bool from_sql);
+					   bool from_sql,
+					   bool use_resowner);
 extern PreparedStatement *FetchPreparedStatement(const char *stmt_name,
 					   bool throwError);
 extern void DropPreparedStatement(const char *stmt_name, bool showError);
diff --git a/src/include/utils/resowner.h b/src/include/utils/resowner.h
index 11034e454e..4ac4a2bbaa 100644
--- a/src/include/utils/resowner.h
+++ b/src/include/utils/resowner.h
@@ -136,4 +136,13 @@ extern void ResourceOwnerRememberFile(ResourceOwner owner,
 extern void ResourceOwnerForgetFile(ResourceOwner owner,
 						File file);
 
+#ifdef XCP
+/* support for prepared statement management */
+extern void ResourceOwnerEnlargePreparedStmts(ResourceOwner owner);
+extern void ResourceOwnerRememberPreparedStmt(ResourceOwner owner,
+						  char *stmt);
+extern void ResourceOwnerForgetPreparedStmt(ResourceOwner owner,
+						char *stmt);
+#endif
+
 #endif   /* RESOWNER_H */