Handling failure with subtransaction

src/backend/distributed/operations/shard_cleaner.c

@@ -169,10 +169,24 @@ TryDropOrphanedShards(bool waitForLocks)
 {
 	int droppedShardCount = 0;
 	MemoryContext savedContext = CurrentMemoryContext;
+
+	/*
+	 * Start a subtransaction so we can rollback database's state to it in case
+	 * of error.
+	 */
+	BeginInternalSubTransaction(NULL);
+
 	PG_TRY();
 	{
 		droppedShardCount = DropOrphanedShardsForMove(waitForLocks);
 		droppedShardCount += DropOrphanedShardsForCleanup();
+
+		/*
+		 * Releasing a subtransaction doesn't free its memory context, since the
+		 * data it contains will be needed at upper commit. See the comments for
+		 * AtSubCommit_Memory() at postgres/src/backend/access/transam/xact.c.
+		 */
+		ReleaseCurrentSubTransaction();
 	}
 	PG_CATCH();
 	{
@@ -180,6 +194,8 @@ TryDropOrphanedShards(bool waitForLocks)
 		ErrorData *edata = CopyErrorData();
 		FlushErrorState();
 
+		RollbackAndReleaseCurrentSubTransaction();
+
 		/* rethrow as WARNING */
 		edata->elevel = WARNING;
 		ThrowErrorData(edata);