dshash: Make it possible to suppress out of memory errors

Introduce dshash_find_or_insert_extended, which is just like
dshash_find_or_insert except that it takes a flags argument.
Currently, the only supported flag is DSHASH_INSERT_NO_OOM, but
I have chosen to use an integer rather than a boolean in case we
end up with more flags in the future.

Reviewed-by: Chao Li <li.evan.chao@gmail.com>
Reviewed-by: Sami Imseih <samimseih@gmail.com>
Discussion: http://postgr.es/m/CA+TgmoaJwUukUZGu7_yL74oMTQQz2=zqucMhF9+9xBmSC5us1w@mail.gmail.com

src/backend/lib/dshash.c

@@ -167,7 +167,8 @@ struct dshash_table
 
 static void delete_item(dshash_table *hash_table,
 						dshash_table_item *item);
-static void resize(dshash_table *hash_table, size_t new_size_log2);
+static bool resize(dshash_table *hash_table, size_t new_size_log2,
+				   int flags);
 static inline void ensure_valid_bucket_pointers(dshash_table *hash_table);
 static inline dshash_table_item *find_in_bucket(dshash_table *hash_table,
 												const void *key,
@@ -178,7 +179,8 @@ static void insert_item_into_bucket(dshash_table *hash_table,
 									dsa_pointer *bucket);
 static dshash_table_item *insert_into_bucket(dshash_table *hash_table,
 											 const void *key,
-											 dsa_pointer *bucket);
+											 dsa_pointer *bucket,
+											 int flags);
 static bool delete_key_from_bucket(dshash_table *hash_table,
 								   const void *key,
 								   dsa_pointer *bucket_head);
@@ -422,19 +424,25 @@ dshash_find(dshash_table *hash_table, const void *key, bool exclusive)
 }
 
 /*
- * Returns a pointer to an exclusively locked item which must be released with
- * dshash_release_lock.  If the key is found in the hash table, 'found' is set
- * to true and a pointer to the existing entry is returned.  If the key is not
- * found, 'found' is set to false, and a pointer to a newly created entry is
- * returned.
+ * Find an existing entry in a dshash_table, or insert a new one.
+ *
+ * DSHASH_INSERT_NO_OOM causes this function to return NULL when no memory is
+ * available for the new entry. Otherwise, such allocations will result in
+ * an ERROR.
+ *
+ * Any entry returned by this function is exclusively locked, and the caller
+ * must release that lock using dshash_release_lock. Notes above dshash_find()
+ * regarding locking and error handling equally apply here.
+ *
+ * On return, *found is set to true if an existing entry was found in the
+ * hash table, and otherwise false.
  *
- * Notes above dshash_find() regarding locking and error handling equally
- * apply here.
  */
 void *
-dshash_find_or_insert(dshash_table *hash_table,
-					  const void *key,
-					  bool *found)
+dshash_find_or_insert_extended(dshash_table *hash_table,
+							   const void *key,
+							   bool *found,
+							   int flags)
 {
 	dshash_hash hash;
 	size_t		partition_index;
@@ -477,14 +485,25 @@ restart:
 			 * reacquire all the locks in the right order to avoid deadlocks.
 			 */
 			LWLockRelease(PARTITION_LOCK(hash_table, partition_index));
-			resize(hash_table, hash_table->size_log2 + 1);
+			if (!resize(hash_table, hash_table->size_log2 + 1, flags))
+			{
+				Assert((flags & DSHASH_INSERT_NO_OOM) != 0);
+				return NULL;
+			}
 
 			goto restart;
 		}
 
 		/* Finally we can try to insert the new item. */
 		item = insert_into_bucket(hash_table, key,
-								  &BUCKET_FOR_HASH(hash_table, hash));
+								  &BUCKET_FOR_HASH(hash_table, hash),
+								  flags);
+		if (item == NULL)
+		{
+			Assert((flags & DSHASH_INSERT_NO_OOM) != 0);
+			LWLockRelease(PARTITION_LOCK(hash_table, partition_index));
+			return NULL;
+		}
 		item->hash = hash;
 		/* Adjust per-lock-partition counter for load factor knowledge. */
 		++partition->count;
@@ -854,10 +873,14 @@ delete_item(dshash_table *hash_table, dshash_table_item *item)
  * Grow the hash table if necessary to the requested number of buckets.  The
  * requested size must be double some previously observed size.
  *
+ * If an out-of-memory condition is observed, this function returns false if
+ * flags includes DSHASH_INSERT_NO_OOM, and otherwise throws an ERROR. In all
+ * other cases, it returns true.
+ *
  * Must be called without any partition lock held.
  */
