Add test functions to neon_test_utils to consume CPU, memory, or local disk.

This is handy for testing VM autoscaling.

pgxn/neon_test_utils/neon_test_utils--1.0.sql

@@ -7,6 +7,36 @@ AS 'MODULE_PATHNAME', 'test_consume_xids'
 LANGUAGE C STRICT
 PARALLEL UNSAFE;
 
+CREATE FUNCTION test_consume_cpu(seconds int)
+RETURNS VOID
+AS 'MODULE_PATHNAME', 'test_consume_cpu'
+LANGUAGE C STRICT
+PARALLEL UNSAFE;
+
+CREATE FUNCTION test_consume_memory(megabytes int)
+RETURNS VOID
+AS 'MODULE_PATHNAME', 'test_consume_memory'
+LANGUAGE C STRICT
+PARALLEL UNSAFE;
+
+CREATE FUNCTION test_release_memory(megabytes int DEFAULT NULL)
+RETURNS VOID
+AS 'MODULE_PATHNAME', 'test_release_memory'
+LANGUAGE C
+PARALLEL UNSAFE;
+
+CREATE FUNCTION test_consume_disk_space(megabytes int)
+RETURNS VOID
+AS 'MODULE_PATHNAME', 'test_consume_disk_space'
+LANGUAGE C STRICT
+PARALLEL UNSAFE;
+
+CREATE FUNCTION test_release_disk_space(megabytes int DEFAULT NULL)
+RETURNS VOID
+AS 'MODULE_PATHNAME', 'test_release_disk_space'
+LANGUAGE C
+PARALLEL UNSAFE;
+
 CREATE FUNCTION clear_buffer_cache()
 RETURNS VOID
 AS 'MODULE_PATHNAME', 'clear_buffer_cache'

pgxn/neon_test_utils/neontest.c

@@ -21,10 +21,12 @@
 #include "miscadmin.h"
 #include "storage/buf_internals.h"
 #include "storage/bufmgr.h"
+#include "storage/fd.h"
 #include "utils/builtins.h"
 #include "utils/pg_lsn.h"
 #include "utils/rel.h"
 #include "utils/varlena.h"
+#include "utils/wait_event.h"
 #include "../neon/pagestore_client.h"
 
 PG_MODULE_MAGIC;
@@ -32,6 +34,11 @@ PG_MODULE_MAGIC;
 extern void _PG_init(void);
 
 PG_FUNCTION_INFO_V1(test_consume_xids);
+PG_FUNCTION_INFO_V1(test_consume_cpu);
+PG_FUNCTION_INFO_V1(test_consume_memory);
+PG_FUNCTION_INFO_V1(test_release_memory);
+PG_FUNCTION_INFO_V1(test_consume_disk_space);
+PG_FUNCTION_INFO_V1(test_release_disk_space);
 PG_FUNCTION_INFO_V1(clear_buffer_cache);
 PG_FUNCTION_INFO_V1(get_raw_page_at_lsn);
 PG_FUNCTION_INFO_V1(get_raw_page_at_lsn_ex);
@@ -97,6 +104,208 @@ test_consume_xids(PG_FUNCTION_ARGS)
 	PG_RETURN_VOID();
 }
 
