Line data    Source code

       1              : /*-------------------------------------------------------------------------
       2              :  *
       3              :  * test_shmem.c
       4              :  *      Helpers to test shmem allocation routines
       5              :  *
       6              :  * Test basic memory allocation in an extension module. One notable feature
       7              :  * that is not exercised by any other module in the repository is the
       8              :  * allocating (non-DSM) shared memory after postmaster startup.
       9              :  *
      10              :  * Copyright (c) 2020-2026, PostgreSQL Global Development Group
      11              :  *
      12              :  * IDENTIFICATION
      13              :  *    src/test/modules/test_shmem/test_shmem.c
      14              :  *
      15              :  *-------------------------------------------------------------------------
      16              :  */
      17              : 
      18              : #include "postgres.h"
      19              : 
      20              : #include "fmgr.h"
      21              : #include "miscadmin.h"
      22              : #include "storage/shmem.h"
      23              : 
      24              : 
      25            4 : PG_MODULE_MAGIC;
      26              : 
      27              : typedef struct TestShmemData
      28              : {
      29              :     int         value;
      30              :     bool        initialized;
      31              :     int         attach_count;
      32              : } TestShmemData;
      33              : 
      34              : static TestShmemData *TestShmem;
      35              : 
      36              : static bool attached_or_initialized = false;
      37              : 
      38              : static void test_shmem_request(void *arg);
      39              : static void test_shmem_init(void *arg);
      40              : static void test_shmem_attach(void *arg);
      41              : 
      42              : static const ShmemCallbacks TestShmemCallbacks = {
      43              :     .flags = SHMEM_CALLBACKS_ALLOW_AFTER_STARTUP,
      44              :     .request_fn = test_shmem_request,
      45              :     .init_fn = test_shmem_init,
      46              :     .attach_fn = test_shmem_attach,
      47              : };
      48              : 
      49              : static void
      50            4 : test_shmem_request(void *arg)
      51              : {
      52            4 :     elog(LOG, "test_shmem_request callback called");
      53              : 
      54            4 :     ShmemRequestStruct(.name = "test_shmem area",
      55              :                        .size = sizeof(TestShmemData),
      56              :                        .ptr = (void **) &TestShmem);
      57            4 : }
      58              : 
      59              : static void
      60            2 : test_shmem_init(void *arg)
      61              : {
      62            2 :     elog(LOG, "init callback called");
      63            2 :     if (TestShmem->initialized)
      64            0 :         elog(ERROR, "shmem area already initialized");
      65            2 :     TestShmem->initialized = true;
      66              : 
      67            2 :     if (attached_or_initialized)
      68            0 :         elog(ERROR, "attach or initialize already called in this process");
      69            2 :     attached_or_initialized = true;
      70            2 : }
      71              : 
      72              : static void
      73            2 : test_shmem_attach(void *arg)
      74              : {
      75            2 :     elog(LOG, "test_shmem_attach callback called");
      76            2 :     if (!TestShmem->initialized)
      77            0 :         elog(ERROR, "shmem area not yet initialized");
      78            2 :     TestShmem->attach_count++;
      79              : 
      80            2 :     if (attached_or_initialized)
      81            0 :         elog(ERROR, "attach or initialize already called in this process");
      82            2 :     attached_or_initialized = true;
      83            2 : }
      84              : 
      85              : void
      86            4 : _PG_init(void)
      87              : {
      88            4 :     elog(LOG, "test_shmem module's _PG_init called");
      89            4 :     RegisterShmemCallbacks(&TestShmemCallbacks);
      90            4 : }
      91              : 
      92            5 : PG_FUNCTION_INFO_V1(get_test_shmem_attach_count);
      93              : Datum
      94            4 : get_test_shmem_attach_count(PG_FUNCTION_ARGS)
      95              : {
      96            4 :     if (!attached_or_initialized)
      97            0 :         elog(ERROR, "shmem area not attached or initialized in this process");
      98            4 :     if (!TestShmem->initialized)
      99            0 :         elog(ERROR, "shmem area not yet initialized");
     100            4 :     PG_RETURN_INT32(TestShmem->attach_count);
     101              : }