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