simplehash: preserve consistency in case of OOM.

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[email protected]
Reviewed-by: Andres Freund
Reviewed-by: Gurjeet Singh

diff --git a/src/include/lib/simplehash.h b/src/include/lib/simplehash.h
index b7adc16b807d4ce66a28c515fd85655cb14f6051..b1a3d7f9277122201628d64c5ce21452186f1c78 100644 (file)
--- 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