plpgsql: pure parser and reentrant scanner

[pgsql.git] / src / bin / pg_upgrade / task.c

/*
 * task.c
 *              framework for parallelizing pg_upgrade's once-in-each-database tasks
 *
 * This framework provides an efficient way of running the various
 * once-in-each-database tasks required by pg_upgrade.  Specifically, it
 * parallelizes these tasks by managing a set of slots that follow a simple
 * state machine and by using libpq's asynchronous APIs to establish the
 * connections and run the queries.  Callers simply need to create a callback
 * function and build/execute an UpgradeTask.  A simple example follows:
 *
 *              static void
 *              my_process_cb(DbInfo *dbinfo, PGresult *res, void *arg)
 *              {
 *                      for (int i = 0; i < PQntuples(res); i++)
 *                      {
 *                              ... process results ...
 *                      }
 *              }
 *
 *              void
 *              my_task(ClusterInfo *cluster)
 *              {
 *                      UpgradeTask *task = upgrade_task_create();
 *
 *                      upgrade_task_add_step(task,
 *                                                                "... query text ...",
 *                                                                my_process_cb,
 *                                                                true,         // let the task free the PGresult
 *                                                                NULL);        // "arg" pointer for callback
 *                      upgrade_task_run(task, cluster);
 *                      upgrade_task_free(task);
 *              }
 *
 * Note that multiple steps can be added to a given task.  When there are
 * multiple steps, the task will run all of the steps consecutively in the same
 * database connection before freeing the connection and moving on.  In other
 * words, it only ever initiates one connection to each database in the
 * cluster for a given run.
 *
 * Copyright (c) 2024-2025, PostgreSQL Global Development Group
 * src/bin/pg_upgrade/task.c
 */

#include "postgres_fe.h"

#include "common/connect.h"
#include "fe_utils/string_utils.h"
#include "pg_upgrade.h"

/*
 * dbs_complete stores the number of databases that we have completed
 * processing.  When this value equals the number of databases in the cluster,
 * the task is finished.
 */
static int      dbs_complete;

/*
 * dbs_processing stores the index of the next database in the cluster's array
 * of databases that will be picked up for processing.  It will always be
 * greater than or equal to dbs_complete.
 */
static int      dbs_processing;

/*
 * This struct stores the information for a single step of a task.  Note that
 * the query string is stored in the "queries" PQExpBuffer for the UpgradeTask.
 * All steps in a task are run in a single connection before moving on to the
 * next database (which requires a new connection).
 */
typedef struct UpgradeTaskStep
{
        UpgradeTaskProcessCB process_cb;        /* processes the results of the query */
        bool            free_result;    /* should we free the result? */
        void       *arg;                        /* pointer passed to process_cb */
} UpgradeTaskStep;

/*
 * This struct is a thin wrapper around an array of steps, i.e.,
 * UpgradeTaskStep, plus a PQExpBuffer for all the query strings.
 */
struct UpgradeTask
{
        UpgradeTaskStep *steps;
        int                     num_steps;
        PQExpBuffer queries;
};

/*
 * The different states for a parallel slot.
 */
typedef enum UpgradeTaskSlotState
{
        FREE,                                           /* slot available for use in a new database */
        CONNECTING,                                     /* waiting for connection to be established */
        RUNNING_QUERIES,                        /* running/processing queries in the task */
} UpgradeTaskSlotState;

/*
 * We maintain an array of user_opts.jobs slots to execute the task.
 */
typedef struct UpgradeTaskSlot
{
        UpgradeTaskSlotState state; /* state of the slot */
        int                     db_idx;                 /* index of the database assigned to slot */
        int                     step_idx;               /* index of the current step of task */
        PGconn     *conn;                       /* current connection managed by slot */
        bool            ready;                  /* slot is ready for processing */
        bool            select_mode;    /* select() mode: true->read, false->write */
        int                     sock;                   /* file descriptor for connection's socket */
} UpgradeTaskSlot;

/*
 * Initializes an UpgradeTask.
 */
UpgradeTask *
upgrade_task_create(void)
{
        UpgradeTask *task = pg_malloc0(sizeof(UpgradeTask));

        task->queries = createPQExpBuffer();

        /* All tasks must first set a secure search_path. */
        upgrade_task_add_step(task, ALWAYS_SECURE_SEARCH_PATH_SQL, NULL, true, NULL);

        return task;
}

/*
 * Frees all storage associated with an UpgradeTask.
 */
void
upgrade_task_free(UpgradeTask *task)
{
        destroyPQExpBuffer(task->queries);
        pg_free(task->steps);
        pg_free(task);
}

