/*-------------------------------------------------------------------------

 *

 * nodeGather.c

 *    Support routines for scanning a plan via multiple workers.

 *

 * Portions Copyright (c) 1996-2025, PostgreSQL Global Development Group

 * Portions Copyright (c) 1994, Regents of the University of California

 *

 * A Gather executor launches parallel workers to run multiple copies of a

 * plan.  It can also run the plan itself, if the workers are not available

 * or have not started up yet.  It then merges all of the results it produces

 * and the results from the workers into a single output stream.  Therefore,

 * it will normally be used with a plan where running multiple copies of the

 * same plan does not produce duplicate output, such as parallel-aware

 * SeqScan.

 *

 * Alternatively, a Gather node can be configured to use just one worker

 * and the single-copy flag can be set.  In this case, the Gather node will

 * run the plan in one worker and will not execute the plan itself.  In

 * this case, it simply returns whatever tuples were returned by the worker.

 * If a worker cannot be obtained, then it will run the plan itself and

 * return the results.  Therefore, a plan used with a single-copy Gather

 * node need not be parallel-aware.

 *

 * IDENTIFICATION

 *    src/backend/executor/nodeGather.c

 *

 *-------------------------------------------------------------------------

 */

#include "postgres.h"

#include "executor/execParallel.h"

#include "executor/executor.h"

#include "executor/nodeGather.h"

#include "executor/tqueue.h"

#include "miscadmin.h"

#include "optimizer/optimizer.h"

#include "utils/wait_event.h"

static TupleTableSlot *ExecGather(PlanState *pstate);

static TupleTableSlot *gather_getnext(GatherState *gatherstate);

static MinimalTuple gather_readnext(GatherState *gatherstate);

static void ExecShutdownGatherWorkers(GatherState *node);

/* ----------------------------------------------------------------

 *    ExecInitGather

 * ----------------------------------------------------------------

 */

GatherState *

ExecInitGather(Gather *node, EState *estate, int eflags)

{

  GatherState *gatherstate;

  Plan     *outerNode;

  TupleDesc tupDesc;

  /* Gather node doesn't have innerPlan node. */

  Assert(innerPlan(node) == NULL);

  /*

   * create state structure

   */

  gatherstate = makeNode(GatherState);

  gatherstate->ps.plan = (Plan *) node;

  gatherstate->ps.state = estate;

  gatherstate->ps.ExecProcNode = ExecGather;

  gatherstate->initialized = false;

  gatherstate->need_to_scan_locally =

    !node->single_copy && parallel_leader_participation;

  gatherstate->tuples_needed = -1;

  /*

   * Miscellaneous initialization

   *

   * create expression context for node

   */

  ExecAssignExprContext(estate, &gatherstate->ps);

  /*

   * now initialize outer plan

   */

  outerNode = outerPlan(node);

  outerPlanState(gatherstate) = ExecInitNode(outerNode, estate, eflags);

  tupDesc = ExecGetResultType(outerPlanState(gatherstate));

  /*

   * Leader may access ExecProcNode result directly (if

   * need_to_scan_locally), or from workers via tuple queue.  So we can't

   * trivially rely on the slot type being fixed for expressions evaluated

   * within this node.

   */

  gatherstate->ps.outeropsset = true;

  gatherstate->ps.outeropsfixed = false;

  /*

   * Initialize result type and projection.

   */

  ExecInitResultTypeTL(&gatherstate->ps);

  ExecConditionalAssignProjectionInfo(&gatherstate->ps, tupDesc, OUTER_VAR);

  /*

   * Without projections result slot type is not trivially known, see

   * comment above.

   */

  if (gatherstate->ps.ps_ProjInfo == NULL)

  {

    gatherstate->ps.resultopsset = true;

    gatherstate->ps.resultopsfixed = false;

  }

  /*

   * Initialize funnel slot to same tuple descriptor as outer plan.

   */

  gatherstate->funnel_slot = ExecInitExtraTupleSlot(estate, tupDesc,

                            &TTSOpsMinimalTuple);

  /*

   * Gather doesn't support checking a qual (it's always more efficient to

   * do it in the child node).

   */

  Assert(!node->plan.qual);

  return gatherstate;

}

/* ----------------------------------------------------------------

 *    ExecGather(node)

 *

 *    Scans the relation via multiple workers and returns

 *    the next qualifying tuple.

 * ----------------------------------------------------------------

 */

static TupleTableSlot *

ExecGather(PlanState *pstate)

