test_custom_fixed_stats.c

Go to the documentation of this file.

/*--------------------------------------------------------------------------
 *
 * test_custom_fixed_stats.c
 *      Test module for fixed-sized custom pgstats
 *
 * Copyright (c) 2025-2026, PostgreSQL Global Development Group
 *
 * IDENTIFICATION
 *      src/test/modules/test_custom_stats/test_custom_fixed_stats.c
 *
 * -------------------------------------------------------------------------
 */
 
#include "postgres.h"
 
#include "access/htup_details.h"
#include "funcapi.h"
#include "pgstat.h"
#include "utils/builtins.h"
#include "utils/pgstat_internal.h"
#include "utils/timestamp.h"
 
PG_MODULE_MAGIC_EXT(
                    .name = "test_custom_fixed_stats",
                    .version = PG_VERSION
);
 
/* Fixed-amount custom statistics entry */
typedef struct PgStat_StatCustomFixedEntry
{
    PgStat_Counter numcalls;    /* # of times update function called */
    TimestampTz stat_reset_timestamp;
} PgStat_StatCustomFixedEntry;
 
typedef struct PgStatShared_CustomFixedEntry
{
    LWLock      lock;           /* protects counters */
    uint32      changecount;    /* for atomic reads */
    PgStat_StatCustomFixedEntry stats;  /* current counters */
    PgStat_StatCustomFixedEntry reset_offset;   /* reset baseline */
} PgStatShared_CustomFixedEntry;
 
/* Callbacks for fixed-amount statistics */
static void test_custom_stats_fixed_init_shmem_cb(void *stats);
static void test_custom_stats_fixed_reset_all_cb(TimestampTz ts);
static void test_custom_stats_fixed_snapshot_cb(void);
 
static const PgStat_KindInfo custom_stats = {
    .name = "test_custom_fixed_stats",
    .fixed_amount = true,       /* exactly one entry */
    .write_to_file = true,      /* persist to stats file */
 
    .shared_size = sizeof(PgStat_StatCustomFixedEntry),
    .shared_data_off = offsetof(PgStatShared_CustomFixedEntry, stats),
    .shared_data_len = sizeof(((PgStatShared_CustomFixedEntry *) 0)->stats),
 
    .init_shmem_cb = test_custom_stats_fixed_init_shmem_cb,
    .reset_all_cb = test_custom_stats_fixed_reset_all_cb,
    .snapshot_cb = test_custom_stats_fixed_snapshot_cb,
};
 
/*
 * Kind ID for test_custom_fixed_stats.
 */
#define PGSTAT_KIND_TEST_CUSTOM_FIXED_STATS 26
 
/*--------------------------------------------------------------------------
 * Module initialization
 *--------------------------------------------------------------------------
 */
 
void
_PG_init(void)
{
    /* Must be loaded via shared_preload_libraries */
    if (!process_shared_preload_libraries_in_progress)
        return;
 
    /* Register custom statistics kind */
    pgstat_register_kind(PGSTAT_KIND_TEST_CUSTOM_FIXED_STATS, &custom_stats);
}
 
/*
 * test_custom_stats_fixed_init_shmem_cb
 *      Initialize shared memory structure
 */
static void
test_custom_stats_fixed_init_shmem_cb(void *stats)
{
    PgStatShared_CustomFixedEntry *stats_shmem =
        (PgStatShared_CustomFixedEntry *) stats;
 
    LWLockInitialize(&stats_shmem->lock, LWTRANCHE_PGSTATS_DATA);
}
 
/*
 * test_custom_stats_fixed_reset_all_cb
 *      Reset the fixed-sized stats
 */
static void
test_custom_stats_fixed_reset_all_cb(TimestampTz ts)
{
    PgStatShared_CustomFixedEntry *stats_shmem =
        pgstat_get_custom_shmem_data(PGSTAT_KIND_TEST_CUSTOM_FIXED_STATS);
 
    /* see explanation above PgStatShared_Archiver for the reset protocol */
    LWLockAcquire(&stats_shmem->lock, LW_EXCLUSIVE);
    pgstat_copy_changecounted_stats(&stats_shmem->reset_offset,
                                    &stats_shmem->stats,
                                    sizeof(stats_shmem->stats),
                                    &stats_shmem->changecount);
    stats_shmem->stats.stat_reset_timestamp = ts;
    LWLockRelease(&stats_shmem->lock);
}
 
