simplehash: preserve consistency in case of OOM.

jeff-davis · jeff-davis · commit b282fa88df78 · 2023-11-17T13:58:16.000-08:00
Compute size first, then allocate, then update the structure. Previously, an out-of-memory when growing could leave the hashtable in an inconsistent state. Discussion: https://postgr.es/m/20231117201334.eyb542qr5yk4gilq@awork3.anarazel.de Reviewed-by: Andres Freund Reviewed-by: Gurjeet Singh
diff --git a/src/include/lib/simplehash.h b/src/include/lib/simplehash.h
@@ -128,7 +128,8 @@
 #define SH_STAT SH_MAKE_NAME(stat)
 
 /* internal helper functions (no externally visible prototypes) */
-#define SH_COMPUTE_PARAMETERS SH_MAKE_NAME(compute_parameters)
+#define SH_COMPUTE_SIZE SH_MAKE_NAME(compute_size)
+#define SH_UPDATE_PARAMETERS SH_MAKE_NAME(update_parameters)
 #define SH_NEXT SH_MAKE_NAME(next)
 #define SH_PREV SH_MAKE_NAME(prev)
 #define SH_DISTANCE_FROM_OPTIMAL SH_MAKE_NAME(distance)
@@ -303,11 +304,11 @@ SH_SCOPE void SH_STAT(SH_TYPE * tb);
 #endif
 
 /*
- * Compute sizing parameters for hashtable. Called when creating and growing
- * the hashtable.
+ * Compute allocation size for hashtable. Result can be passed to
+ * SH_UPDATE_PARAMETERS.
  */
-static inline void
-SH_COMPUTE_PARAMETERS(SH_TYPE * tb, uint64 newsize)
+static inline uint64
+SH_COMPUTE_SIZE(uint64 newsize)
 {
 	uint64		size;
 
@@ -325,6 +326,18 @@ SH_COMPUTE_PARAMETERS(SH_TYPE * tb, uint64 newsize)
 	if (unlikely((((uint64) sizeof(SH_ELEMENT_TYPE)) * size) >= SIZE_MAX / 2))
 		sh_error("hash table too large");
 
+	return size;
+}
+
+/*
+ * Update sizing parameters for hashtable. Called when creating and growing
+ * the hashtable.
+ */
+static inline void
+SH_UPDATE_PARAMETERS(SH_TYPE * tb, uint64 newsize)
+{
+	uint64		size = SH_COMPUTE_SIZE(newsize);
+
 	/* now set size */
 	tb->size = size;
 	tb->sizemask = (uint32) (size - 1);
@@ -446,10 +459,11 @@ SH_CREATE(MemoryContext ctx, uint32 nelements, void *private_data)
 	/* increase nelements by fillfactor, want to store nelements elements */
 	size = Min((double) SH_MAX_SIZE, ((double) nelements) / SH_FILLFACTOR);
 
-	SH_COMPUTE_PARAMETERS(tb, size);
+	size = SH_COMPUTE_SIZE(size);
 
-	tb->data = (SH_ELEMENT_TYPE *) SH_ALLOCATE(tb, sizeof(SH_ELEMENT_TYPE) * tb->size);
+	tb->data = (SH_ELEMENT_TYPE *) SH_ALLOCATE(tb, sizeof(SH_ELEMENT_TYPE) * size);
 
+	SH_UPDATE_PARAMETERS(tb, size);
 	return tb;
 }
 
@@ -490,10 +504,15 @@ SH_GROW(SH_TYPE * tb, uint64 newsize)
 	Assert(oldsize != SH_MAX_SIZE);
 	Assert(oldsize < newsize);
 
-	/* compute parameters for new table */
-	SH_COMPUTE_PARAMETERS(tb, newsize);
+	newsize = SH_COMPUTE_SIZE(newsize);
 
-	tb->data = (SH_ELEMENT_TYPE *) SH_ALLOCATE(tb, sizeof(SH_ELEMENT_TYPE) * tb->size);
+	tb->data = (SH_ELEMENT_TYPE *) SH_ALLOCATE(tb, sizeof(SH_ELEMENT_TYPE) * newsize);
+
+	/*
+	 * Update parameters for new table after allocation succeeds to avoid
+	 * inconsistent state on OOM.
+	 */
+	SH_UPDATE_PARAMETERS(tb, newsize);
 
 	newdata = tb->data;
 
@@ -1173,7 +1192,8 @@ SH_STAT(SH_TYPE * tb)
 #undef SH_STAT
 
 /* internal function names */
-#undef SH_COMPUTE_PARAMETERS
+#undef SH_COMPUTE_SIZE
+#undef SH_UPDATE_PARAMETERS
 #undef SH_COMPARE_KEYS
 #undef SH_INITIAL_BUCKET
 #undef SH_NEXT
