diff --git a/src/backend/storage/ipc/shmem.c b/src/backend/storage/ipc/shmem.c
index 0e827344462..6c3519ac6aa 100644
--- a/src/backend/storage/ipc/shmem.c
+++ b/src/backend/storage/ipc/shmem.c
@@ -348,33 +348,27 @@ 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 == SHMEM_ATTACH_UNKNOWN_SIZE)
 	{
-		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)
+		if (ShmemIndex == NULL)
 			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);
 	}
+	else if (options->size <= 0)
+		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)
 	{
@@ -790,6 +784,8 @@ void
 ResetShmemAllocator(void)
 {
 	Assert(!IsUnderPostmaster);
+	ShmemAllocator = NULL;
+	ShmemIndex = NULL;
 	shmem_request_state = SRS_INITIAL;
 
 	pending_shmem_requests = NIL;
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..0e14e7b15d0 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
@@ -7,71 +7,18 @@ use PostgreSQL::Test::Cluster;
 use PostgreSQL::Test::Utils;
 use Test::More;
 
-# Initialize a cluster with the extension installed.  The tests will
-# call the function that comes with the extension to load it.
+# Initialize a cluster with the extension installed.  Some tests below will
+# load the library via shared_preload_libraries, while others will load it
+# after startup by calling the extension's functions.
 my $node = PostgreSQL::Test::Cluster->new('main');
 $node->init;
 $node->start;
 $node->safe_psql("postgres", "CREATE EXTENSION test_shmem");
 $node->stop;
 
-###
-# Test allocating memory after startup, i.e. when the library is not
-# in shared_preload_libraries
-###
-$node->start;
-
-# Check that the attach counter is incremented on a new connection
-my $attach_count1 =
-  $node->safe_psql("postgres", "SELECT get_test_shmem_attach_count();");
-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;
 
 ###
-# 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;
-
-my (undef, undef, $stderr) =
-  $node->psql("postgres", "SELECT get_test_shmem_attach_count();");
-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");
-
-$node->stop;
-$node->adjust_conf('postgresql.conf', 'test_shmem.area_size', undef);
-
-###
-# 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);
-
-	ok($result, "shmem area is initialized in single-user mode");
-}
-
-###
-# 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'");
@@ -81,9 +28,9 @@ $node->start;
 # called or not, depending on whether this is an EXEC_BACKEND build.
 my $exec_backend =
   $node->safe_psql("postgres", "SHOW debug_exec_backend;") eq 'on';
-$attach_count1 =
+my $attach_count1 =
   $node->safe_psql("postgres", "SELECT get_test_shmem_attach_count();");
-$attach_count2 =
+my $attach_count2 =
   $node->safe_psql("postgres", "SELECT get_test_shmem_attach_count();");
 
 if ($exec_backend)
@@ -99,16 +46,70 @@ else
 	);
 }
 
-# clean up
 $node->stop;
-$node->adjust_conf('postgresql.conf', "shared_preload_libraries", undef);
+$node->adjust_conf('postgresql.conf', 'shared_preload_libraries', undef);
+
+
+###
+# Test allocating memory after startup, i.e. when the library is not
+# in shared_preload_libraries
+###
+$node->start;
+
+# This first call to the function after startup loads the library
+# and initializes the shmem area.
+$attach_count1 =
+  $node->safe_psql("postgres", "SELECT get_test_shmem_attach_count();");
+
+# Check that the attach counter is incremented on a new connection
+$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");
+
+# Allocate another shmem area, after the library has been loaded
+my $res = $node->safe_psql("postgres",
+	"SELECT test_shmem_register('test_shmem after startup', 20, 1);");
+is($res, 0, 'allocate after startup');
+
+# 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');
+
+# If the size doesn't match when attaching, you get an error
+my (undef, undef, $stderr) =
+  $node->psql("postgres",
+	"SELECT test_shmem_register('test_shmem after startup', 25, 3);");
+like(
+	$stderr,
+	qr/ERROR:  shared memory struct "test_shmem after startup" was created with different size: existing 20, requested 25/,
+	"attaching with different size fails");
+
+# 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');
+
+# Cannot use SHMEM_ATTACH_UNKNOWN_SIZE with a segment that doesn't exist
+(undef, undef, $stderr) =
+  $node->psql("postgres",
+	"SELECT test_shmem_register('test_shmem missing area', -1, 5);");
+like(
+	$stderr,
+	qr/ERROR:  cannot attach to shared memory struct "test_shmem missing area" because it does not exist/,
+	"unknown-size request for a nonexistent area fails");
+
+$node->stop;
+
 
 ###
 # Test a failure in initializing the shared memory area
 ###
 SKIP:
 {
-	skip "injection points not supported by this build",
+	skip "injection points not supported by this build", 2
 	  if $ENV{enable_injection_points} ne 'yes';
 	$node->start;
 	$node->safe_psql("postgres", "CREATE EXTENSION injection_points;");
@@ -139,9 +140,11 @@ SKIP:
 	$node->stop;
 }
 
+
 ###
 # Test "out of shared memory" in an after-startup request
 ###
+
 # Huge pages round up the main shared memory segment, which can leave more
 # unused space than the request below.  Disable them to make the test reliable.
 $node->append_conf('postgresql.conf', 'huge_pages = off');
@@ -167,4 +170,59 @@ $session->query_safe("SELECT get_test_shmem_attach_count();");
 $session->quit;
 $node->stop;
 
+
+###
+# Test shmem allocations 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', 4
+	  if $windows_os;
+
+	my @command = (
+		'postgres', '--single', '-F',
+		'-c' => 'exit_on_error=true',
+		'-D' => $node->data_dir);
+
+	# The call to get_test_shmem_attach_count() loads the library and
+	# allocates the shmem area
+	my $queries = "SELECT get_test_shmem_attach_count();\n";
+	my $result = run_log([ @command, 'postgres' ], '<' => \$queries);
+	ok($result, "shmem area is initialized in single-user mode");
+
+	# Test allocate and attach requests in a single-user session
+	$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, 'postgres' ], '<' => \$queries);
+	ok($result, "allocate and attach in single-user mode");
+
+	# Cannot use SHMEM_ATTACH_UNKNOWN_SIZE with a segment that doesn't exist.
+	# We tested this in the after-startup tests already, but test it at startup,
+	# in single-user mode, too.
+	my $startup_stderr;
+	$result = run_log(
+		[
+			@command,
+			'-c' => 'shared_preload_libraries=test_shmem',
+			'-c' => 'test_shmem.area_size=-1',
+			'postgres'
+		],
+		'<' => \$queries,
+		'2>' => \$startup_stderr);
+	ok(!$result,
+		"unknown-size request during single-user startup is rejected");
+	like(
+		$startup_stderr,
+		qr/SHMEM_ATTACH_UNKNOWN_SIZE cannot be used during startup/,
+		"error message on rejected unknown-size request");
+}
+
 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..bb822ea30a8 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,57 @@ 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 ShmemRequestStruct.
+ */
+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 shmem structure, with the caller-supplied name and
+ * size.
+ *
+ * The given integer 'new_value' is stored in the area, and the old value
+ * 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,
+	};
+
+	RegisterShmemCallbacks(&callbacks);
+	if (attached == NULL)
+		elog(ERROR, "could not attach to shared memory");
+
+	old_value = *attached;
+	*attached = new_value;
+
+	PG_RETURN_INT32(old_value);
+}