+
+/*
+ * test_consume_cpu(seconds int). Keeps one CPU busy for the given number of seconds.
+ */
+Datum
+test_consume_cpu(PG_FUNCTION_ARGS)
+{
+	int32		seconds = PG_GETARG_INT32(0);
+	TimestampTz start;
+	uint64		total_iterations = 0;
+
+	start = GetCurrentTimestamp();
+
+	for (;;)
+	{
+		TimestampTz elapsed;
+
+		elapsed = GetCurrentTimestamp() - start;
+		if (elapsed > (TimestampTz) seconds * USECS_PER_SEC)
+			break;
+
+		/* keep spinning */
+		for (int i = 0; i < 1000000; i++)
+			total_iterations++;
+		elog(DEBUG2, "test_consume_cpu(): %lld iterations in total", total_iterations);
+
+		CHECK_FOR_INTERRUPTS();
+	}
+
+	PG_RETURN_VOID();
+}
+
+static MemoryContext consume_cxt = NULL;
+static slist_head consumed_memory_chunks;
+static int64 num_memory_chunks;
+
+/*
+ * test_consume_memory(megabytes int).
+ *
+ * Consume given amount of memory. The allocation is made in TopMemoryContext,
+ * so it outlives the function, until you call test_release_memory to
+ * explicitly release it, or close the session.
+ */
+Datum
+test_consume_memory(PG_FUNCTION_ARGS)
+{
+	int32		megabytes = PG_GETARG_INT32(0);
+
+	/*
+	 * Consume the memory in a new memory context, so that it's convenient to
+	 * release and to display it separately in a possible memory context dump.
+	 */
+	if (consume_cxt == NULL)
+		consume_cxt = AllocSetContextCreate(TopMemoryContext,
+											"test_consume_memory",
+											ALLOCSET_DEFAULT_SIZES);
+
+	for (int32 i = 0; i < megabytes; i++)
+	{
+		char	   *p;
+
+		p = MemoryContextAllocZero(consume_cxt, 1024 * 1024);
+
+		/* touch the memory, so that it's really allocated by the kernel */
+		for (int j = 0; j < 1024 * 1024; j += 1024)
+			p[j] = j % 0xFF;
+
+		slist_push_head(&consumed_memory_chunks, (slist_node *) p);
+		num_memory_chunks++;
+	}
+
+	PG_RETURN_VOID();
+}
+
+/*
+ * test_release_memory(megabytes int). NULL releases all
+ */
+Datum
+test_release_memory(PG_FUNCTION_ARGS)
+{
+	TimestampTz start;
+
+	if (PG_ARGISNULL(0))
+	{
+		if (consume_cxt)
+		{
+			MemoryContextDelete(consume_cxt);
+			consume_cxt = NULL;
+			num_memory_chunks = 0;
+		}
+	}
+	else
+	{
+		int32		chunks_to_release = PG_GETARG_INT32(0);
+
+		if (chunks_to_release > num_memory_chunks)
+		{
+			elog(WARNING, "only %lld MB is consumed, releasing it all", num_memory_chunks);
+			chunks_to_release = num_memory_chunks;
+		}
+
+		for (int32 i = 0; i < chunks_to_release; i++)
+		{
+			slist_node *chunk = slist_pop_head_node(&consumed_memory_chunks);
+
+			pfree(chunk);
+			num_memory_chunks--;
+		}
+	}
+
+	PG_RETURN_VOID();
+}
+
+static File consume_file = -1;
+
+/*
+ * test_consume_disk_space(megabytes int).
+ *
+ * Like test_consume_memory, but cretes temporary files adding up to given
+ * size, instead of allocating memory.
+ */
+Datum
+test_consume_disk_space(PG_FUNCTION_ARGS)
+{
+	int32		megabytes = PG_GETARG_INT32(0);
+	off_t		current_size;
+	off_t		desired_size;
+	char	   *zerobuf;
+
+	if (consume_file == -1)
+	{
+		consume_file = OpenTemporaryFile(true);
+		current_size = 0;
+	}
+	else
+		current_size = FileSize(consume_file);
+
+	desired_size = current_size + (off_t) megabytes * 1024 * 1024;
+	zerobuf = palloc0(1024 * 1024);
+
+	while (current_size < desired_size)
+	{
+		int			bytes_written;
+
+		bytes_written = FileWrite(consume_file, zerobuf,
+								  Min(desired_size - current_size, 1024 * 1024),
+								  current_size, PG_WAIT_EXTENSION);
+		if (bytes_written < 0)
+			elog(ERROR, "could not extend temporary file: %m");
+		current_size += bytes_written;
+
+		CHECK_FOR_INTERRUPTS();
+	}
+	pfree(zerobuf);
+
+	PG_RETURN_VOID();
+}
+
+/*
+ * test_release_disk_space(megabytes int). NULL releases all
+ */
+Datum
+test_release_disk_space(PG_FUNCTION_ARGS)
+{
+	TimestampTz start;
+
+	if (PG_ARGISNULL(0))
+	{
+		if (consume_file != -1)
+		{
+			FileClose(consume_file);
+			consume_file = -1;
+		}
+	}
+	else
+	{
+		int32		megabytes_to_release = PG_GETARG_INT32(0);
+		off_t		current_size;
+		off_t		desired_size;
+
+		if (consume_file == -1)
+		{
+			elog(WARNING, "no extra disk space is consumed at the moment");
+			PG_RETURN_VOID();
+		}
+		current_size = FileSize(consume_file);
+
+		if ((int64) megabytes_to_release * 1024 * 1024 > current_size)
+		{
+			elog(WARNING, "only %lld MB of disk space is currently consumed, releasing it all", (long long) current_size / (1024 * 1024));
+			desired_size = 0;
+		}
+		else
+		{
+			desired_size = current_size - (int64) megabytes_to_release * 1024 * 1024;
+		}
+		FileTruncate(consume_file, desired_size, PG_WAIT_EXTENSION);
+	}
+
+	PG_RETURN_VOID();
+}
+
 /*
  * Flush the buffer cache, evicting all pages that are not currently pinned.
  */

test_runner/sql_regress/expected/neon-test-utils.out

@@ -0,0 +1,46 @@
+-- Test the test utils in pgxn/neon_test_utils. We don't test that
+-- these actually consume resources like they should - that would be
+-- tricky - but at least we check that they don't crash.
+CREATE EXTENSION neon_test_utils;
+select test_consume_cpu(1);
+ test_consume_cpu 
+------------------
+ 
+(1 row)
+
+select test_consume_memory(20); -- Allocate 20 MB
+ test_consume_memory 
+---------------------
+ 
+(1 row)
+
+select test_release_memory(5);  -- Release 5 MB
+ test_release_memory 
+---------------------
+ 
+(1 row)
+
+select test_release_memory();   -- Release the remaining 15 MB
+ test_release_memory 
+---------------------
+ 
+(1 row)
+
+select test_consume_disk_space(20); -- Allocate 20 MB of disk space
+ test_consume_disk_space 
+-------------------------
+ 
+(1 row)
+
+select test_release_disk_space(5);  -- Release 5 MB
+ test_release_disk_space 
+-------------------------
+ 
+(1 row)
+
+select test_release_disk_space();   -- Release the remaining 15 MB
+ test_release_disk_space 
+-------------------------
+ 
+(1 row)
+

test_runner/sql_regress/parallel_schedule

@@ -7,4 +7,5 @@
 test: neon-cid
 test: neon-rel-truncate
 test: neon-clog
+test: neon-test-utils
 test: neon-vacuum-full

test_runner/sql_regress/sql/neon-test-utils.sql

@@ -0,0 +1,15 @@
+-- Test the test utils in pgxn/neon_test_utils. We don't test that
+-- these actually consume resources like they should - that would be
+-- tricky - but at least we check that they don't crash.
+
+CREATE EXTENSION neon_test_utils;
+
+select test_consume_cpu(1);
+
+select test_consume_memory(20); -- Allocate 20 MB
+select test_release_memory(5);  -- Release 5 MB
+select test_release_memory();   -- Release the remaining 15 MB
+
+select test_consume_disk_space(20); -- Allocate 20 MB of disk space
+select test_release_disk_space(5);  -- Release 5 MB
+select test_release_disk_space();   -- Release the remaining 15 MB