have_memoize_track_cachekey_memory.patch

diff --git a/src/backend/executor/nodeMemoize.c b/src/backend/executor/nodeMemoize.c
index 33bcff4783..e7d240542d 100644
--- a/src/backend/executor/nodeMemoize.c
+++ b/src/backend/executor/nodeMemoize.c
@@ -363,9 +363,10 @@ remove_cache_entry(MemoizeState *mstate, MemoizeEntry *entry)
 	/*
 	 * Update memory accounting. entry_purge_tuples should have already
 	 * subtracted the memory used for each cached tuple.  Here we just update
-	 * the amount used by the entry itself.
+	 * the amount used by the entry itself and the key.
 	 */
 	mstate->mem_used -= EMPTY_ENTRY_MEMORY_BYTES(entry);
+	mstate->mem_used -= sizeof(MemoizeKey) + GetMemoryChunkSpace(key->params);
 
 	/* Remove the entry from the cache */
 	memoize_delete_item(mstate->hashtable, entry);
@@ -535,6 +536,12 @@ cache_lookup(MemoizeState *mstate, bool *found)
 	entry->key = key = (MemoizeKey *) palloc(sizeof(MemoizeKey));
 	key->params = ExecCopySlotMinimalTuple(mstate->probeslot);
 
+	/*
+	 * Account for the memory used for storing the params tuple and the
+	 * MemoizeKey struct
+	 */
+	mstate->mem_used += sizeof(MemoizeKey) + GetMemoryChunkSpace(key->params);
+
 	/* Update the total cache memory utilization */
 	mstate->mem_used += EMPTY_ENTRY_MEMORY_BYTES(entry);
 
@@ -1057,6 +1064,7 @@ ExecEndMemoize(MemoizeState *node)
 			MemoizeTuple *tuple = entry->tuplehead;
 
 			mem += EMPTY_ENTRY_MEMORY_BYTES(entry);
