injection_point.c

Go to the documentation of this file.

/*-------------------------------------------------------------------------
 *
 * injection_point.c
 *    Routines to control and run injection points in the code.
 *
 * Injection points can be used to run arbitrary code by attaching callbacks
 * that would be executed in place of the named injection point.
 *
 * Portions Copyright (c) 1996-2024, PostgreSQL Global Development Group
 * Portions Copyright (c) 1994, Regents of the University of California
 *
 *
 * IDENTIFICATION
 *    src/backend/utils/misc/injection_point.c
 *
 *-------------------------------------------------------------------------
 */
#include "postgres.h"
 
#include <sys/stat.h>
 
#include "fmgr.h"
#include "miscadmin.h"
#include "storage/fd.h"
#include "storage/lwlock.h"
#include "storage/shmem.h"
#include "utils/hsearch.h"
#include "utils/injection_point.h"
#include "utils/memutils.h"
 
#ifdef USE_INJECTION_POINTS
 
/* Field sizes */
#define INJ_NAME_MAXLEN     64
#define INJ_LIB_MAXLEN      128
#define INJ_FUNC_MAXLEN     128
#define INJ_PRIVATE_MAXLEN  1024
 
/* Single injection point stored in shared memory */
typedef struct InjectionPointEntry
{
    /*
     * Because injection points need to be usable without LWLocks, we use a
     * generation counter on each entry to allow safe, lock-free reading.
     *
     * To read an entry, first read the current 'generation' value.  If it's
     * even, then the slot is currently unused, and odd means it's in use.
     * When reading the other fields, beware that they may change while
     * reading them, if the entry is released and reused!  After reading the
     * other fields, read 'generation' again: if its value hasn't changed, you
     * can be certain that the other fields you read are valid.  Otherwise,
     * the slot was concurrently recycled, and you should ignore it.
     *
     * When adding an entry, you must store all the other fields first, and
     * then update the generation number, with an appropriate memory barrier
     * in between. In addition to that protocol, you must also hold
     * InjectionPointLock, to prevent two backends from modifying the array at
     * the same time.
     */
    pg_atomic_uint64 generation;
 
    char        name[INJ_NAME_MAXLEN];  /* hash key */
    char        library[INJ_LIB_MAXLEN];    /* library */
    char        function[INJ_FUNC_MAXLEN];  /* function */
 
    /*
     * Opaque data area that modules can use to pass some custom data to
     * callbacks, registered when attached.
     */
    char        private_data[INJ_PRIVATE_MAXLEN];
} InjectionPointEntry;
 
#define MAX_INJECTION_POINTS    128
 
/*
 * Shared memory array of active injection points.
 *
 * 'max_inuse' is the highest index currently in use, plus one.  It's just an
 * optimization to avoid scanning through the whole entry, in the common case
 * that there are no injection points, or only a few.
 */
typedef struct InjectionPointsCtl
{
    pg_atomic_uint32 max_inuse;
    InjectionPointEntry entries[MAX_INJECTION_POINTS];
} InjectionPointsCtl;
 
NON_EXEC_STATIC InjectionPointsCtl *ActiveInjectionPoints;
 
/*
 * Backend local cache of injection callbacks already loaded, stored in
 * TopMemoryContext.
 */
typedef struct InjectionPointCacheEntry
{
    char        name[INJ_NAME_MAXLEN];
    char        private_data[INJ_PRIVATE_MAXLEN];
    InjectionPointCallback callback;
 
    /*
     * Shmem slot and copy of its generation number when this cache entry was
     * created.  They can be used to validate if the cached entry is still
     * valid.
     */
    int         slot_idx;
    uint64      generation;
} InjectionPointCacheEntry;
 
static HTAB *InjectionPointCache = NULL;
 
/*
 * injection_point_cache_add
 *
 * Add an injection point to the local cache.
 */
static InjectionPointCacheEntry *
injection_point_cache_add(const char *name,
                          int slot_idx,
                          uint64 generation,
                          InjectionPointCallback callback,
                          const void *private_data)
{
    InjectionPointCacheEntry *entry;
    bool        found;
 
    /* If first time, initialize */
    if (InjectionPointCache == NULL)
    {
        HASHCTL     hash_ctl;
 
        hash_ctl.keysize = sizeof(char[INJ_NAME_MAXLEN]);
        hash_ctl.entrysize = sizeof(InjectionPointCacheEntry);
        hash_ctl.hcxt = TopMemoryContext;
 
        InjectionPointCache = hash_create("InjectionPoint cache hash",
                                          MAX_INJECTION_POINTS,
                                          &hash_ctl,
                                          HASH_ELEM | HASH_STRINGS | HASH_CONTEXT);
    }
 
    entry = (InjectionPointCacheEntry *)
        hash_search(InjectionPointCache, name, HASH_ENTER, &found);
 
    Assert(!found);
    strlcpy(entry->name, name, sizeof(entry->name));
    entry->slot_idx = slot_idx;
    entry->generation = generation;
    entry->callback = callback;
    memcpy(entry->private_data, private_data, INJ_PRIVATE_MAXLEN);
 
    return entry;
}
 
