Line data    Source code

       1             : /*--------------------------------------------------------------------------
       2             :  *
       3             :  * test_dsm_registry.c
       4             :  *    Test the dynamic shared memory registry.
       5             :  *
       6             :  * Copyright (c) 2024, 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             : 
      19           4 : PG_MODULE_MAGIC;
      20             : 
      21             : typedef struct TestDSMRegistryStruct
      22             : {
      23             :     int         val;
      24             :     LWLock      lck;
      25             : } TestDSMRegistryStruct;
      26             : 
      27             : static TestDSMRegistryStruct *tdr_state;
      28             : 
      29             : static void
      30           2 : tdr_init_shmem(void *ptr)
      31             : {
      32           2 :     TestDSMRegistryStruct *state = (TestDSMRegistryStruct *) ptr;
      33             : 
      34           2 :     LWLockInitialize(&state->lck, LWLockNewTrancheId());
      35           2 :     state->val = 0;
      36           2 : }
      37             : 
      38             : static void
      39           4 : tdr_attach_shmem(void)
      40             : {
      41             :     bool        found;
      42             : 
      43           4 :     tdr_state = GetNamedDSMSegment("test_dsm_registry",
      44             :                                    sizeof(TestDSMRegistryStruct),
      45             :                                    tdr_init_shmem,
      46             :                                    &found);
      47           4 :     LWLockRegisterTranche(tdr_state->lck.tranche, "test_dsm_registry");
      48           4 : }
      49             : 
      50           4 : PG_FUNCTION_INFO_V1(set_val_in_shmem);
      51             : Datum
      52           2 : set_val_in_shmem(PG_FUNCTION_ARGS)
      53             : {
      54           2 :     tdr_attach_shmem();
      55             : 
      56           2 :     LWLockAcquire(&tdr_state->lck, LW_EXCLUSIVE);
      57           2 :     tdr_state->val = PG_GETARG_UINT32(0);
      58           2 :     LWLockRelease(&tdr_state->lck);
      59             : 
      60           2 :     PG_RETURN_VOID();
      61             : }
      62             : 
      63           4 : PG_FUNCTION_INFO_V1(get_val_in_shmem);
      64             : Datum
      65           2 : get_val_in_shmem(PG_FUNCTION_ARGS)
      66             : {
      67             :     int         ret;
      68             : 
      69           2 :     tdr_attach_shmem();
      70             : 
      71           2 :     LWLockAcquire(&tdr_state->lck, LW_SHARED);
      72           2 :     ret = tdr_state->val;
      73           2 :     LWLockRelease(&tdr_state->lck);
      74             : 
      75           2 :     PG_RETURN_UINT32(ret);
      76             : }