pg_pacs/pacs.c

/*--------------------------------------------------------------------------
 *
 * pacs.c
 *		Code for easier management of "measures".
 *
 * Portions Copyright (c) 1996-2024, PostgreSQL Global Development Group
 * Portions Copyright (c) 1994, Regents of the University of California
 *
 * -------------------------------------------------------------------------
 */

#include "postgres.h"
#include "fmgr.h"

#include "common/hashfn.h"
#include "pgstat.h"
#include "utils/builtins.h"
#include "utils/pgstat_internal.h"

#include "pacs.h"

PG_MODULE_MAGIC;

/* structure for measure */
typedef struct MeasureStatEntry
{
    PgStat_Counter measure;
} MeasureStatEntry;

/* boilerplace structure for all measures */
typedef struct MeasuresStatShared
{
    PgStatShared_Common header;
    MeasureStatEntry stats;
} MeasuresStatShared;

/* boilerplate declaration */
static bool measure_flush_cb(PgStat_EntryRef *entry_ref, bool nowait);

/* declare our measures stats */
static const PgStat_KindInfo measures = {
    .name = "measures",
    .fixed_amount = false,

    /* Measures are system-wide */
    .accessed_across_databases = true,

    .shared_size = sizeof(MeasuresStatShared),
    .shared_data_off = offsetof(MeasuresStatShared, stats),
    .shared_data_len = sizeof(((MeasuresStatShared *) 0)->stats),
    .pending_size = sizeof(MeasureStatEntry),
    .flush_pending_cb = measure_flush_cb,
};

/*
 * Compute stats entry idx from measure name with an 8-byte hash.
 */
#define PGSTAT_MEASURE_IDX(name) hash_bytes_extended((const unsigned char *) name, strlen(name), 0)

/*
 * Kind ID reserved for statistics of measures.
 */
#define PGSTAT_KIND_MEASURE	(PGSTAT_KIND_CUSTOM_MIN + 1)

/* Track if stats are loaded */
static bool measures_loaded = false;

void
_PG_init(void)
{
	if (!process_shared_preload_libraries_in_progress)
		return;

   ereport(LOG,
            (errmsg("Module pacs is starting"),
             errdetail("providing advanced stats structures.")));
    pgstat_register_kind(PGSTAT_KIND_MEASURE, &measures);

    /* mark stats as loaded */
    measures_loaded = true;
}

/*
 * Callback for flushing a measure
 */
static bool
measure_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
{
    MeasureStatEntry *localent;
    MeasuresStatShared *shfuncent;

    localent = (MeasureStatEntry *) entry_ref->pending;
    shfuncent = (MeasuresStatShared *) entry_ref->shared_stats;

    if (!pgstat_lock_entry(entry_ref, nowait))
        return false;

    shfuncent->stats.measure += localent->measure;
    return true;
}

static void
check_loaded()
{
    if (!measures_loaded)
        ereport(ERROR,
                (errcode(ERRCODE_UNDEFINED_OBJECT),
                 errmsg("measures module not loaded"),
                 errhint("check shared_preload_libraries")));
}

/*
 * The following code contains:
 * - the exported C functions
 * - the SQL API
 */
/*
 * Create a measure
 */
PGDLLEXPORT void
pacsCreateMeasure(const char *name)
{
    PgStat_EntryRef *entry_ref;
    MeasuresStatShared *shstatent;

    /* error if disabled */
    check_loaded();

    entry_ref = pgstat_get_entry_ref_locked(PGSTAT_KIND_MEASURE, InvalidOid,
                                            PGSTAT_MEASURE_IDX(name), false);
    shstatent = (MeasuresStatShared *) entry_ref->shared_stats;

    /* initialize shared memory data */
    memset(&shstatent->stats, 0, sizeof(shstatent->stats));
    pgstat_unlock_entry(entry_ref);
}

/*
 * SQL API
 */
PG_FUNCTION_INFO_V1(pacs_create_measure);
Datum
pacs_create_measure(PG_FUNCTION_ARGS)
{
    Name              name;

    name = PG_GETARG_NAME(0);
    pacsCreateMeasure(NameStr(*name));
    PG_RETURN_VOID();
}

/*
 * Drop a measure
 *
 * XXX register to create, drop to unregister...
 */
PGDLLEXPORT void
pacsDropMeasure(const char *name)
{
    /* error if disabled */
    check_loaded();

    /* drop and on faillure inform for later garbage collection */
    if (!pgstat_drop_entry(PGSTAT_KIND_MEASURE, InvalidOid,
                           PGSTAT_MEASURE_IDX(name)))
        pgstat_request_entry_refs_gc();
}

