Branch data     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                 :             : #include "utils/guc.h"
      24                 :             : #include "utils/injection_point.h"
      25                 :             : 
      26                 :             : 
      27                 :           9 : PG_MODULE_MAGIC;
      28                 :             : 
      29                 :             : typedef struct TestShmemData
      30                 :             : {
      31                 :             :     bool        initialized;
      32                 :             :     int         attach_count;
      33                 :             :     char        dummy_data[FLEXIBLE_ARRAY_MEMBER];
      34                 :             : } TestShmemData;
      35                 :             : 
      36                 :             : static TestShmemData *TestShmem;
      37                 :             : 
      38                 :             : #define MIN_TEST_AREA_BYTES sizeof(TestShmemData)
      39                 :             : #define DEFAULT_TEST_AREA_BYTES MIN_TEST_AREA_BYTES
      40                 :             : #define MAX_TEST_AREA_BYTES 1000000
      41                 :             : 
      42                 :             : static bool attached_or_initialized = false;
      43                 :             : static int  test_shmem_area_size = MIN_TEST_AREA_BYTES;
      44                 :             : static bool test_shmem_guc_defined = false;
      45                 :             : 
      46                 :             : static void test_shmem_request(void *arg);
      47                 :             : static void test_shmem_init(void *arg);
      48                 :             : static void test_shmem_attach(void *arg);
      49                 :             : 
      50                 :             : static const ShmemCallbacks TestShmemCallbacks = {
      51                 :             :     .flags = SHMEM_CALLBACKS_ALLOW_AFTER_STARTUP,
      52                 :             :     .request_fn = test_shmem_request,
      53                 :             :     .init_fn = test_shmem_init,
      54                 :             :     .attach_fn = test_shmem_attach,
      55                 :             : };
      56                 :             : 
      57                 :             : static void
      58                 :           9 : test_shmem_request(void *arg)
      59                 :             : {
      60         [ +  - ]:           9 :     elog(LOG, "test_shmem_request callback called");
      61                 :             : 
      62                 :           9 :     ShmemRequestStruct(.name = "test_shmem area",
      63                 :             :                        .size = test_shmem_area_size,
      64                 :             :                        .ptr = (void **) &TestShmem);
      65                 :           9 : }
      66                 :             : 
      67                 :             : static void
      68                 :           6 : test_shmem_init(void *arg)
      69                 :             : {
      70         [ +  - ]:           6 :     elog(LOG, "init callback called");
      71                 :             : 
      72                 :           6 :     INJECTION_POINT("test-shmem-init", NULL);
      73                 :             : 
      74         [ -  + ]:           5 :     if (TestShmem->initialized)
      75         [ #  # ]:           0 :         elog(ERROR, "shmem area already initialized");
      76                 :           5 :     TestShmem->initialized = true;
      77                 :             : 
      78         [ -  + ]:           5 :     if (attached_or_initialized)
      79         [ #  # ]:           0 :         elog(ERROR, "attach or initialize already called in this process");
      80                 :           5 :     attached_or_initialized = true;
      81                 :           5 : }
      82                 :             : 
      83                 :             : static void
      84                 :           1 : test_shmem_attach(void *arg)
      85                 :             : {
      86         [ +  - ]:           1 :     elog(LOG, "test_shmem_attach callback called");
      87         [ -  + ]:           1 :     if (!TestShmem->initialized)
      88         [ #  # ]:           0 :         elog(ERROR, "shmem area not yet initialized");
      89                 :           1 :     TestShmem->attach_count++;
      90                 :             : 
      91         [ -  + ]:           1 :     if (attached_or_initialized)
      92         [ #  # ]:           0 :         elog(ERROR, "attach or initialize already called in this process");
      93                 :           1 :     attached_or_initialized = true;
      94                 :           1 : }
      95                 :             : 
      96                 :             : void
      97                 :           9 : _PG_init(void)
      98                 :             : {
      99         [ +  - ]:           9 :     elog(LOG, "test_shmem module's _PG_init called");
     100                 :             : 
     101         [ +  + ]:           9 :     if (!test_shmem_guc_defined)
     102                 :             :     {
     103                 :           8 :         DefineCustomIntVariable("test_shmem.area_size",
     104                 :             :                                 "Size of the shmem area to request.",
     105                 :             :                                 NULL,
     106                 :             :                                 &test_shmem_area_size,
     107                 :             :                                 DEFAULT_TEST_AREA_BYTES,
     108                 :             :                                 MIN_TEST_AREA_BYTES,
     109                 :             :                                 MAX_TEST_AREA_BYTES,
     110                 :             :                                 PGC_USERSET,
     111                 :             :                                 GUC_UNIT_BYTE,
     112                 :             :                                 NULL, NULL, NULL);
     113                 :           8 :         MarkGUCPrefixReserved("test_shmem");
     114                 :           8 :         test_shmem_guc_defined = true;
     115                 :             :     }
     116                 :           9 :     RegisterShmemCallbacks(&TestShmemCallbacks);
     117                 :           6 : }
     118                 :             : 
     119                 :           7 : PG_FUNCTION_INFO_V1(get_test_shmem_attach_count);
     120                 :             : Datum
     121                 :           6 : get_test_shmem_attach_count(PG_FUNCTION_ARGS)
     122                 :             : {
     123         [ -  + ]:           6 :     if (!attached_or_initialized)
     124         [ #  # ]:           0 :         elog(ERROR, "shmem area not attached or initialized in this process");
     125         [ -  + ]:           6 :     if (!TestShmem->initialized)
     126         [ #  # ]:           0 :         elog(ERROR, "shmem area not yet initialized");
     127                 :           6 :     PG_RETURN_INT32(TestShmem->attach_count);
     128                 :             : }