/*
 * injection_point_cache_remove
 *
 * Remove entry from the local cache.  Note that this leaks a callback
 * loaded but removed later on, which should have no consequence from
 * a testing perspective.
 */
static void
injection_point_cache_remove(const char *name)
{
    bool        found PG_USED_FOR_ASSERTS_ONLY;
 
    (void) hash_search(InjectionPointCache, name, HASH_REMOVE, &found);
    Assert(found);
}
 
/*
 * injection_point_cache_load
 *
 * Load an injection point into the local cache.
 */
static InjectionPointCacheEntry *
injection_point_cache_load(InjectionPointEntry *entry, int slot_idx, uint64 generation)
{
    char        path[MAXPGPATH];
    void       *injection_callback_local;
 
    snprintf(path, MAXPGPATH, "%s/%s%s", pkglib_path,
             entry->library, DLSUFFIX);
 
    if (!pg_file_exists(path))
        elog(ERROR, "could not find library \"%s\" for injection point \"%s\"",
             path, entry->name);
 
    injection_callback_local = (void *)
        load_external_function(path, entry->function, false, NULL);
 
    if (injection_callback_local == NULL)
        elog(ERROR, "could not find function \"%s\" in library \"%s\" for injection point \"%s\"",
             entry->function, path, entry->name);
 
    /* add it to the local cache */
    return injection_point_cache_add(entry->name,
                                     slot_idx,
                                     generation,
                                     injection_callback_local,
                                     entry->private_data);
}
 
/*
 * injection_point_cache_get
 *
 * Retrieve an injection point from the local cache, if any.
 */
static InjectionPointCacheEntry *
injection_point_cache_get(const char *name)
{
    bool        found;
    InjectionPointCacheEntry *entry;
 
    /* no callback if no cache yet */
    if (InjectionPointCache == NULL)
        return NULL;
 
    entry = (InjectionPointCacheEntry *)
        hash_search(InjectionPointCache, name, HASH_FIND, &found);
 
    if (found)
        return entry;
 
    return NULL;
}
#endif                          /* USE_INJECTION_POINTS */
 
/*
 * Return the space for dynamic shared hash table.
 */
Size
InjectionPointShmemSize(void)
{
#ifdef USE_INJECTION_POINTS
    Size        sz = 0;
 
    sz = add_size(sz, sizeof(InjectionPointsCtl));
    return sz;
#else
    return 0;
#endif
}
 
/*
 * Allocate shmem space for dynamic shared hash.
 */
void
InjectionPointShmemInit(void)
{
#ifdef USE_INJECTION_POINTS
    bool        found;
 
    ActiveInjectionPoints = ShmemInitStruct("InjectionPoint hash",
                                            sizeof(InjectionPointsCtl),
                                            &found);
    if (!IsUnderPostmaster)
    {
        Assert(!found);
        pg_atomic_init_u32(&ActiveInjectionPoints->max_inuse, 0);
        for (int i = 0; i < MAX_INJECTION_POINTS; i++)
            pg_atomic_init_u64(&ActiveInjectionPoints->entries[i].generation, 0);
    }
    else
        Assert(found);
#endif
}
 
/*
 * Attach a new injection point.
 */