PG_FUNCTION_INFO_V1(pacs_drop_measure);
Datum
pacs_drop_measure(PG_FUNCTION_ARGS)
{
    Name              name;

    name = PG_GETARG_NAME(0);
    pacsDropMeasure(NameStr(*name));
    PG_RETURN_VOID();
}

/*
 * Increment a measure
 *
 * By design, PostgreSQL does create the measure if absent!
 *
 * we do not manage ereport() here to not add overhead.
 * XXX is it accurate ?
 */
PGDLLEXPORT void
pacsIncrementMeasure(const char *name)
{
    PgStat_EntryRef *entry_ref;
    MeasuresStatShared *shstatent;
    MeasureStatEntry *statent;

    /* quick leave if disabled */
    if (!measures_loaded)
      return;

    entry_ref = pgstat_get_entry_ref_locked(PGSTAT_KIND_MEASURE, InvalidOid,
                                            PGSTAT_MEASURE_IDX(name), false);

    shstatent = (MeasuresStatShared *) entry_ref->shared_stats;
    statent = &shstatent->stats;

    /* Update the measure */
    statent->measure++;

    pgstat_unlock_entry(entry_ref);
}

PG_FUNCTION_INFO_V1(pacs_increment_measure);
Datum
pacs_increment_measure(PG_FUNCTION_ARGS)
{
    Name              name;

    name = PG_GETARG_NAME(0);
    pacsIncrementMeasure(NameStr(*name));
    PG_RETURN_VOID();
}

/*
 * Support function for the SQL-callable functions. Returns
 * a pointer to the measure struct.
 */
static MeasureStatEntry *
pacsGetMeasure(const char *name)
{
    MeasureStatEntry *entry = NULL;

    /* error if disabled */
    check_loaded();

    /* Compile the lookup key as a hash of the measure name */
    entry = (MeasureStatEntry *) pgstat_fetch_entry(PGSTAT_KIND_MEASURE,
                                                    InvalidOid,
                                                    PGSTAT_MEASURE_IDX(name));
    return entry;
}

/*
 * SQL function returning the measure current value.
 */
PG_FUNCTION_INFO_V1(pacs_get_measure);
Datum
pacs_get_measure(PG_FUNCTION_ARGS)
{
    MeasureStatEntry  *entry;
    Name              name;

    name = PG_GETARG_NAME(0);

    entry = pacsGetMeasure(NameStr(*name));
    if (entry == NULL)
        ereport(ERROR,
                (errcode(ERRCODE_UNDEFINED_OBJECT),
                errmsg("measure \"%s\" does not exist",
                       NameStr(*name)),
                errhint("ensure the measure name is correct or "
                        "$$ select create_measure('%s')$$ to create it",
                        NameStr(*name))));

    PG_RETURN_INT64(entry->measure);
}

/*
 * Reset a measure
 */
PGDLLEXPORT void
pacsResetMeasure(const char *name)
{
    /* error if disabled */
    check_loaded();
    pgstat_reset_entry(PGSTAT_KIND_MEASURE, InvalidOid,
                       PGSTAT_MEASURE_IDX(name), 0);
    return;
}

/*
 * SQL function reseting the measure(s)
 */
PG_FUNCTION_INFO_V1(pacs_reset_measure);
Datum
pacs_reset_measure(PG_FUNCTION_ARGS)
{
    Name    name;
    name = PG_GETARG_NAME(0);
    pacsResetMeasure(NameStr(*name));
    PG_RETURN_VOID();
}

/*
 * Reset a measure
 */
PGDLLEXPORT void
pacsResetAllMeasures(void)
{
    /* error if disabled */
    check_loaded();
    pgstat_reset_of_kind(PGSTAT_KIND_MEASURE);
    return;
}

PG_FUNCTION_INFO_V1(pacs_reset_all_measures);
Datum
pacs_reset_all_measures(PG_FUNCTION_ARGS)
{
    pacsResetAllMeasures();
    PG_RETURN_VOID();
}

// PG_FUNCTION_INFO_V1(pacs_list_all_measures);
// PGDLLEXPORT void
// pacs_list_all_measures(PgStat_Kind kind)
// {
//     HASH_SEQ_STATUS status;
//     PgStat_EntryRef *entry_ref;
//
//     const PgStat_KindInfo *kind_info = pgstat_get_kind_info(kind);
//     if (!kind_info || !kind_info->hash_table)
//     {
//         elog(ERROR, "Invalid or unsupported statistics kind.");
//         return;
//     }
//
//     hash_seq_init(&status, kind_info->hash_table);
//
//     while ((entry_ref = hash_seq_search(&status)) != NULL)
//     {
//         elog(INFO, "Key for kind %d: %u/%llu", kind,
//             entry_ref->key.dboid, (unsigned long long)entry_ref->key.objid);
//     }
// }