Line data    Source code

       1             : /*--------------------------------------------------------------------------
       2             :  *
       3             :  * test_custom_fixed_stats.c
       4             :  *      Test module for fixed-sized custom pgstats
       5             :  *
       6             :  * Copyright (c) 2025, PostgreSQL Global Development Group
       7             :  *
       8             :  * IDENTIFICATION
       9             :  *      src/test/modules/test_custom_stats/test_custom_fixed_stats.c
      10             :  *
      11             :  * -------------------------------------------------------------------------
      12             :  */
      13             : 
      14             : #include "postgres.h"
      15             : 
      16             : #include "access/htup_details.h"
      17             : #include "funcapi.h"
      18             : #include "pgstat.h"
      19             : #include "utils/builtins.h"
      20             : #include "utils/pgstat_internal.h"
      21             : 
      22           6 : PG_MODULE_MAGIC_EXT(
      23             :                     .name = "test_custom_fixed_stats",
      24             :                     .version = PG_VERSION
      25             : );
      26             : 
      27             : /* Fixed-amount custom statistics entry */
      28             : typedef struct PgStat_StatCustomFixedEntry
      29             : {
      30             :     PgStat_Counter numcalls;    /* # of times update function called */
      31             :     TimestampTz stat_reset_timestamp;
      32             : } PgStat_StatCustomFixedEntry;
      33             : 
      34             : typedef struct PgStatShared_CustomFixedEntry
      35             : {
      36             :     LWLock      lock;           /* protects counters */
      37             :     uint32      changecount;    /* for atomic reads */
      38             :     PgStat_StatCustomFixedEntry stats;  /* current counters */
      39             :     PgStat_StatCustomFixedEntry reset_offset;   /* reset baseline */
      40             : } PgStatShared_CustomFixedEntry;
      41             : 
      42             : /* Callbacks for fixed-amount statistics */
      43             : static void test_custom_stats_fixed_init_shmem_cb(void *stats);
      44             : static void test_custom_stats_fixed_reset_all_cb(TimestampTz ts);
      45             : static void test_custom_stats_fixed_snapshot_cb(void);
      46             : 
      47             : static const PgStat_KindInfo custom_stats = {
      48             :     .name = "test_custom_fixed_stats",
      49             :     .fixed_amount = true,       /* exactly one entry */
      50             :     .write_to_file = true,      /* persist to stats file */
      51             : 
      52             :     .shared_size = sizeof(PgStat_StatCustomFixedEntry),
      53             :     .shared_data_off = offsetof(PgStatShared_CustomFixedEntry, stats),
      54             :     .shared_data_len = sizeof(((PgStatShared_CustomFixedEntry *) 0)->stats),
      55             : 
      56             :     .init_shmem_cb = test_custom_stats_fixed_init_shmem_cb,
      57             :     .reset_all_cb = test_custom_stats_fixed_reset_all_cb,
      58             :     .snapshot_cb = test_custom_stats_fixed_snapshot_cb,
      59             : };
      60             : 
      61             : /*
      62             :  * Kind ID for test_custom_fixed_stats.
      63             :  */
      64             : #define PGSTAT_KIND_TEST_CUSTOM_FIXED_STATS 26
      65             : 
      66             : /*--------------------------------------------------------------------------
      67             :  * Module initialization
      68             :  *--------------------------------------------------------------------------
      69             :  */
      70             : 
      71             : void
      72           6 : _PG_init(void)
      73             : {
      74             :     /* Must be loaded via shared_preload_libraries */
      75           6 :     if (!process_shared_preload_libraries_in_progress)
      76           0 :         return;
      77             : 
      78             :     /* Register custom statistics kind */
      79           6 :     pgstat_register_kind(PGSTAT_KIND_TEST_CUSTOM_FIXED_STATS, &custom_stats);
      80             : }
      81             : 
      82             : /*
      83             :  * test_custom_stats_fixed_init_shmem_cb
      84             :  *      Initialize shared memory structure
      85             :  */
      86             : static void
      87           6 : test_custom_stats_fixed_init_shmem_cb(void *stats)
      88             : {
      89           6 :     PgStatShared_CustomFixedEntry *stats_shmem =
      90             :         (PgStatShared_CustomFixedEntry *) stats;
      91             : 
      92           6 :     LWLockInitialize(&stats_shmem->lock, LWTRANCHE_PGSTATS_DATA);
      93           6 : }
      94             : 
      95             : /*
      96             :  * test_custom_stats_fixed_reset_all_cb
      97             :  *      Reset the fixed-sized stats
      98             :  */
      99             : static void
     100           2 : test_custom_stats_fixed_reset_all_cb(TimestampTz ts)
     101             : {
     102             :     PgStatShared_CustomFixedEntry *stats_shmem =
     103           2 :         pgstat_get_custom_shmem_data(PGSTAT_KIND_TEST_CUSTOM_FIXED_STATS);
     104             : 
     105             :     /* see explanation above PgStatShared_Archiver for the reset protocol */
     106           2 :     LWLockAcquire(&stats_shmem->lock, LW_EXCLUSIVE);
     107           2 :     pgstat_copy_changecounted_stats(&stats_shmem->reset_offset,
     108           2 :                                     &stats_shmem->stats,
     109             :                                     sizeof(stats_shmem->stats),
     110             :                                     &stats_shmem->changecount);
     111           2 :     stats_shmem->stats.stat_reset_timestamp = ts;
     112           2 :     LWLockRelease(&stats_shmem->lock);
     113           2 : }
     114             : 
     115             : /*
     116             :  * test_custom_stats_fixed_snapshot_cb
     117             :  *      Copy current stats to snapshot area
     118             :  */
     119             : static void
     120           6 : test_custom_stats_fixed_snapshot_cb(void)
     121             : {
     122             :     PgStatShared_CustomFixedEntry *stats_shmem =
     123           6 :         pgstat_get_custom_shmem_data(PGSTAT_KIND_TEST_CUSTOM_FIXED_STATS);
     124             :     PgStat_StatCustomFixedEntry *stat_snap =
     125           6 :         pgstat_get_custom_snapshot_data(PGSTAT_KIND_TEST_CUSTOM_FIXED_STATS);
     126           6 :     PgStat_StatCustomFixedEntry *reset_offset = &stats_shmem->reset_offset;
     127             :     PgStat_StatCustomFixedEntry reset;
     128             : 
     129           6 :     pgstat_copy_changecounted_stats(stat_snap,
     130           6 :                                     &stats_shmem->stats,
     131             :                                     sizeof(stats_shmem->stats),
     132             :                                     &stats_shmem->changecount);
     133             : 
     134           6 :     LWLockAcquire(&stats_shmem->lock, LW_SHARED);
     135           6 :     memcpy(&reset, reset_offset, sizeof(stats_shmem->stats));
     136           6 :     LWLockRelease(&stats_shmem->lock);
     137             : 
     138             :     /* Apply reset offsets */
     139             : #define FIXED_COMP(fld) stat_snap->fld -= reset.fld;
     140           6 :     FIXED_COMP(numcalls);
     141             : #undef FIXED_COMP
     142           6 : }
     143             : 
     144             : /*--------------------------------------------------------------------------
     145             :  * SQL-callable functions
     146             :  *--------------------------------------------------------------------------
     147             :  */
     148             : 
     149             : /*
     150             :  * test_custom_stats_fixed_update
     151             :  *      Increment call counter
     152             :  */
     153          14 : PG_FUNCTION_INFO_V1(test_custom_stats_fixed_update);
     154             : Datum
     155          12 : test_custom_stats_fixed_update(PG_FUNCTION_ARGS)
     156             : {
     157             :     PgStatShared_CustomFixedEntry *stats_shmem;
     158             : 
     159          12 :     stats_shmem = pgstat_get_custom_shmem_data(PGSTAT_KIND_TEST_CUSTOM_FIXED_STATS);
     160             : 
     161          12 :     LWLockAcquire(&stats_shmem->lock, LW_EXCLUSIVE);
     162             : 
     163          12 :     pgstat_begin_changecount_write(&stats_shmem->changecount);
     164          12 :     stats_shmem->stats.numcalls++;
     165          12 :     pgstat_end_changecount_write(&stats_shmem->changecount);
     166             : 
     167          12 :     LWLockRelease(&stats_shmem->lock);
     168             : 
     169          12 :     PG_RETURN_VOID();
     170             : }
     171             : 
     172             : /*
     173             :  * test_custom_stats_fixed_reset
     174             :  *      Reset statistics by calling pgstat system
     175             :  */
     176           2 : PG_FUNCTION_INFO_V1(test_custom_stats_fixed_reset);
     177             : Datum
     178           0 : test_custom_stats_fixed_reset(PG_FUNCTION_ARGS)
     179             : {
     180           0 :     pgstat_reset_of_kind(PGSTAT_KIND_TEST_CUSTOM_FIXED_STATS);
     181             : 
     182           0 :     PG_RETURN_VOID();
     183             : }
     184             : 
     185             : /*
     186             :  * test_custom_stats_fixed_report
     187             :  *      Return current counter values
     188             :  */
     189           6 : PG_FUNCTION_INFO_V1(test_custom_stats_fixed_report);
     190             : Datum
     191           4 : test_custom_stats_fixed_report(PG_FUNCTION_ARGS)
     192             : {
     193             :     TupleDesc   tupdesc;
     194           4 :     Datum       values[2] = {0};
     195           4 :     bool        nulls[2] = {false};
     196             :     PgStat_StatCustomFixedEntry *stats;
     197             : 
     198             :     /* Take snapshot (applies reset offsets) */
     199           4 :     pgstat_snapshot_fixed(PGSTAT_KIND_TEST_CUSTOM_FIXED_STATS);
     200           4 :     stats = pgstat_get_custom_snapshot_data(PGSTAT_KIND_TEST_CUSTOM_FIXED_STATS);
     201             : 
     202             :     /* Build return tuple */
     203           4 :     tupdesc = CreateTemplateTupleDesc(2);
     204           4 :     TupleDescInitEntry(tupdesc, (AttrNumber) 1, "numcalls",
     205             :                        INT8OID, -1, 0);
     206           4 :     TupleDescInitEntry(tupdesc, (AttrNumber) 2, "stats_reset",
     207             :                        TIMESTAMPTZOID, -1, 0);
     208           4 :     BlessTupleDesc(tupdesc);
     209             : 
     210           4 :     values[0] = Int64GetDatum(stats->numcalls);
     211             : 
     212             :     /* Handle uninitialized timestamp (no reset yet) */
     213           4 :     if (stats->stat_reset_timestamp == 0)
     214             :     {
     215           2 :         nulls[1] = true;
     216             :     }
     217             :     else
     218             :     {
     219           2 :         values[1] = TimestampTzGetDatum(stats->stat_reset_timestamp);
     220             :     }
     221             : 
     222             :     /* Return as tuple */
     223           4 :     PG_RETURN_DATUM(HeapTupleGetDatum(heap_form_tuple(tupdesc, values, nulls)));
     224             : }