/*
 * test_custom_stats_fixed_snapshot_cb
 *      Copy current stats to snapshot area
 */
static void
test_custom_stats_fixed_snapshot_cb(void)
{
    PgStatShared_CustomFixedEntry *stats_shmem =
        pgstat_get_custom_shmem_data(PGSTAT_KIND_TEST_CUSTOM_FIXED_STATS);
    PgStat_StatCustomFixedEntry *stat_snap =
        pgstat_get_custom_snapshot_data(PGSTAT_KIND_TEST_CUSTOM_FIXED_STATS);
    PgStat_StatCustomFixedEntry *reset_offset = &stats_shmem->reset_offset;
    PgStat_StatCustomFixedEntry reset;
 
    pgstat_copy_changecounted_stats(stat_snap,
                                    &stats_shmem->stats,
                                    sizeof(stats_shmem->stats),
                                    &stats_shmem->changecount);
 
    LWLockAcquire(&stats_shmem->lock, LW_SHARED);
    memcpy(&reset, reset_offset, sizeof(stats_shmem->stats));
    LWLockRelease(&stats_shmem->lock);
 
    /* Apply reset offsets */
#define FIXED_COMP(fld) stat_snap->fld -= reset.fld;
    FIXED_COMP(numcalls);
#undef FIXED_COMP
}
 
/*--------------------------------------------------------------------------
 * SQL-callable functions
 *--------------------------------------------------------------------------
 */
 
/*
 * test_custom_stats_fixed_update
 *      Increment call counter
 */
PG_FUNCTION_INFO_V1(test_custom_stats_fixed_update);
Datum
test_custom_stats_fixed_update(PG_FUNCTION_ARGS)
{
    PgStatShared_CustomFixedEntry *stats_shmem;
 
    stats_shmem = pgstat_get_custom_shmem_data(PGSTAT_KIND_TEST_CUSTOM_FIXED_STATS);
 
    LWLockAcquire(&stats_shmem->lock, LW_EXCLUSIVE);
 
    pgstat_begin_changecount_write(&stats_shmem->changecount);
    stats_shmem->stats.numcalls++;
    pgstat_end_changecount_write(&stats_shmem->changecount);
 
    LWLockRelease(&stats_shmem->lock);
 
    PG_RETURN_VOID();
}
 
/*
 * test_custom_stats_fixed_reset
 *      Reset statistics by calling pgstat system
 */
PG_FUNCTION_INFO_V1(test_custom_stats_fixed_reset);
Datum
test_custom_stats_fixed_reset(PG_FUNCTION_ARGS)
{
    pgstat_reset_of_kind(PGSTAT_KIND_TEST_CUSTOM_FIXED_STATS);
 
    PG_RETURN_VOID();
}
 
/*
 * test_custom_stats_fixed_report
 *      Return current counter values
 */
PG_FUNCTION_INFO_V1(test_custom_stats_fixed_report);
Datum
test_custom_stats_fixed_report(PG_FUNCTION_ARGS)
{
    TupleDesc   tupdesc;
    Datum       values[2] = {0};
    bool        nulls[2] = {false};
    PgStat_StatCustomFixedEntry *stats;
 
    /* Take snapshot (applies reset offsets) */
    pgstat_snapshot_fixed(PGSTAT_KIND_TEST_CUSTOM_FIXED_STATS);
    stats = pgstat_get_custom_snapshot_data(PGSTAT_KIND_TEST_CUSTOM_FIXED_STATS);
 
    /* Build return tuple */
    tupdesc = CreateTemplateTupleDesc(2);
    TupleDescInitEntry(tupdesc, (AttrNumber) 1, "numcalls",
                       INT8OID, -1, 0);
    TupleDescInitEntry(tupdesc, (AttrNumber) 2, "stats_reset",
                       TIMESTAMPTZOID, -1, 0);
    BlessTupleDesc(tupdesc);
 
    values[0] = Int64GetDatum(stats->numcalls);
 
    /* Handle uninitialized timestamp (no reset yet) */
    if (stats->stat_reset_timestamp == 0)
    {
        nulls[1] = true;
    }
    else
    {
        values[1] = TimestampTzGetDatum(stats->stat_reset_timestamp);
    }
 
    /* Return as tuple */
    PG_RETURN_DATUM(HeapTupleGetDatum(heap_form_tuple(tupdesc, values, nulls)));
}