+			mem += sizeof(MemoizeKey) + GetMemoryChunkSpace(entry->key->params);
 			while (tuple != NULL)
 			{
 				mem += CACHE_TUPLE_BYTES(tuple);
diff --git a/src/backend/optimizer/path/costsize.c b/src/backend/optimizer/path/costsize.c
index 7918bb6f0d..f87a900ac2 100644
--- a/src/backend/optimizer/path/costsize.c
+++ b/src/backend/optimizer/path/costsize.c
@@ -2481,6 +2481,7 @@ cost_memoize_rescan(PlannerInfo *root, MemoizePath *mpath,
 					Cost *rescan_startup_cost, Cost *rescan_total_cost)
 {
 	EstimationInfo estinfo;
+	ListCell   *lc;
 	Cost		input_startup_cost = mpath->subpath->startup_cost;
 	Cost		input_total_cost = mpath->subpath->total_cost;
 	double		tuples = mpath->subpath->rows;
@@ -2495,7 +2496,6 @@ cost_memoize_rescan(PlannerInfo *root, MemoizePath *mpath,
 	double		hit_ratio;
 	Cost		startup_cost;
 	Cost		total_cost;
-
 	/* available cache space */
 	hash_mem_bytes = get_hash_memory_limit();
 
@@ -2504,12 +2504,14 @@ cost_memoize_rescan(PlannerInfo *root, MemoizePath *mpath,
 	 * To provide us with better estimations on how many cache entries we can
 	 * store at once, we make a call to the executor here to ask it what
 	 * memory overheads there are for a single cache entry.
-	 *
-	 * XXX we also store the cache key, but that's not accounted for here.
 	 */
 	est_entry_bytes = relation_byte_size(tuples, width) +
 		ExecEstimateCacheEntryOverheadBytes(tuples);
 
+	/* include the estimated width for the cache keys */
+	foreach(lc, mpath->param_exprs)
+		est_entry_bytes += get_expr_width(root, (Expr *) lfirst(lc));
+
 	/* estimate on the upper limit of cache entries we can hold at once */
 	est_cache_entries = floor(hash_mem_bytes / est_entry_bytes);
 
@@ -5995,6 +5997,48 @@ set_rel_width(PlannerInfo *root, RelOptInfo *rel)
 	rel->reltarget->width = tuple_width;
 }
 
+/*
+ * get_var_width
+ *		Estimate width of 'var' by attempting to lookup the width cached in
+ *		the owning RelOptInfo's attr_widths array.  Fall back on getting the
+ *		average width for the given Var's type.
+ *
+ * Note 'var' must be from the current query level.
+ */
+static int32
+get_var_width(PlannerInfo *root, Var *var)
+{
+	int32		var_width;
+
+	/* We should not see any upper-level Vars here */
+	Assert(var->varlevelsup == 0);
+
+	/* Try to get data from RelOptInfo cache */
+	if (!IS_SPECIAL_VARNO(var->varno) &&
+		var->varno < root->simple_rel_array_size)
+	{
+		RelOptInfo *rel = root->simple_rel_array[var->varno];
+
+		if (rel != NULL &&
+			var->varattno >= rel->min_attr &&
+			var->varattno <= rel->max_attr)
+		{
+			int			ndx = var->varattno - rel->min_attr;
+
+			if (rel->attr_widths[ndx] > 0)
+				return rel->attr_widths[ndx];
+		}
+	}
+
+	/*
+	 * No cached data available, so estimate using just the type info.
+	 */
+	var_width = get_typavgwidth(var->vartype, var->vartypmod);
+	Assert(var_width > 0);
+
+	return var_width;
+}
+
 /*
  * set_pathtarget_cost_width
  *		Set the estimated eval cost and output width of a PathTarget tlist.
@@ -6022,40 +6066,7 @@ set_pathtarget_cost_width(PlannerInfo *root, PathTarget *target)
 		Node	   *node = (Node *) lfirst(lc);
 
 		if (IsA(node, Var))
-		{
-			Var		   *var = (Var *) node;
-			int32		item_width;
-
-			/* We should not see any upper-level Vars here */
-			Assert(var->varlevelsup == 0);
-
-			/* Try to get data from RelOptInfo cache */
-			if (!IS_SPECIAL_VARNO(var->varno) &&
-				var->varno < root->simple_rel_array_size)
-			{
-				RelOptInfo *rel = root->simple_rel_array[var->varno];
-
-				if (rel != NULL &&
-					var->varattno >= rel->min_attr &&
-					var->varattno <= rel->max_attr)
-				{
-					int			ndx = var->varattno - rel->min_attr;
-
-					if (rel->attr_widths[ndx] > 0)
-					{
-						tuple_width += rel->attr_widths[ndx];
-						continue;
-					}
-				}
-			}
-
-			/*
-			 * No cached data available, so estimate using just the type info.
-			 */
-			item_width = get_typavgwidth(var->vartype, var->vartypmod);
-			Assert(item_width > 0);
-			tuple_width += item_width;
-		}
+			tuple_width += get_var_width(root, (Var *) node);
 		else
 		{
 			/*
@@ -6081,6 +6092,25 @@ set_pathtarget_cost_width(PlannerInfo *root, PathTarget *target)
 	return target;
 }
 
+/*
+ * get_expr_width
+ *		Estimate the width of the given expr attempting to use the width
+ *		cached in a Var's owning RelOptInfo, else fallback on the type's
+ *		average width when unable to or when the given Expr is not a Var.
+ */
+int32
+get_expr_width(PlannerInfo *root, Expr *expr)
+{
+	int32		width;
+
+	if (IsA(expr, Var))
+		return get_var_width(root, (Var *) expr);
+
+	width = get_typavgwidth(exprType((Node *) expr), exprTypmod((Node *) expr));
+	Assert(width > 0);
+	return width;
+}
+
 /*
  * relation_byte_size
  *	  Estimate the storage space in bytes for a given number of tuples
diff --git a/src/include/optimizer/cost.h b/src/include/optimizer/cost.h
index 6cf49705d3..5b0d093dcc 100644
--- a/src/include/optimizer/cost.h
+++ b/src/include/optimizer/cost.h
@@ -209,6 +209,7 @@ extern void set_namedtuplestore_size_estimates(PlannerInfo *root, RelOptInfo *re
 extern void set_result_size_estimates(PlannerInfo *root, RelOptInfo *rel);
 extern void set_foreign_size_estimates(PlannerInfo *root, RelOptInfo *rel);
 extern PathTarget *set_pathtarget_cost_width(PlannerInfo *root, PathTarget *target);
+extern int32 get_expr_width(PlannerInfo *root, Expr *expr);
 extern double compute_bitmap_pages(PlannerInfo *root, RelOptInfo *baserel,
 								   Path *bitmapqual, int loop_count, Cost *cost, double *tuple);