SSL 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);
 }
 if (!protocol_negotiation_supported(cluster))
  appendPQExpBufferStr(&conn_opts, " max_protocol_version=3.0");

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

 if (!slot->conn)
  pg_fatal("out of memory");

 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("%s() failed: %m", "select");

 /*
  * 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);
}