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