void
InjectionPointAttach(const char *name,
                     const char *library,
                     const char *function,
                     const void *private_data,
                     int private_data_size)
{
#ifdef USE_INJECTION_POINTS
    InjectionPointEntry *entry;
    uint64      generation;
    uint32      max_inuse;
    int         free_idx;
 
    if (strlen(name) >= INJ_NAME_MAXLEN)
        elog(ERROR, "injection point name %s too long (maximum of %u)",
             name, INJ_NAME_MAXLEN);
    if (strlen(library) >= INJ_LIB_MAXLEN)
        elog(ERROR, "injection point library %s too long (maximum of %u)",
             library, INJ_LIB_MAXLEN);
    if (strlen(function) >= INJ_FUNC_MAXLEN)
        elog(ERROR, "injection point function %s too long (maximum of %u)",
             function, INJ_FUNC_MAXLEN);
    if (private_data_size >= INJ_PRIVATE_MAXLEN)
        elog(ERROR, "injection point data too long (maximum of %u)",
             INJ_PRIVATE_MAXLEN);
 
    /*
     * Allocate and register a new injection point.  A new point should not
     * exist.  For testing purposes this should be fine.
     */
    LWLockAcquire(InjectionPointLock, LW_EXCLUSIVE);
    max_inuse = pg_atomic_read_u32(&ActiveInjectionPoints->max_inuse);
    free_idx = -1;
 
    for (int idx = 0; idx < max_inuse; idx++)
    {
        entry = &ActiveInjectionPoints->entries[idx];
        generation = pg_atomic_read_u64(&entry->generation);
        if (generation % 2 == 0)
        {
            /*
             * Found a free slot where we can add the new entry, but keep
             * going so that we will find out if the entry already exists.
             */
            if (free_idx == -1)
                free_idx = idx;
        }
 
        if (strcmp(entry->name, name) == 0)
            elog(ERROR, "injection point \"%s\" already defined", name);
    }
    if (free_idx == -1)
    {
        if (max_inuse == MAX_INJECTION_POINTS)
            elog(ERROR, "too many injection points");
        free_idx = max_inuse;
    }
    entry = &ActiveInjectionPoints->entries[free_idx];
    generation = pg_atomic_read_u64(&entry->generation);
    Assert(generation % 2 == 0);
 
    /* Save the entry */
    strlcpy(entry->name, name, sizeof(entry->name));
    entry->name[INJ_NAME_MAXLEN - 1] = '\0';
    strlcpy(entry->library, library, sizeof(entry->library));
    entry->library[INJ_LIB_MAXLEN - 1] = '\0';
    strlcpy(entry->function, function, sizeof(entry->function));
    entry->function[INJ_FUNC_MAXLEN - 1] = '\0';
    if (private_data != NULL)
        memcpy(entry->private_data, private_data, private_data_size);
 
    pg_write_barrier();
    pg_atomic_write_u64(&entry->generation, generation + 1);
 
    if (free_idx + 1 > max_inuse)
        pg_atomic_write_u32(&ActiveInjectionPoints->max_inuse, free_idx + 1);
 
    LWLockRelease(InjectionPointLock);
 
#else
    elog(ERROR, "injection points are not supported by this build");
#endif
}
 
/*
 * Detach an existing injection point.
 *
 * Returns true if the injection point was detached, false otherwise.
 */
bool
InjectionPointDetach(const char *name)
{
#ifdef USE_INJECTION_POINTS
    bool        found = false;
    int         idx;
    int         max_inuse;
 
    LWLockAcquire(InjectionPointLock, LW_EXCLUSIVE);
 
    /* Find it in the shmem array, and mark the slot as unused */
    max_inuse = (int) pg_atomic_read_u32(&ActiveInjectionPoints->max_inuse);
    for (idx = max_inuse - 1; idx >= 0; --idx)
    {
        InjectionPointEntry *entry = &ActiveInjectionPoints->entries[idx];
        uint64      generation;
 
        generation = pg_atomic_read_u64(&entry->generation);
        if (generation % 2 == 0)
            continue;           /* empty slot */
 
        if (strcmp(entry->name, name) == 0)
        {
            Assert(!found);
            found = true;
            pg_atomic_write_u64(&entry->generation, generation + 1);
            break;
        }
    }
 
    /* If we just removed the highest-numbered entry, update 'max_inuse' */
    if (found && idx == max_inuse - 1)
    {
        for (; idx >= 0; --idx)
        {
            InjectionPointEntry *entry = &ActiveInjectionPoints->entries[idx];
            uint64      generation;
 
            generation = pg_atomic_read_u64(&entry->generation);
            if (generation % 2 != 0)
                break;
        }
        pg_atomic_write_u32(&ActiveInjectionPoints->max_inuse, idx + 1);
    }
    LWLockRelease(InjectionPointLock);
 
    return found;
#else
    elog(ERROR, "Injection points are not supported by this build");
    return true;                /* silence compiler */
#endif
}
 
#ifdef USE_INJECTION_POINTS
/*
 * Common workhorse of InjectionPointRun() and InjectionPointLoad()
 *
 * Checks if an injection point exists in shared memory, and update
 * the local cache entry accordingly.
 */
