From c18c6625d97e06dc056ab6f24a347a311400efd6 Mon Sep 17 00:00:00 2001
From: Hadi Moshayedi <hadi@moshayedi.net>
Date: Tue, 31 Oct 2017 21:20:18 -0400
Subject: [PATCH] Lock relations before calling citus_table_size().

This is to make sure they don't get dropped.
---
 .../distributed/utils/statistics_collection.c        | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/src/backend/distributed/utils/statistics_collection.c b/src/backend/distributed/utils/statistics_collection.c
index 91a05252a..55881c584 100644
--- a/src/backend/distributed/utils/statistics_collection.c
+++ b/src/backend/distributed/utils/statistics_collection.c
@@ -154,6 +154,17 @@ DistributedTablesSize(List *distTableOids)
 		Oid relationId = lfirst_oid(distTableOidCell);
 		Datum tableSizeDatum = 0;
 
+		/*
+		 * Relations can get dropped after getting the Oid list and before we
+		 * reach here. Acquire a lock to make sure the relation is available
+		 * while we are getting its size.
+		 */
+		Relation relation = try_relation_open(relationId, AccessShareLock);
+		if (relation == NULL)
+		{
+			continue;
+		}
+
 		/*
 		 * Ignore hash partitioned tables with size greater than 1, since
 		 * citus_table_size() doesn't work on them.
@@ -167,6 +178,7 @@ DistributedTablesSize(List *distTableOids)
 		tableSizeDatum = DirectFunctionCall1(citus_table_size,
 											 ObjectIdGetDatum(relationId));
 		totalSize += DatumGetInt64(tableSizeDatum);
+		heap_close(relation, AccessShareLock);
 	}
 
 	return totalSize;