Refactor FindShardInterval to use cacheEntry

All callers fetch a cache entry and extract/compute arguments for the
eventual FindShardInterval call, so it makes more sense to refactor
into that function itself; this solves the use-after-free bug, too.

src/backend/distributed/commands/multi_copy.c

@@ -1753,12 +1753,6 @@ CitusCopyDestReceiverStartup(DestReceiver *dest, int operation,
 	/* keep the table metadata to avoid looking it up for every tuple */
 	copyDest->tableMetadata = cacheEntry;
 
-	/* determine whether to use binary search */
-	if (partitionMethod != DISTRIBUTE_BY_HASH || !cacheEntry->hasUniformHashDistribution)
-	{
-		copyDest->useBinarySearch = true;
-	}
-
 	if (cacheEntry->replicationModel == REPLICATION_MODEL_2PC)
 	{
 		CoordinatedTransactionUse2PC();
@@ -1835,19 +1829,10 @@ CitusCopyDestReceiverReceive(TupleTableSlot *slot, DestReceiver *dest)
 {
 	CitusCopyDestReceiver *copyDest = (CitusCopyDestReceiver *) dest;
 
-	DistTableCacheEntry *tableMetadata = copyDest->tableMetadata;
-	char partitionMethod = tableMetadata->partitionMethod;
 	int partitionColumnIndex = copyDest->partitionColumnIndex;
 	TupleDesc tupleDescriptor = copyDest->tupleDescriptor;
 	CopyStmt *copyStatement = copyDest->copyStatement;
 
-	int shardCount = tableMetadata->shardIntervalArrayLength;
-	ShardInterval **shardIntervalCache = tableMetadata->sortedShardIntervalArray;
-
-	bool useBinarySearch = copyDest->useBinarySearch;
-	FmgrInfo *hashFunction = tableMetadata->hashFunction;
-	FmgrInfo *compareFunction = tableMetadata->shardIntervalCompareFunction;
-
 	HTAB *shardConnectionHash = copyDest->shardConnectionHash;
 	CopyOutState copyOutState = copyDest->copyOutState;
 	FmgrInfo *columnOutputFunctions = copyDest->columnOutputFunctions;
@@ -1907,10 +1892,7 @@ CitusCopyDestReceiverReceive(TupleTableSlot *slot, DestReceiver *dest)
 	 * For reference table, this function blindly returns the tables single
 	 * shard.
 	 */
-	shardInterval = FindShardInterval(partitionColumnValue, shardIntervalCache,
+	shardInterval = FindShardInterval(partitionColumnValue, copyDest->tableMetadata);
-									  shardCount, partitionMethod,
-									  compareFunction, hashFunction,
-									  useBinarySearch);
 	if (shardInterval == NULL)
 	{
 		ereport(ERROR, (errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),

src/backend/distributed/planner/multi_router_planner.c

@@ -2057,35 +2057,13 @@ ShardInterval *
 FastShardPruning(Oid distributedTableId, Datum partitionValue)
 {
 	DistTableCacheEntry *cacheEntry = DistributedTableCacheEntry(distributedTableId);
-	int shardCount = cacheEntry->shardIntervalArrayLength;
-	ShardInterval **sortedShardIntervalArray = cacheEntry->sortedShardIntervalArray;
-	bool useBinarySearch = false;
-	char partitionMethod = cacheEntry->partitionMethod;
-	FmgrInfo *shardIntervalCompareFunction = cacheEntry->shardIntervalCompareFunction;
-	bool hasUniformHashDistribution = cacheEntry->hasUniformHashDistribution;
-	FmgrInfo *hashFunction = NULL;
 	ShardInterval *shardInterval = NULL;
 
-	/* determine whether to use binary search */
-	if (partitionMethod != DISTRIBUTE_BY_HASH || !hasUniformHashDistribution)
-	{
-		useBinarySearch = true;
-	}
-
-	/* we only need hash functions for hash distributed tables */
-	if (partitionMethod == DISTRIBUTE_BY_HASH)
-	{
-		hashFunction = cacheEntry->hashFunction;
-	}
-
 	/*
 	 * Call FindShardInterval to find the corresponding shard interval for the
 	 * given partition value.
 	 */
-	shardInterval = FindShardInterval(partitionValue, sortedShardIntervalArray,
+	shardInterval = FindShardInterval(partitionValue, cacheEntry);
-									  shardCount, partitionMethod,
-									  shardIntervalCompareFunction, hashFunction,
-									  useBinarySearch);
 
 	return shardInterval;
 }

src/backend/distributed/utils/shardinterval_utils.c

@@ -23,9 +23,7 @@
 #include "utils/memutils.h"
 
 
-static int FindShardIntervalIndex(Datum searchedValue, ShardInterval **shardIntervalCache,
+static int FindShardIntervalIndex(Datum searchedValue, DistTableCacheEntry *cacheEntry);
-								  int shardCount, char partitionMethod,
-								  FmgrInfo *compareFunction, bool useBinarySearch);
 static int SearchCachedShardInterval(Datum partitionColumnValue,
 									 ShardInterval **shardIntervalCache,
 									 int shardCount, FmgrInfo *compareFunction);
@@ -188,12 +186,7 @@ ShardIndex(ShardInterval *shardInterval)
 	Datum shardMinValue = shardInterval->minValue;
 
 	DistTableCacheEntry *cacheEntry = DistributedTableCacheEntry(distributedTableId);
-	ShardInterval **shardIntervalCache = cacheEntry->sortedShardIntervalArray;
-	int shardCount = cacheEntry->shardIntervalArrayLength;
 	char partitionMethod = cacheEntry->partitionMethod;
-	FmgrInfo *compareFunction = cacheEntry->shardIntervalCompareFunction;
-	bool hasUniformHashDistribution = cacheEntry->hasUniformHashDistribution;
-	bool useBinarySearch = false;
 
 	/*
 	 * Note that, we can also support append and range distributed tables, but
@@ -215,15 +208,7 @@ ShardIndex(ShardInterval *shardInterval)
 		return shardIndex;
 	}
 
-	/* determine whether to use binary search */
+	shardIndex = FindShardIntervalIndex(shardMinValue, cacheEntry);
-	if (partitionMethod != DISTRIBUTE_BY_HASH || !hasUniformHashDistribution)
-	{
-		useBinarySearch = true;
-	}
-
-	shardIndex = FindShardIntervalIndex(shardMinValue, shardIntervalCache,
-										shardCount, partitionMethod,
-										compareFunction, useBinarySearch);
 
 	return shardIndex;
 }
@@ -236,28 +221,24 @@ ShardIndex(ShardInterval *shardInterval)
  * as NULL for them.
  */
 ShardInterval *
-FindShardInterval(Datum partitionColumnValue, ShardInterval **shardIntervalCache,
+FindShardInterval(Datum partitionColumnValue, DistTableCacheEntry *cacheEntry)
-				  int shardCount, char partitionMethod, FmgrInfo *compareFunction,
-				  FmgrInfo *hashFunction, bool useBinarySearch)
 {
 	Datum searchedValue = partitionColumnValue;
 	int shardIndex = INVALID_SHARD_INDEX;
 
-	if (partitionMethod == DISTRIBUTE_BY_HASH)
+	if (cacheEntry->partitionMethod == DISTRIBUTE_BY_HASH)
 	{
-		searchedValue = FunctionCall1(hashFunction, partitionColumnValue);
+		searchedValue = FunctionCall1(cacheEntry->hashFunction, partitionColumnValue);
 	}
 
-	shardIndex = FindShardIntervalIndex(searchedValue, shardIntervalCache,
+	shardIndex = FindShardIntervalIndex(searchedValue, cacheEntry);
-										shardCount, partitionMethod,
-										compareFunction, useBinarySearch);
 
 	if (shardIndex == INVALID_SHARD_INDEX)
 	{
 		return NULL;
 	}
 
-	return shardIntervalCache[shardIndex];
+	return cacheEntry->sortedShardIntervalArray[shardIndex];
 }
 
 
@@ -273,10 +254,14 @@ FindShardInterval(Datum partitionColumnValue, ShardInterval **shardIntervalCache
  * fire this.
  */
 static int
-FindShardIntervalIndex(Datum searchedValue, ShardInterval **shardIntervalCache,
+FindShardIntervalIndex(Datum searchedValue, DistTableCacheEntry *cacheEntry)
-					   int shardCount, char partitionMethod, FmgrInfo *compareFunction,
-					   bool useBinarySearch)
 {
+	ShardInterval **shardIntervalCache = cacheEntry->sortedShardIntervalArray;
+	int shardCount = cacheEntry->shardIntervalArrayLength;
+	char partitionMethod = cacheEntry->partitionMethod;
+	FmgrInfo *compareFunction = cacheEntry->shardIntervalCompareFunction;
+	bool useBinarySearch = (partitionMethod != DISTRIBUTE_BY_HASH ||
+							!cacheEntry->hasUniformHashDistribution);
 	int shardIndex = INVALID_SHARD_INDEX;
 
 	if (partitionMethod == DISTRIBUTE_BY_HASH)

src/include/distributed/multi_copy.h

@@ -61,7 +61,6 @@ typedef struct CitusCopyDestReceiver
 
 	/* distributed table metadata */
 	DistTableCacheEntry *tableMetadata;
-	bool useBinarySearch;
 
 	/* open relation handle */
 	Relation distributedRelation;

src/include/distributed/shardinterval_utils.h

@@ -13,6 +13,7 @@
 #define SHARDINTERVAL_UTILS_H_
 
 #include "distributed/master_metadata_utility.h"
+#include "distributed/metadata_cache.h"
 #include "nodes/primnodes.h"
 
 #define INVALID_SHARD_INDEX -1
@@ -33,10 +34,7 @@ extern int CompareRelationShards(const void *leftElement,
 								 const void *rightElement);
 extern int ShardIndex(ShardInterval *shardInterval);
 extern ShardInterval * FindShardInterval(Datum partitionColumnValue,
-										 ShardInterval **shardIntervalCache,
+										 DistTableCacheEntry *cacheEntry);
-										 int shardCount, char partitionMethod,
-										 FmgrInfo *compareFunction,
-										 FmgrInfo *hashFunction, bool useBinarySearch);
 extern bool SingleReplicatedTable(Oid relationId);
 
 #endif /* SHARDINTERVAL_UTILS_H_ */