static InjectionPointCacheEntry *
InjectionPointCacheRefresh(const char *name)
{
    uint32      max_inuse;
    int         namelen;
    InjectionPointEntry local_copy;
    InjectionPointCacheEntry *cached;
 
    /*
     * First read the number of in-use slots.  More entries can be added or
     * existing ones can be removed while we're reading them.  If the entry
     * we're looking for is concurrently added or removed, we might or might
     * not see it.  That's OK.
     */
    max_inuse = pg_atomic_read_u32(&ActiveInjectionPoints->max_inuse);
    if (max_inuse == 0)
    {
        if (InjectionPointCache)
        {
            hash_destroy(InjectionPointCache);
            InjectionPointCache = NULL;
        }
        return NULL;
    }
 
    /*
     * If we have this entry in the local cache already, check if the cached
     * entry is still valid.
     */
    cached = injection_point_cache_get(name);
    if (cached)
    {
        int         idx = cached->slot_idx;
        InjectionPointEntry *entry = &ActiveInjectionPoints->entries[idx];
 
        if (pg_atomic_read_u64(&entry->generation) == cached->generation)
        {
            /* still good */
            return cached;
        }
        injection_point_cache_remove(name);
        cached = NULL;
    }
 
    /*
     * Search the shared memory array.
     *
     * It's possible that the entry we're looking for is concurrently detached
     * or attached.  Or detached *and* re-attached, to the same slot or a
     * different slot.  Detach and re-attach is not an atomic operation, so
     * it's OK for us to return the old value, NULL, or the new value in such
     * cases.
     */
    namelen = strlen(name);
    for (int idx = 0; idx < max_inuse; idx++)
    {
        InjectionPointEntry *entry = &ActiveInjectionPoints->entries[idx];
        uint64      generation;
 
        /*
         * Read the generation number so that we can detect concurrent
         * modifications.  The read barrier ensures that the generation number
         * is loaded before any of the other fields.
         */
        generation = pg_atomic_read_u64(&entry->generation);
        if (generation % 2 == 0)
            continue;           /* empty slot */
        pg_read_barrier();
 
        /* Is this the injection point we're looking for? */
        if (memcmp(entry->name, name, namelen + 1) != 0)
            continue;
 
        /*
         * The entry can change at any time, if the injection point is
         * concurrently detached.  Copy it to local memory, and re-check the
         * generation.  If the generation hasn't changed, we know our local
         * copy is coherent.
         */
        memcpy(&local_copy, entry, sizeof(InjectionPointEntry));
 
        pg_read_barrier();
        if (pg_atomic_read_u64(&entry->generation) != generation)
        {
            /*
             * The entry was concurrently detached.
             *
             * Continue the search, because if the generation number changed,
             * we cannot trust the result of the name comparison we did above.
             * It's theoretically possible that it falsely matched a mixed-up
             * state of the old and new name, if the slot was recycled with a
             * different name.
             */
            continue;
        }
 
        /* Success! Load it into the cache and return it */
        return injection_point_cache_load(&local_copy, idx, generation);
    }
    return NULL;
}
#endif
 
/*
 * Load an injection point into the local cache.
 *
 * This is useful to be able to load an injection point before running it,
 * especially if the injection point is called in a code path where memory
 * allocations cannot happen, like critical sections.
 */
void
InjectionPointLoad(const char *name)
{
#ifdef USE_INJECTION_POINTS
    InjectionPointCacheRefresh(name);
#else
    elog(ERROR, "Injection points are not supported by this build");
#endif
}
 
/*
 * Execute an injection point, if defined.
 */
void
InjectionPointRun(const char *name)
{
#ifdef USE_INJECTION_POINTS
    InjectionPointCacheEntry *cache_entry;
 
    cache_entry = InjectionPointCacheRefresh(name);
    if (cache_entry)
        cache_entry->callback(name, cache_entry->private_data);
#else
    elog(ERROR, "Injection points are not supported by this build");
#endif
}
 
/*
 * Execute an injection point directly from the cache, if defined.
 */
void
InjectionPointCached(const char *name)
{
#ifdef USE_INJECTION_POINTS
    InjectionPointCacheEntry *cache_entry;
 
    cache_entry = injection_point_cache_get(name);
    if (cache_entry)
        cache_entry->callback(name, cache_entry->private_data);
#else
    elog(ERROR, "Injection points are not supported by this build");
#endif
}
 
/*
 * Test if an injection point is defined.
 */
bool
IsInjectionPointAttached(const char *name)
{
#ifdef USE_INJECTION_POINTS
    return InjectionPointCacheRefresh(name) != NULL;
#else
    elog(ERROR, "Injection points are not supported by this build");
    return false;               /* silence compiler */
#endif
}