From f8d46a43b5cedfa1516f016155fed59e807748bb Mon Sep 17 00:00:00 2001 From: Ayush Tiwari Date: Mon, 28 Sep 2026 11:58:12 +0530 Subject: [PATCH v3] Allow unknown-size shmem attach in single-user mode ShmemRequestInternal() rejects SHMEM_ATTACH_UNKNOWN_SIZE in a standalone backend even when attaching after startup. Allow the sentinel when registering a request, and reject it instead when calculating the size of a new shared memory segment. This covers postmaster crash restart too, where ShmemIndex is still non-NULL while the request callbacks run again. Exercise late attachments in both normal and standalone backends, including missing areas and recovery after a failed attachment. Keep the GUC-driven failure and same-session retry test. Discussion: https://postgr.es/m/CAJTYsWXREGXCnWdZJsXt8pT2UY_Y2GSx+HOUu2X3s8YDCLj0_A@mail.gmail.com Co-authored-by: Heikki Linnakangas Backpatch-through: 19 --- src/backend/storage/ipc/shmem.c | 32 ++-- .../test_shmem/t/001_late_shmem_alloc.pl | 140 +++++++++++++----- .../modules/test_shmem/test_shmem--1.0.sql | 4 + src/test/modules/test_shmem/test_shmem.c | 58 ++++++++ 4 files changed, 173 insertions(+), 61 deletions(-) diff --git a/src/backend/storage/ipc/shmem.c b/src/backend/storage/ipc/shmem.c index 0e827344462..9fe3ea251bc 100644 --- a/src/backend/storage/ipc/shmem.c +++ b/src/backend/storage/ipc/shmem.c @@ -348,33 +348,22 @@ ShmemRequestInternal(ShmemStructOpts *options, ShmemRequestKind kind) MemoryContext oldcontext; ShmemRequest *request; + /* Check that we're in the right state */ + if (shmem_request_state != SRS_REQUESTING) + elog(ERROR, "ShmemRequestStruct can only be called from a shmem_request callback"); + /* Check the options */ if (options->name == NULL) elog(ERROR, "shared memory request is missing 'name' option"); - if (IsUnderPostmaster) - { - if (options->size <= 0 && options->size != SHMEM_ATTACH_UNKNOWN_SIZE) - elog(ERROR, "invalid size %zd for shared memory request for \"%s\"", - options->size, options->name); - } - else - { - if (options->size == SHMEM_ATTACH_UNKNOWN_SIZE) - elog(ERROR, "SHMEM_ATTACH_UNKNOWN_SIZE cannot be used during startup"); - if (options->size <= 0) - elog(ERROR, "invalid size %zd for shared memory request for \"%s\"", - options->size, options->name); - } + if (options->size <= 0 && options->size != SHMEM_ATTACH_UNKNOWN_SIZE) + elog(ERROR, "invalid size %zd for shared memory request for \"%s\"", + options->size, options->name); if (options->alignment != 0 && pg_nextpower2_size_t(options->alignment) != options->alignment) elog(ERROR, "invalid alignment %zu for shared memory request for \"%s\"", options->alignment, options->name); - /* Check that we're in the right state */ - if (shmem_request_state != SRS_REQUESTING) - elog(ERROR, "ShmemRequestStruct can only be called from a shmem_request callback"); - /* Check that it's not already registered in this process */ foreach_ptr(ShmemRequest, existing, pending_shmem_requests) { @@ -398,8 +387,8 @@ ShmemRequestInternal(ShmemStructOpts *options, ShmemRequestKind kind) * ShmemGetRequestedSize() --- estimate the total size of all registered shared * memory structures. * - * This is called at postmaster startup, before the shared memory segment has - * been created. + * This is called when sizing a new segment at postmaster or standalone + * startup, including a postmaster crash restart. */ size_t ShmemGetRequestedSize(void) @@ -416,6 +405,9 @@ ShmemGetRequestedSize(void) { size_t alignment = request->options->alignment; + if (request->options->size == SHMEM_ATTACH_UNKNOWN_SIZE) + elog(ERROR, "SHMEM_ATTACH_UNKNOWN_SIZE cannot be used during startup"); + /* pad the start address for alignment like ShmemAllocRaw() does */ if (alignment < PG_CACHE_LINE_SIZE) alignment = PG_CACHE_LINE_SIZE; diff --git a/src/test/modules/test_shmem/t/001_late_shmem_alloc.pl b/src/test/modules/test_shmem/t/001_late_shmem_alloc.pl index 7d961fe2790..18e4b2f33bd 100644 --- a/src/test/modules/test_shmem/t/001_late_shmem_alloc.pl +++ b/src/test/modules/test_shmem/t/001_late_shmem_alloc.pl @@ -21,57 +21,56 @@ $node->stop; ### $node->start; -# Check that the attach counter is incremented on a new connection +# This first call to the function after startup loads the library +# and initializes the shmem area. my $attach_count1 = $node->safe_psql("postgres", "SELECT get_test_shmem_attach_count();"); + +# Check that the attach counter is incremented on a new connection my $attach_count2 = $node->safe_psql("postgres", "SELECT get_test_shmem_attach_count();"); cmp_ok($attach_count2, '>', $attach_count1, "attach callback is called in each backend"); -$node->stop; +# Allocate another shmem area, after the library is loaded. +my $stderr; +my $res = $node->safe_psql("postgres", + "SELECT test_shmem_register('test_shmem after startup', 20, 1);"); +is($res, 0, 'allocate after startup'); -### -# Test that trying to allocate a new shmem area with size = -# SHMEM_ATTACH_UNKNOWN_SIZE (-1) fails. -### -$node->append_conf('postgresql.conf', "test_shmem.area_size = -1"); -$node->start; +# Test attaching to it again +$res = $node->safe_psql("postgres", + "SELECT test_shmem_register('test_shmem after startup', 20, 2);"); +is($res, 1, 'attach after startup'); -my (undef, undef, $stderr) = - $node->psql("postgres", "SELECT get_test_shmem_attach_count();"); +# If the size doesn't match when attaching, you get an error +(undef, undef, $stderr) = + $node->psql("postgres", "SELECT test_shmem_register('test_shmem after startup', 25, 3);"); like( $stderr, - qr/cannot attach to shared memory struct "test_shmem area" because it does not exist/, - "unknown size request for a nonexistent area fails"); + qr/ERROR: shared memory struct "test_shmem after startup" was created with different size: existing 20, requested 25/, + "attaching with different size fails"); -$node->stop; -$node->adjust_conf('postgresql.conf', 'test_shmem.area_size', undef); +# Test attaching with SHMEM_ATTACH_UNKNOWN_SIZE +$res = + $node->safe_psql("postgres", "SELECT test_shmem_register('test_shmem after startup', -1, 4);"); +is($res, 2, 'attach with SHMEM_ATTACH_UNKNOWN_SIZE'); -### -# Test allocating memory after startup in single-user mode -### -SKIP: -{ - # Skip the test on Windows, as single-user mode would fail on permission - # failure with privileged accounts. - skip 'single-user test is not supported by this platform', 1 - if $windows_os; - my $query = "SELECT get_test_shmem_attach_count();\n"; - my $result = run_log( - [ - 'postgres', '--single', '-F', - '-c' => 'exit_on_error=true', - '-D' => $node->data_dir, - 'postgres' - ], - '<' => \$query); +(undef, undef, $stderr) = $node->psql("postgres", + "SELECT test_shmem_register('test_shmem missing area', -1, 3);"); +like($stderr, + qr/cannot attach to shared memory struct "test_shmem missing area" because it does not exist/, + "unknown-size request for a nonexistent area fails"); - ok($result, "shmem area is initialized in single-user mode"); -} +(undef, undef, $stderr) = $node->psql("postgres", + "SELECT test_shmem_register('test_shmem too small', 1, 3);"); +like($stderr, qr/shared memory size must be at least \d+ bytes/, + "request smaller than an integer fails"); + +$node->stop; ### -# Test that loading via shared_preload_libraries also works +# Test that loading via shared_preload_libraries works ### $node->append_conf('postgresql.conf', "shared_preload_libraries = 'test_shmem'"); @@ -99,9 +98,8 @@ else ); } -# clean up $node->stop; -$node->adjust_conf('postgresql.conf', "shared_preload_libraries", undef); +$node->adjust_conf('postgresql.conf', 'shared_preload_libraries', undef); ### # Test a failure in initializing the shared memory area @@ -150,9 +148,9 @@ my $session = $node->background_psql('postgres', on_error_stop => 0); # make the request larger than the memory reserved for after-startup # requests. -$session->query(q[SET test_shmem.area_size = '128kB';]); +$session->query(q[SET test_shmem.area_size = '128kB';], verbose => 0); -$session->query("SELECT get_test_shmem_attach_count();"); +$session->query("SELECT get_test_shmem_attach_count();", verbose => 0); like( $session->{stderr}, qr/not enough shared memory/, @@ -162,9 +160,69 @@ like( # requested, it gets cleaned up on allocation failure. Verify that a # request for a smaller area succeeds in the same session. $session->{stderr} = ''; -$session->query("SET test_shmem.area_size = default;"); -$session->query_safe("SELECT get_test_shmem_attach_count();"); +$session->query("SET test_shmem.area_size = default;", verbose => 0); +$session->query_safe("SELECT get_test_shmem_attach_count();", verbose => 0); $session->quit; $node->stop; +### +# Test allocating memory after startup in single-user mode +### +SKIP: +{ + # Skip the test on Windows, as single-user mode would fail on permission + # failure with privileged accounts. + skip 'single-user test is not supported by this platform', 5 + if $windows_os; + + my @command = ( + 'postgres', '--single', '-F', + '-D' => $node->data_dir); + + my $queries = "SELECT get_test_shmem_attach_count();\n"; + my $result = run_log([@command, '-c' => 'exit_on_error=true', 'postgres'], + '<' => \$queries); + ok($result, "shmem area is initialized in single-user mode"); + + $queries = qq{ +-- allocate +SELECT test_shmem_register('test_shmem after startup', 25, 1); +-- attach +SELECT test_shmem_register('test_shmem after startup', 25, 2); +-- attach with SHMEM_ATTACH_UNKNOWN_SIZE +SELECT test_shmem_register('test_shmem after startup', -1, 3); +}; + $result = run_log([@command, '-c' => 'exit_on_error=true', 'postgres'], + '<' => \$queries); + ok($result, "allocate and attach in single-user mode"); + + $queries = qq{ +SELECT test_shmem_register('test_shmem missing in single user', -1, 3); +SELECT test_shmem_register('test_shmem after error', 25, 7); +SELECT test_shmem_register('test_shmem after error', -1, 8); +}; + my ($stdout, $stderr); + $result = run_log([@command, '-c' => 'exit_on_error=false', 'postgres'], + '<' => \$queries, '>' => \$stdout, '2>' => \$stderr); + like($stderr, + qr/cannot attach to shared memory struct "test_shmem missing in single user" because it does not exist/, + "unknown-size request for a nonexistent area fails in single-user mode"); + ok($result && $stdout =~ /test_shmem_register = "7"/, + "requests succeed after a failed attachment in the same process"); + + my $startup_stderr; + $result = run_log( + [@command, '-c' => 'exit_on_error=true', + '-c' => 'shared_preload_libraries=test_shmem', + '-c' => 'test_shmem.area_size=-1', 'postgres'], + '<' => \$queries, '2>' => \$startup_stderr); + ok(!$result && + $startup_stderr =~ /SHMEM_ATTACH_UNKNOWN_SIZE cannot be used during startup/, + "unknown-size requests are rejected during single-user startup"); +} + +# clean up +$node->stop; +$node->adjust_conf('postgresql.conf', "shared_preload_libraries", undef); + done_testing(); diff --git a/src/test/modules/test_shmem/test_shmem--1.0.sql b/src/test/modules/test_shmem/test_shmem--1.0.sql index 2d01fd9256c..1dec3ce0d68 100644 --- a/src/test/modules/test_shmem/test_shmem--1.0.sql +++ b/src/test/modules/test_shmem/test_shmem--1.0.sql @@ -7,3 +7,7 @@ CREATE FUNCTION get_test_shmem_attach_count() RETURNS pg_catalog.int4 STRICT AS 'MODULE_PATHNAME' LANGUAGE C; + +CREATE FUNCTION test_shmem_register(pg_catalog.text, pg_catalog.int8, pg_catalog.int4) +RETURNS pg_catalog.int4 STRICT +AS 'MODULE_PATHNAME' LANGUAGE C; diff --git a/src/test/modules/test_shmem/test_shmem.c b/src/test/modules/test_shmem/test_shmem.c index 2a7e13512bf..f7c48f94653 100644 --- a/src/test/modules/test_shmem/test_shmem.c +++ b/src/test/modules/test_shmem/test_shmem.c @@ -20,6 +20,7 @@ #include "fmgr.h" #include "miscadmin.h" #include "storage/shmem.h" +#include "utils/builtins.h" #include "utils/guc.h" #include "utils/injection_point.h" @@ -129,3 +130,60 @@ get_test_shmem_attach_count(PG_FUNCTION_ARGS) elog(ERROR, "shmem area not yet initialized"); PG_RETURN_INT32(TestShmem->attach_count); } + + +/* + * Callback for test_shmem_register(). test_shmem_register() provides the + * options, we just pass them through to ShmemRequestStructWithOpts. + */ +static void +test_shmem_after_startup_request(void *arg) +{ + ShmemStructOpts *opts = (ShmemStructOpts *) arg; + + elog(LOG, "test_shmem_after_startup_request callback called"); + + ShmemRequestStructWithOpts(opts); +} + +/* + * Allocate or attach to a shared memory structure, with the caller-supplied + * name and size. + * + * The given integer 'new_value' is stored at the beginning of the structure, + * and the old value there is returned. + */ +PG_FUNCTION_INFO_V1(test_shmem_register); +Datum +test_shmem_register(PG_FUNCTION_ARGS) +{ + char *name = text_to_cstring(PG_GETARG_TEXT_PP(0)); + int64 size = PG_GETARG_INT64(1); + int new_value = PG_GETARG_INT32(2); + int old_value; + int *attached = NULL; + + ShmemStructOpts opts = { + .name = name, + .size = size, + .ptr = (void **) &attached, + }; + + ShmemCallbacks callbacks = { + .flags = SHMEM_CALLBACKS_ALLOW_AFTER_STARTUP, + .request_fn = test_shmem_after_startup_request, + .opaque_arg = &opts, + }; + + if (size > 0 && size < (int64) sizeof(int)) + elog(ERROR, "shared memory size must be at least %zu bytes", sizeof(int)); + + RegisterShmemCallbacks(&callbacks); + if (attached == NULL) + elog(ERROR, "could not attach to shared memory"); + + old_value = *attached; + *attached = new_value; + + PG_RETURN_INT32(old_value); +} -- 2.34.1