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                 :             : 
      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                 :             : }