Line data    Source code

       1              : /*--------------------------------------------------------------------------
       2              :  *
       3              :  * test_dsm_registry.c
       4              :  *    Test the dynamic shared memory registry.
       5              :  *
       6              :  * Copyright (c) 2024-2026, PostgreSQL Global Development Group
       7              :  *
       8              :  * IDENTIFICATION
       9              :  *      src/test/modules/test_dsm_registry/test_dsm_registry.c
      10              :  *
      11              :  * -------------------------------------------------------------------------
      12              :  */
      13              : #include "postgres.h"
      14              : 
      15              : #include "fmgr.h"
      16              : #include "storage/dsm_registry.h"
      17              : #include "storage/lwlock.h"
      18              : #include "utils/builtins.h"
      19              : 
      20            2 : PG_MODULE_MAGIC;
      21              : 
      22              : typedef struct TestDSMRegistryStruct
      23              : {
      24              :     int         val;
      25              :     LWLock      lck;
      26              : } TestDSMRegistryStruct;
      27              : 
      28              : typedef struct TestDSMRegistryHashEntry
      29              : {
      30              :     char        key[64];
      31              :     dsa_pointer val;
      32              : } TestDSMRegistryHashEntry;
      33              : 
      34              : static TestDSMRegistryStruct *tdr_dsm;
      35              : static dsa_area *tdr_dsa;
      36              : static dshash_table *tdr_hash;
      37              : 
      38              : static const dshash_parameters dsh_params = {
      39              :     offsetof(TestDSMRegistryHashEntry, val),
      40              :     sizeof(TestDSMRegistryHashEntry),
      41              :     dshash_strcmp,
      42              :     dshash_strhash,
      43              :     dshash_strcpy
      44              : };
      45              : 
      46              : static void
      47            1 : init_tdr_dsm(void *ptr, void *arg)
      48              : {
      49            1 :     TestDSMRegistryStruct *dsm = (TestDSMRegistryStruct *) ptr;
      50              : 
      51            1 :     if ((int) (intptr_t) arg != 5432)
      52            0 :         elog(ERROR, "unexpected arg value %d", (int) (intptr_t) arg);
      53              : 
      54            1 :     LWLockInitialize(&dsm->lck, LWLockNewTrancheId("test_dsm_registry"));
      55            1 :     dsm->val = 0;
      56            1 : }
      57              : 
      58              : static void
      59            4 : tdr_attach_shmem(void)
      60              : {
      61              :     bool        found;
      62              : 
      63            4 :     tdr_dsm = GetNamedDSMSegment("test_dsm_registry_dsm",
      64              :                                  sizeof(TestDSMRegistryStruct),
      65              :                                  init_tdr_dsm,
      66              :                                  &found, (void *) (intptr_t) 5432);
      67              : 
      68            4 :     if (tdr_dsa == NULL)
      69            2 :         tdr_dsa = GetNamedDSA("test_dsm_registry_dsa", &found);
      70              : 
      71            4 :     if (tdr_hash == NULL)
      72            2 :         tdr_hash = GetNamedDSHash("test_dsm_registry_hash", &dsh_params, &found);
      73            4 : }
      74              : 
      75            2 : PG_FUNCTION_INFO_V1(set_val_in_shmem);
      76              : Datum
      77            1 : set_val_in_shmem(PG_FUNCTION_ARGS)
      78              : {
      79            1 :     tdr_attach_shmem();
      80              : 
      81            1 :     LWLockAcquire(&tdr_dsm->lck, LW_EXCLUSIVE);
      82            1 :     tdr_dsm->val = PG_GETARG_INT32(0);
      83            1 :     LWLockRelease(&tdr_dsm->lck);
      84              : 
      85            1 :     PG_RETURN_VOID();
      86              : }
      87              : 
      88            2 : PG_FUNCTION_INFO_V1(get_val_in_shmem);
      89              : Datum
      90            1 : get_val_in_shmem(PG_FUNCTION_ARGS)
      91              : {
      92              :     int         ret;
      93              : 
      94            1 :     tdr_attach_shmem();
      95              : 
      96            1 :     LWLockAcquire(&tdr_dsm->lck, LW_SHARED);
      97            1 :     ret = tdr_dsm->val;
      98            1 :     LWLockRelease(&tdr_dsm->lck);
      99              : 
     100            1 :     PG_RETURN_INT32(ret);
     101              : }
     102              : 
     103            2 : PG_FUNCTION_INFO_V1(set_val_in_hash);
     104              : Datum
     105            1 : set_val_in_hash(PG_FUNCTION_ARGS)
     106              : {
     107              :     TestDSMRegistryHashEntry *entry;
     108            1 :     char       *key = TextDatumGetCString(PG_GETARG_DATUM(0));
     109            1 :     char       *val = TextDatumGetCString(PG_GETARG_DATUM(1));
     110              :     bool        found;
     111              : 
     112            1 :     if (strlen(key) >= offsetof(TestDSMRegistryHashEntry, val))
     113            0 :         ereport(ERROR,
     114              :                 (errmsg("key too long")));
     115              : 
     116            1 :     tdr_attach_shmem();
     117              : 
     118            1 :     entry = dshash_find_or_insert(tdr_hash, key, &found);
     119            1 :     if (found)
     120            0 :         dsa_free(tdr_dsa, entry->val);
     121              : 
     122            1 :     entry->val = dsa_allocate(tdr_dsa, strlen(val) + 1);
     123            1 :     strcpy(dsa_get_address(tdr_dsa, entry->val), val);
     124              : 
     125            1 :     dshash_release_lock(tdr_hash, entry);
     126              : 
     127            1 :     PG_RETURN_VOID();
     128              : }
     129              : 
     130            2 : PG_FUNCTION_INFO_V1(get_val_in_hash);
     131              : Datum
     132            1 : get_val_in_hash(PG_FUNCTION_ARGS)
     133              : {
     134              :     TestDSMRegistryHashEntry *entry;
     135            1 :     char       *key = TextDatumGetCString(PG_GETARG_DATUM(0));
     136            1 :     text       *val = NULL;
     137              : 
     138            1 :     tdr_attach_shmem();
     139              : 
     140            1 :     entry = dshash_find(tdr_hash, key, false);
     141            1 :     if (entry == NULL)
     142            0 :         PG_RETURN_NULL();
     143              : 
     144            1 :     val = cstring_to_text(dsa_get_address(tdr_dsa, entry->val));
     145              : 
     146            1 :     dshash_release_lock(tdr_hash, entry);
     147              : 
     148            1 :     PG_RETURN_TEXT_P(val);
     149              : }