{

  GatherState *node = castNode(GatherState, pstate);

  TupleTableSlot *slot;

  ExprContext *econtext;

  CHECK_FOR_INTERRUPTS();

  /*

   * Initialize the parallel context and workers on first execution. We do

   * this on first execution rather than during node initialization, as it

   * needs to allocate a large dynamic segment, so it is better to do it

   * only if it is really needed.

   */

  if (!node->initialized)

  {

    EState     *estate = node->ps.state;

    Gather     *gather = (Gather *) node->ps.plan;

    /*

     * Sometimes we might have to run without parallelism; but if parallel

     * mode is active then we can try to fire up some workers.

     */

    if (gather->num_workers > 0 && estate->es_use_parallel_mode)

    {

      ParallelContext *pcxt;

      /* Initialize, or re-initialize, shared state needed by workers. */

      if (!node->pei)

        node->pei = ExecInitParallelPlan(outerPlanState(node),

                         estate,

                         gather->initParam,

                         gather->num_workers,

                         node->tuples_needed);

      else

        ExecParallelReinitialize(outerPlanState(node),

                     node->pei,

                     gather->initParam);

      /*

       * Register backend workers. We might not get as many as we

       * requested, or indeed any at all.

       */

      pcxt = node->pei->pcxt;

      LaunchParallelWorkers(pcxt);

      /* We save # workers launched for the benefit of EXPLAIN */

      node->nworkers_launched = pcxt->nworkers_launched;

      /*

       * Count number of workers originally wanted and actually

       * launched.

       */

      estate->es_parallel_workers_to_launch += pcxt->nworkers_to_launch;

      estate->es_parallel_workers_launched += pcxt->nworkers_launched;

      /* Set up tuple queue readers to read the results. */

      if (pcxt->nworkers_launched > 0)

      {

        ExecParallelCreateReaders(node->pei);

        /* Make a working array showing the active readers */

        node->nreaders = pcxt->nworkers_launched;

        node->reader = (TupleQueueReader **)

          palloc(node->nreaders * sizeof(TupleQueueReader *));

        memcpy(node->reader, node->pei->reader,

             node->nreaders * sizeof(TupleQueueReader *));

      }

      else

      {

        /* No workers?  Then never mind. */

        node->nreaders = 0;

        node->reader = NULL;

      }

      node->nextreader = 0;

    }

    /* Run plan locally if no workers or enabled and not single-copy. */

    node->need_to_scan_locally = (node->nreaders == 0)

      || (!gather->single_copy && parallel_leader_participation);

    node->initialized = true;

  }

  /*

   * Reset per-tuple memory context to free any expression evaluation

   * storage allocated in the previous tuple cycle.

   */

  econtext = node->ps.ps_ExprContext;

  ResetExprContext(econtext);

  /*

   * Get next tuple, either from one of our workers, or by running the plan

   * ourselves.

   */

  slot = gather_getnext(node);

  if (TupIsNull(slot))

    return NULL;

  /* If no projection is required, we're done. */

  if (node->ps.ps_ProjInfo == NULL)

    return slot;

  /*

   * Form the result tuple using ExecProject(), and return it.

   */

  econtext->ecxt_outertuple = slot;

  return ExecProject(node->ps.ps_ProjInfo);

}

/* ----------------------------------------------------------------

 *    ExecEndGather

 *

 *    frees any storage allocated through C routines.

 * ----------------------------------------------------------------

 */

void

ExecEndGather(GatherState *node)

{

  ExecEndNode(outerPlanState(node));  /* let children clean up first */

  ExecShutdownGather(node);

}

/*

 * Read the next tuple.  We might fetch a tuple from one of the tuple queues

 * using gather_readnext, or if no tuple queue contains a tuple and the

 * single_copy flag is not set, we might generate one locally instead.

 */

static TupleTableSlot *

gather_getnext(GatherState *gatherstate)

{

  PlanState  *outerPlan = outerPlanState(gatherstate);

  TupleTableSlot *outerTupleSlot;

  TupleTableSlot *fslot = gatherstate->funnel_slot;

  MinimalTuple tup;

  while (gatherstate->nreaders > 0 || gatherstate->need_to_scan_locally)

  {

    CHECK_FOR_INTERRUPTS();

    if (gatherstate->nreaders > 0)

    {

      tup = gather_readnext(gatherstate);

      if (HeapTupleIsValid(tup))

      {

        ExecStoreMinimalTuple(tup,  /* tuple to store */

                    fslot,  /* slot to store the tuple */

                    false); /* don't pfree tuple  */

        return fslot;

      }

    }

    if (gatherstate->need_to_scan_locally)

    {

      EState     *estate = gatherstate->ps.state;

      /* Install our DSA area while executing the plan. */

      estate->es_query_dsa =

        gatherstate->pei ? gatherstate->pei->area : NULL;

      outerTupleSlot = ExecProcNode(outerPlan);

      estate->es_query_dsa = NULL;

      if (!TupIsNull(outerTupleSlot))

        return outerTupleSlot;

      gatherstate->need_to_scan_locally = false;

    }

  }

  return ExecClearTuple(fslot);

}

/*

 * Attempt to read a tuple from one of our parallel workers.

 */

static MinimalTuple

gather_readnext(GatherState *gatherstate)

