From 66a8bbe5dcce8f73a621ca3d2456783908ee8787 Mon Sep 17 00:00:00 2001 From: Heikki Linnakangas Date: Fri, 17 Nov 2023 00:59:42 +0100 Subject: [PATCH] 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 | 30 +++ pgxn/neon_test_utils/neontest.c | 209 ++++++++++++++++++ .../sql_regress/expected/neon-test-utils.out | 46 ++++ test_runner/sql_regress/parallel_schedule | 1 + .../sql_regress/sql/neon-test-utils.sql | 15 ++ 5 files changed, 301 insertions(+) create mode 100644 test_runner/sql_regress/expected/neon-test-utils.out create mode 100644 test_runner/sql_regress/sql/neon-test-utils.sql diff --git a/pgxn/neon_test_utils/neon_test_utils--1.0.sql b/pgxn/neon_test_utils/neon_test_utils--1.0.sql index 402981a9a6..b09a22062c 100644 --- a/pgxn/neon_test_utils/neon_test_utils--1.0.sql +++ b/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' diff --git a/pgxn/neon_test_utils/neontest.c b/pgxn/neon_test_utils/neontest.c index aa644efd40..9c37a78936 100644 --- a/pgxn/neon_test_utils/neontest.c +++ b/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. */ diff --git a/test_runner/sql_regress/expected/neon-test-utils.out b/test_runner/sql_regress/expected/neon-test-utils.out new file mode 100644 index 0000000000..8962d5d43c --- /dev/null +++ b/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) + diff --git a/test_runner/sql_regress/parallel_schedule b/test_runner/sql_regress/parallel_schedule index 569c7b5066..d9508d1c90 100644 --- a/test_runner/sql_regress/parallel_schedule +++ b/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 diff --git a/test_runner/sql_regress/sql/neon-test-utils.sql b/test_runner/sql_regress/sql/neon-test-utils.sql new file mode 100644 index 0000000000..ed8e656eb0 --- /dev/null +++ b/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