-static void
-resize(dshash_table *hash_table, size_t new_size_log2)
+static bool
+resize(dshash_table *hash_table, size_t new_size_log2, int flags)
 {
 	dsa_pointer old_buckets;
 	dsa_pointer new_buckets_shared;
@@ -865,6 +888,7 @@ resize(dshash_table *hash_table, size_t new_size_log2)
 	size_t		size;
 	size_t		new_size = ((size_t) 1) << new_size_log2;
 	size_t		i;
+	int			dsa_flags = DSA_ALLOC_HUGE | DSA_ALLOC_ZERO;
 
 	/*
 	 * Acquire the locks for all lock partitions.  This is expensive, but we
@@ -882,23 +906,34 @@ resize(dshash_table *hash_table, size_t new_size_log2)
 			 * obtaining all the locks and return early.
 			 */
 			LWLockRelease(PARTITION_LOCK(hash_table, 0));
-			return;
+			return true;
 		}
 	}
 
 	Assert(new_size_log2 == hash_table->control->size_log2 + 1);
 
 	/* Allocate the space for the new table. */
+	if (flags & DSHASH_INSERT_NO_OOM)
+		dsa_flags |= DSA_ALLOC_NO_OOM;
 	new_buckets_shared =
 		dsa_allocate_extended(hash_table->area,
 							  sizeof(dsa_pointer) * new_size,
-							  DSA_ALLOC_HUGE | DSA_ALLOC_ZERO);
-	new_buckets = dsa_get_address(hash_table->area, new_buckets_shared);
+							  dsa_flags);
+
+	/* If DSHASH_INSERT_NO_OOM was specified, allocation may have failed. */
+	if (!DsaPointerIsValid(new_buckets_shared))
+	{
+		/* Release all the locks and return without resizing. */
+		for (i = 0; i < DSHASH_NUM_PARTITIONS; ++i)
+			LWLockRelease(PARTITION_LOCK(hash_table, i));
+		return false;
+	}
 
 	/*
 	 * We've allocated the new bucket array; all that remains to do now is to
 	 * reinsert all items, which amounts to adjusting all the pointers.
 	 */
+	new_buckets = dsa_get_address(hash_table->area, new_buckets_shared);
 	size = ((size_t) 1) << hash_table->control->size_log2;
 	for (i = 0; i < size; ++i)
 	{
@@ -928,6 +963,8 @@ resize(dshash_table *hash_table, size_t new_size_log2)
 	/* Release all the locks. */
 	for (i = 0; i < DSHASH_NUM_PARTITIONS; ++i)
 		LWLockRelease(PARTITION_LOCK(hash_table, i));
+
+	return true;
 }
 
 /*
@@ -982,19 +1019,26 @@ insert_item_into_bucket(dshash_table *hash_table,
 
 /*
  * Allocate space for an entry with the given key and insert it into the
- * provided bucket.
+ * provided bucket.  Returns NULL if out of memory and DSHASH_INSERT_NO_OOM
+ * was specified in flags.
  */
 static dshash_table_item *
 insert_into_bucket(dshash_table *hash_table,
 				   const void *key,
-				   dsa_pointer *bucket)
+				   dsa_pointer *bucket,
+				   int flags)
 {
 	dsa_pointer item_pointer;
 	dshash_table_item *item;
+	int			dsa_flags;
 
-	item_pointer = dsa_allocate(hash_table->area,
-								hash_table->params.entry_size +
-								MAXALIGN(sizeof(dshash_table_item)));
+	dsa_flags = (flags & DSHASH_INSERT_NO_OOM) ? DSA_ALLOC_NO_OOM : 0;
+	item_pointer = dsa_allocate_extended(hash_table->area,
+										 hash_table->params.entry_size +
+										 MAXALIGN(sizeof(dshash_table_item)),
+										 dsa_flags);
+	if (!DsaPointerIsValid(item_pointer))
+		return NULL;
 	item = dsa_get_address(hash_table->area, item_pointer);
 	copy_key(hash_table, ENTRY_FROM_ITEM(item), key);
 	insert_item_into_bucket(hash_table, item_pointer, item, bucket);

src/include/lib/dshash.h

@@ -92,15 +92,23 @@ extern void dshash_detach(dshash_table *hash_table);
 extern dshash_table_handle dshash_get_hash_table_handle(dshash_table *hash_table);
 extern void dshash_destroy(dshash_table *hash_table);
 
+/* Flags for dshash_find_or_insert_extended. */
+#define DSHASH_INSERT_NO_OOM	0x01	/* no failure if out-of-memory */
+
 /* Finding, creating, deleting entries. */
 extern void *dshash_find(dshash_table *hash_table,
 						 const void *key, bool exclusive);
-extern void *dshash_find_or_insert(dshash_table *hash_table,
-								   const void *key, bool *found);
+extern void *dshash_find_or_insert_extended(dshash_table *hash_table,
+											const void *key, bool *found,
+											int flags);
 extern bool dshash_delete_key(dshash_table *hash_table, const void *key);
 extern void dshash_delete_entry(dshash_table *hash_table, void *entry);
 extern void dshash_release_lock(dshash_table *hash_table, void *entry);
 
+/* Find or insert with error on out-of-memory. */
+#define dshash_find_or_insert(hash_table, key, found) \
+	dshash_find_or_insert_extended(hash_table, key, found, 0)
+
 /* seq scan support */
 extern void dshash_seq_init(dshash_seq_status *status, dshash_table *hash_table,
 							bool exclusive);