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            0 : 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            0 : init_tdr_dsm(void *ptr, void *arg)
      48              : {
      49            0 :         TestDSMRegistryStruct *dsm = (TestDSMRegistryStruct *) ptr;
      50              : 
      51            0 :         if ((int) (intptr_t) arg != 5432)
      52            0 :                 elog(ERROR, "unexpected arg value %d", (int) (intptr_t) arg);
      53              : 
      54            0 :         LWLockInitialize(&dsm->lck, LWLockNewTrancheId("test_dsm_registry"));
      55            0 :         dsm->val = 0;
      56            0 : }
      57              : 
      58              : static void
      59            0 : tdr_attach_shmem(void)
      60              : {
      61            0 :         bool            found;
      62              : 
      63            0 :         tdr_dsm = GetNamedDSMSegment("test_dsm_registry_dsm",
      64              :                                                                  sizeof(TestDSMRegistryStruct),
      65              :                                                                  init_tdr_dsm,
      66              :                                                                  &found, (void *) (intptr_t) 5432);
      67              : 
      68            0 :         if (tdr_dsa == NULL)
      69            0 :                 tdr_dsa = GetNamedDSA("test_dsm_registry_dsa", &found);
      70              : 
      71            0 :         if (tdr_hash == NULL)
      72            0 :                 tdr_hash = GetNamedDSHash("test_dsm_registry_hash", &dsh_params, &found);
      73            0 : }
      74              : 
      75            0 : PG_FUNCTION_INFO_V1(set_val_in_shmem);
      76              : Datum
      77            0 : set_val_in_shmem(PG_FUNCTION_ARGS)
      78              : {
      79            0 :         tdr_attach_shmem();
      80              : 
      81            0 :         LWLockAcquire(&tdr_dsm->lck, LW_EXCLUSIVE);
      82            0 :         tdr_dsm->val = PG_GETARG_INT32(0);
      83            0 :         LWLockRelease(&tdr_dsm->lck);
      84              : 
      85            0 :         PG_RETURN_VOID();
      86              : }
      87              : 
      88            0 : PG_FUNCTION_INFO_V1(get_val_in_shmem);
      89              : Datum
      90            0 : get_val_in_shmem(PG_FUNCTION_ARGS)
      91              : {
      92            0 :         int                     ret;
      93              : 
      94            0 :         tdr_attach_shmem();
      95              : 
      96            0 :         LWLockAcquire(&tdr_dsm->lck, LW_SHARED);
      97            0 :         ret = tdr_dsm->val;
      98            0 :         LWLockRelease(&tdr_dsm->lck);
      99              : 
     100            0 :         PG_RETURN_INT32(ret);
     101            0 : }
     102              : 
     103            0 : PG_FUNCTION_INFO_V1(set_val_in_hash);
     104              : Datum
     105            0 : set_val_in_hash(PG_FUNCTION_ARGS)
     106              : {
     107            0 :         TestDSMRegistryHashEntry *entry;
     108            0 :         char       *key = TextDatumGetCString(PG_GETARG_DATUM(0));
     109            0 :         char       *val = TextDatumGetCString(PG_GETARG_DATUM(1));
     110            0 :         bool            found;
     111              : 
     112            0 :         if (strlen(key) >= offsetof(TestDSMRegistryHashEntry, val))
     113            0 :                 ereport(ERROR,
     114              :                                 (errmsg("key too long")));
     115              : 
     116            0 :         tdr_attach_shmem();
     117              : 
     118            0 :         entry = dshash_find_or_insert(tdr_hash, key, &found);
     119            0 :         if (found)
     120            0 :                 dsa_free(tdr_dsa, entry->val);
     121              : 
     122            0 :         entry->val = dsa_allocate(tdr_dsa, strlen(val) + 1);
     123            0 :         strcpy(dsa_get_address(tdr_dsa, entry->val), val);
     124              : 
     125            0 :         dshash_release_lock(tdr_hash, entry);
     126              : 
     127            0 :         PG_RETURN_VOID();
     128            0 : }
     129              : 
     130            0 : PG_FUNCTION_INFO_V1(get_val_in_hash);
     131              : Datum
     132            0 : get_val_in_hash(PG_FUNCTION_ARGS)
     133              : {
     134            0 :         TestDSMRegistryHashEntry *entry;
     135            0 :         char       *key = TextDatumGetCString(PG_GETARG_DATUM(0));
     136            0 :         text       *val = NULL;
     137              : 
     138            0 :         tdr_attach_shmem();
     139              : 
     140            0 :         entry = dshash_find(tdr_hash, key, false);
     141            0 :         if (entry == NULL)
     142            0 :                 PG_RETURN_NULL();
     143              : 
     144            0 :         val = cstring_to_text(dsa_get_address(tdr_dsa, entry->val));
     145              : 
     146            0 :         dshash_release_lock(tdr_hash, entry);
     147              : 
     148            0 :         PG_RETURN_TEXT_P(val);
     149            0 : }