/*
 * Adds a step to an UpgradeTask.  The steps will be executed in each database
 * in the order in which they are added.
 *
 *      task: task object that must have been initialized via upgrade_task_create()
 *      query: the query text
 *      process_cb: function that processes the results of the query
 *      free_result: should we free the PGresult, or leave it to the caller?
 *      arg: pointer to task-specific data that is passed to each callback
 */
void
upgrade_task_add_step(UpgradeTask *task, const char *query,
                                          UpgradeTaskProcessCB process_cb, bool free_result,
                                          void *arg)
{
        UpgradeTaskStep *new_step;

        task->steps = pg_realloc(task->steps,
                                                         ++task->num_steps * sizeof(UpgradeTaskStep));

        new_step = &task->steps[task->num_steps - 1];
        new_step->process_cb = process_cb;
        new_step->free_result = free_result;
        new_step->arg = arg;

        appendPQExpBuffer(task->queries, "%s;", query);
}

/*
 * Build a connection string for the slot's current database and asynchronously
 * start a new connection, but do not wait for the connection to be
 * established.
 */
static void
start_conn(const ClusterInfo *cluster, UpgradeTaskSlot *slot)
{
        PQExpBufferData conn_opts;
        DbInfo     *dbinfo = &cluster->dbarr.dbs[slot->db_idx];

        /* Build connection string with proper quoting */
        initPQExpBuffer(&conn_opts);
        appendPQExpBufferStr(&conn_opts, "dbname=");
        appendConnStrVal(&conn_opts, dbinfo->db_name);
        appendPQExpBufferStr(&conn_opts, " user=");
        appendConnStrVal(&conn_opts, os_info.user);
        appendPQExpBuffer(&conn_opts, " port=%d", cluster->port);
        if (cluster->sockdir)
        {
                appendPQExpBufferStr(&conn_opts, " host=");
                appendConnStrVal(&conn_opts, cluster->sockdir);
        }

        slot->conn = PQconnectStart(conn_opts.data);

        if (!slot->conn)
                pg_fatal("failed to create connection with connection string: \"%s\"",
                                 conn_opts.data);

        termPQExpBuffer(&conn_opts);
}

/*
 * Run the process_cb callback function to process the result of a query, and
 * free the result if the caller indicated we should do so.
 */
static void
process_query_result(const ClusterInfo *cluster, UpgradeTaskSlot *slot,
                                         const UpgradeTask *task)
{
        UpgradeTaskStep *steps = &task->steps[slot->step_idx];
        UpgradeTaskProcessCB process_cb = steps->process_cb;
        DbInfo     *dbinfo = &cluster->dbarr.dbs[slot->db_idx];
        PGresult   *res = PQgetResult(slot->conn);

        if (PQstatus(slot->conn) == CONNECTION_BAD ||
                (PQresultStatus(res) != PGRES_TUPLES_OK &&
                 PQresultStatus(res) != PGRES_COMMAND_OK))
                pg_fatal("connection failure: %s", PQerrorMessage(slot->conn));

        /*
         * We assume that a NULL process_cb callback function means there's
         * nothing to process.  This is primarily intended for the initial step in
         * every task that sets a safe search_path.
         */
        if (process_cb)
                (*process_cb) (dbinfo, res, steps->arg);

        if (steps->free_result)
                PQclear(res);
}

/*
 * Advances the state machine for a given slot as necessary.
 */
