From 80e2edf1a5d82735d9abfb85f62149d60005de03 Mon Sep 17 00:00:00 2001
From: Nils Dijk <nils@citusdata.com>
Date: Wed, 27 Jul 2022 14:54:45 +0200
Subject: [PATCH] refactor inner loop of background worker thing

---
 src/backend/distributed/utils/maintenanced.c | 79 ++++++++++----------
 1 file changed, 41 insertions(+), 38 deletions(-)

diff --git a/src/backend/distributed/utils/maintenanced.c b/src/backend/distributed/utils/maintenanced.c
index 40f59270f..23306c29c 100644
--- a/src/backend/distributed/utils/maintenanced.c
+++ b/src/backend/distributed/utils/maintenanced.c
@@ -912,57 +912,60 @@ RebalanceJobsBackgroundWorkerMain(Datum arg)
 			RebalanceJob *job = GetRunableRebalanceJob();
 			MemoryContextSwitchTo(oldContext);
 
-			if (job)
+			if (!job)
 			{
-				ereport(LOG, (errmsg("found job with jobid: %ld", job->jobid)));
-				MemoryContext savedContext = CurrentMemoryContext;
-
-				UpdateJobStatus(job, REBALANCE_JOB_STATUS_RUNNING);
 				PopActiveSnapshot();
 				CommitTransactionCommand();
 
+				hasJobs = false;
+				break;
+			}
+
+			ereport(LOG, (errmsg("found job with jobid: %ld", job->jobid)));
+
+			/* Update job status to indicate it is running */
+			UpdateJobStatus(job, REBALANCE_JOB_STATUS_RUNNING);
+
+			PopActiveSnapshot();
+			CommitTransactionCommand();
+
+			MemoryContext savedContext = CurrentMemoryContext;
+			PG_TRY();
+			{
+				StartTransactionCommand();
+				PushActiveSnapshot(GetTransactionSnapshot());
+				if (ExecuteRebalanceJob(job))
+				{
+					UpdateJobStatus(job, REBALANCE_JOB_STATUS_DONE);
+
+					PopActiveSnapshot();
+					CommitTransactionCommand();
+				}
+			}
+			PG_CATCH();
+			{
+				MemoryContextSwitchTo(savedContext);
+
+				ErrorData *edata = CopyErrorData();
+				FlushErrorState();
+
 				StartTransactionCommand();
 				PushActiveSnapshot(GetTransactionSnapshot());
 
-				BeginInternalSubTransaction(NULL);
+				UpdateJobError(job, edata);
 
-				PG_TRY();
-				{
-					if (ExecuteRebalanceJob(job))
-					{
-						UpdateJobStatus(job, REBALANCE_JOB_STATUS_DONE);
-					}
+				PopActiveSnapshot();
+				CommitTransactionCommand();
 
-					ReleaseCurrentSubTransaction();
-				}
-				PG_CATCH();
-				{
-					MemoryContextSwitchTo(savedContext);
+				FreeErrorData(edata);
+				edata = NULL;
 
-					ErrorData *edata = CopyErrorData();
-					FlushErrorState();
-
-					RollbackAndReleaseCurrentSubTransaction();
-
-					UpdateJobError(job, edata);
-
-					FreeErrorData(edata);
-					edata = NULL;
-
-					/* TODO log that there was an error */
-				}
-				PG_END_TRY();
-			}
-			else
-			{
-				hasJobs = false;
+				/* TODO log that there was an error */
 			}
+			PG_END_TRY();
 		}
-
-		PopActiveSnapshot();
-		CommitTransactionCommand();
-		ProcessCompletedNotifies();
 	}
+
 	MemoryContextSwitchTo(oldContextPerJob);
 	MemoryContextDelete(perJobContext);
 }