{

  int     nvisited = 0;

  for (;;)

  {

    TupleQueueReader *reader;

    MinimalTuple tup;

    bool    readerdone;

    /* Check for async events, particularly messages from workers. */

    CHECK_FOR_INTERRUPTS();

    /*

     * Attempt to read a tuple, but don't block if none is available.

     *

     * Note that TupleQueueReaderNext will just return NULL for a worker

     * which fails to initialize.  We'll treat that worker as having

     * produced no tuples; WaitForParallelWorkersToFinish will error out

     * when we get there.

     */

    Assert(gatherstate->nextreader < gatherstate->nreaders);

    reader = gatherstate->reader[gatherstate->nextreader];

    tup = TupleQueueReaderNext(reader, true, &readerdone);

    /*

     * If this reader is done, remove it from our working array of active

     * readers.  If all readers are done, we're outta here.

     */

    if (readerdone)

    {

      Assert(!tup);

      --gatherstate->nreaders;

      if (gatherstate->nreaders == 0)

      {

        ExecShutdownGatherWorkers(gatherstate);

        return NULL;

      }

      memmove(&gatherstate->reader[gatherstate->nextreader],

          &gatherstate->reader[gatherstate->nextreader + 1],

          sizeof(TupleQueueReader *)

          * (gatherstate->nreaders - gatherstate->nextreader));

      if (gatherstate->nextreader >= gatherstate->nreaders)

        gatherstate->nextreader = 0;

      continue;

    }

    /* If we got a tuple, return it. */

    if (tup)

      return tup;

    /*

     * Advance nextreader pointer in round-robin fashion.  Note that we

     * only reach this code if we weren't able to get a tuple from the

     * current worker.  We used to advance the nextreader pointer after

     * every tuple, but it turns out to be much more efficient to keep

     * reading from the same queue until that would require blocking.

     */

    gatherstate->nextreader++;

    if (gatherstate->nextreader >= gatherstate->nreaders)

      gatherstate->nextreader = 0;

    /* Have we visited every (surviving) TupleQueueReader? */

    nvisited++;

    if (nvisited >= gatherstate->nreaders)

    {

      /*

       * If (still) running plan locally, return NULL so caller can

       * generate another tuple from the local copy of the plan.

       */

      if (gatherstate->need_to_scan_locally)

        return NULL;

      /* Nothing to do except wait for developments. */

      (void) WaitLatch(MyLatch, WL_LATCH_SET | WL_EXIT_ON_PM_DEATH, 0,

               WAIT_EVENT_EXECUTE_GATHER);

      ResetLatch(MyLatch);

      nvisited = 0;

    }

  }

}

/* ----------------------------------------------------------------

 *    ExecShutdownGatherWorkers

 *

 *    Stop all the parallel workers.

 * ----------------------------------------------------------------

 */

static void

ExecShutdownGatherWorkers(GatherState *node)

{

  if (node->pei != NULL)

    ExecParallelFinish(node->pei);

  /* Flush local copy of reader array */

  if (node->reader)

    pfree(node->reader);

  node->reader = NULL;

}

/* ----------------------------------------------------------------

 *    ExecShutdownGather

 *

 *    Destroy the setup for parallel workers including parallel context.

 * ----------------------------------------------------------------

 */

void

ExecShutdownGather(GatherState *node)

{

  ExecShutdownGatherWorkers(node);

  /* Now destroy the parallel context. */

  if (node->pei != NULL)

  {

    ExecParallelCleanup(node->pei);

    node->pei = NULL;

  }

}

/* ----------------------------------------------------------------

 *            Join Support

 * ----------------------------------------------------------------

 */

/* ----------------------------------------------------------------

 *    ExecReScanGather

 *

 *    Prepare to re-scan the result of a Gather.

 * ----------------------------------------------------------------

 */

void

ExecReScanGather(GatherState *node)

{

  Gather     *gather = (Gather *) node->ps.plan;

  PlanState  *outerPlan = outerPlanState(node);

  /* Make sure any existing workers are gracefully shut down */

  ExecShutdownGatherWorkers(node);

  /* Mark node so that shared state will be rebuilt at next call */

  node->initialized = false;

  /*

   * Set child node's chgParam to tell it that the next scan might deliver a

   * different set of rows within the leader process.  (The overall rowset

   * shouldn't change, but the leader process's subset might; hence nodes

   * between here and the parallel table scan node mustn't optimize on the

   * assumption of an unchanging rowset.)

   */

  if (gather->rescan_param >= 0)

    outerPlan->chgParam = bms_add_member(outerPlan->chgParam,

                       gather->rescan_param);

  /*

   * If chgParam of subnode is not null then plan will be re-scanned by

   * first ExecProcNode.  Note: because this does nothing if we have a

   * rescan_param, it's currently guaranteed that parallel-aware child nodes

   * will not see a ReScan call until after they get a ReInitializeDSM call.

   * That ordering might not be something to rely on, though.  A good rule

   * of thumb is that ReInitializeDSM should reset only shared state, ReScan

   * should reset only local state, and anything that depends on both of

   * those steps being finished must wait until the first ExecProcNode call.

   */

  if (outerPlan->chgParam == NULL)

    ExecReScan(outerPlan);

}