static void
process_slot(const ClusterInfo *cluster, UpgradeTaskSlot *slot, const UpgradeTask *task)
{
        PostgresPollingStatusType status;

        if (!slot->ready)
                return;

        switch (slot->state)
        {
                case FREE:

                        /*
                         * If all of the databases in the cluster have been processed or
                         * are currently being processed by other slots, we are done.
                         */
                        if (dbs_processing >= cluster->dbarr.ndbs)
                                return;

                        /*
                         * Claim the next database in the cluster's array and initiate a
                         * new connection.
                         */
                        slot->db_idx = dbs_processing++;
                        slot->state = CONNECTING;
                        start_conn(cluster, slot);

                        return;

                case CONNECTING:

                        /* Check for connection failure. */
                        status = PQconnectPoll(slot->conn);
                        if (status == PGRES_POLLING_FAILED)
                                pg_fatal("connection failure: %s", PQerrorMessage(slot->conn));

                        /* Check whether the connection is still establishing. */
                        if (status != PGRES_POLLING_OK)
                        {
                                slot->select_mode = (status == PGRES_POLLING_READING);
                                return;
                        }

                        /*
                         * Move on to running/processing the queries in the task.
                         */
                        slot->state = RUNNING_QUERIES;
                        slot->select_mode = true;       /* wait until ready for reading */
                        if (!PQsendQuery(slot->conn, task->queries->data))
                                pg_fatal("connection failure: %s", PQerrorMessage(slot->conn));

                        return;

                case RUNNING_QUERIES:

                        /*
                         * Consume any available data and clear the read-ready indicator
                         * for the connection.
                         */
                        if (!PQconsumeInput(slot->conn))
                                pg_fatal("connection failure: %s", PQerrorMessage(slot->conn));

                        /*
                         * Process any results that are ready so that we can free up this
                         * slot for another database as soon as possible.
                         */
                        for (; slot->step_idx < task->num_steps; slot->step_idx++)
                        {
                                /* If no more results are available yet, move on. */
                                if (PQisBusy(slot->conn))
                                        return;

                                process_query_result(cluster, slot, task);
                        }

                        /*
                         * If we just finished processing the result of the last step in
                         * the task, free the slot.  We recursively call this function on
                         * the newly-freed slot so that we can start initiating the next
                         * connection immediately instead of waiting for the next loop
                         * through the slots.
                         */
                        dbs_complete++;
                        PQfinish(slot->conn);
                        memset(slot, 0, sizeof(UpgradeTaskSlot));
                        slot->ready = true;

                        process_slot(cluster, slot, task);

                        return;
        }
}

/*
 * Returns -1 on error, else the number of ready descriptors.
 */
static int
select_loop(int maxFd, fd_set *input, fd_set *output)
{
        fd_set          save_input = *input;
        fd_set          save_output = *output;

        if (maxFd == 0)
                return 0;

        for (;;)
        {
                int                     i;

                *input = save_input;
                *output = save_output;

                i = select(maxFd + 1, input, output, NULL, NULL);

#ifndef WIN32
                if (i < 0 && errno == EINTR)
                        continue;
#else
                if (i == SOCKET_ERROR && WSAGetLastError() == WSAEINTR)
                        continue;
#endif
                return i;
        }
}

/*
 * Wait on the slots to either finish connecting or to receive query results if
 * possible.  This avoids a tight loop in upgrade_task_run().
 */
static void
wait_on_slots(UpgradeTaskSlot *slots, int numslots)
{
        fd_set          input;
        fd_set          output;
        int                     maxFd = 0;

        FD_ZERO(&input);
        FD_ZERO(&output);

        for (int i = 0; i < numslots; i++)
        {
                /*
                 * We assume the previous call to process_slot() handled everything
                 * that was marked ready in the previous call to wait_on_slots(), if
                 * any.
                 */
                slots[i].ready = false;

                /*
                 * This function should only ever see free slots as we are finishing
                 * processing the last few databases, at which point we don't have any
                 * databases left for them to process.  We'll never use these slots
                 * again, so we can safely ignore them.
                 */
                if (slots[i].state == FREE)
                        continue;

                /*
                 * Add the socket to the set.
                 */
                slots[i].sock = PQsocket(slots[i].conn);
                if (slots[i].sock < 0)
                        pg_fatal("invalid socket");
                FD_SET(slots[i].sock, slots[i].select_mode ? &input : &output);
                maxFd = Max(maxFd, slots[i].sock);
        }

        /*
         * If we found socket(s) to wait on, wait.
         */
        if (select_loop(maxFd, &input, &output) == -1)
                pg_fatal("select() failed: %m");

        /*
         * Mark which sockets appear to be ready.
         */
        for (int i = 0; i < numslots; i++)
                slots[i].ready |= (FD_ISSET(slots[i].sock, &input) ||
                                                   FD_ISSET(slots[i].sock, &output));
}

/*
 * Runs all the steps of the task in every database in the cluster using
 * user_opts.jobs parallel slots.
 */
void
upgrade_task_run(const UpgradeTask *task, const ClusterInfo *cluster)
{
        int                     jobs = Max(1, user_opts.jobs);
        UpgradeTaskSlot *slots = pg_malloc0(sizeof(UpgradeTaskSlot) * jobs);

        dbs_complete = 0;
        dbs_processing = 0;

        /*
         * Process every slot the first time round.
         */
        for (int i = 0; i < jobs; i++)
                slots[i].ready = true;

        while (dbs_complete < cluster->dbarr.ndbs)
        {
                for (int i = 0; i < jobs; i++)
                        process_slot(cluster, &slots[i], task);

                wait_on_slots(slots, jobs);
        }

        pg_free(slots);
}