/*-------------------------------------------------------------------------
 *
 * sharedtuplestore.c
 *   Simple mechanism for sharing tuples between backends.
 *
 * This module contains a shared temporary tuple storage mechanism providing
 * a parallel-aware subset of the features of tuplestore.c.  Multiple backends
 * can write to a SharedTuplestore, and then multiple backends can later scan
 * the stored tuples.  Currently, the only scan type supported is a parallel
 * scan where each backend reads an arbitrary subset of the tuples that were
 * written.
 *
 * Portions Copyright (c) 1996-2025, PostgreSQL Global Development Group
 * Portions Copyright (c) 1994, Regents of the University of California
 *
 * IDENTIFICATION
 *   src/backend/utils/sort/sharedtuplestore.c
 *
 *-------------------------------------------------------------------------
 */

#include "postgres.h"

#include "access/htup.h"
#include "access/htup_details.h"
#include "storage/buffile.h"
#include "storage/lwlock.h"
#include "storage/sharedfileset.h"
#include "utils/sharedtuplestore.h"

/*
 * The size of chunks, in pages.  This is somewhat arbitrarily set to match
 * the size of HASH_CHUNK, so that Parallel Hash obtains new chunks of tuples
 * at approximately the same rate as it allocates new chunks of memory to
 * insert them into.
 */
#define STS_CHUNK_PAGES 4
#define STS_CHUNK_HEADER_SIZE offsetof(SharedTuplestoreChunk, data)
#define STS_CHUNK_DATA_SIZE (STS_CHUNK_PAGES * BLCKSZ - STS_CHUNK_HEADER_SIZE)

/* Chunk written to disk. */
typedef struct SharedTuplestoreChunk
{
 int   ntuples;  /* Number of tuples in this chunk. */
 int   overflow;  /* If overflow, how many including this one? */
 char  data[FLEXIBLE_ARRAY_MEMBER];
} SharedTuplestoreChunk;

/* Per-participant shared state. */
typedef struct SharedTuplestoreParticipant
{
 LWLock  lock;
 BlockNumber read_page;  /* Page number for next read. */
 BlockNumber npages;   /* Number of pages written. */
 bool  writing;  /* Used only for assertions. */
} SharedTuplestoreParticipant;

/* The control object that lives in shared memory. */
struct SharedTuplestore
{
 int   nparticipants; /* Number of participants that can write. */
 int   flags;   /* Flag bits from SHARED_TUPLESTORE_XXX */
 size_t  meta_data_size; /* Size of per-tuple header. */
 char  name[NAMEDATALEN]; /* A name for this tuplestore. */

 /* Followed by per-participant shared state. */
 SharedTuplestoreParticipant participants[FLEXIBLE_ARRAY_MEMBER];
};

/* Per-participant state that lives in backend-local memory. */
struct SharedTuplestoreAccessor
{
 int   participant; /* My participant number. */
 SharedTuplestore *sts;  /* The shared state. */
 SharedFileSet *fileset;  /* The SharedFileSet holding files. */
 MemoryContext context;  /* Memory context for buffers. */

 /* State for reading. */
 int   read_participant; /* The current participant to read from. */
 BufFile    *read_file;  /* The current file to read from. */
 int   read_ntuples_available; /* The number of tuples in chunk. */
 int   read_ntuples; /* How many tuples have we read from chunk? */
 size_t  read_bytes;  /* How many bytes have we read from chunk? */
 char    *read_buffer; /* A buffer for loading tuples. */
 size_t  read_buffer_size;
 BlockNumber read_next_page; /* Lowest block we'll consider reading. */

 /* State for writing. */
 SharedTuplestoreChunk *write_chunk; /* Buffer for writing. */
 BufFile    *write_file;  /* The current file to write to. */
 BlockNumber write_page;  /* The next page to write to. */
 char    *write_pointer; /* Current write pointer within chunk. */
 char    *write_end;  /* One past the end of the current chunk. */
};

static void sts_filename(char *name, SharedTuplestoreAccessor *accessor,
       int participant);

/*
 * Return the amount of shared memory required to hold SharedTuplestore for a
 * given number of participants.
 */
size_t
sts_estimate(int participants)
{
 return offsetof(SharedTuplestore, participants) +
  sizeof(SharedTuplestoreParticipant) * participants;
}

/*
 * Initialize a SharedTuplestore in existing shared memory.  There must be
 * space for sts_estimate(participants) bytes.  If flags includes the value
 * SHARED_TUPLESTORE_SINGLE_PASS, the files may in future be removed more
 * eagerly (but this isn't yet implemented).
 *
 * Tuples that are stored may optionally carry a piece of fixed sized
 * meta-data which will be retrieved along with the tuple.  This is useful for
 * the hash values used in multi-batch hash joins, but could have other
 * applications.
 *
 * The caller must supply a SharedFileSet, which is essentially a directory
 * that will be cleaned up automatically, and a name which must be unique
 * across all SharedTuplestores created in the same SharedFileSet.
 */
SharedTuplestoreAccessor *
sts_initialize(SharedTuplestore *sts, int participants,
      int my_participant_number,
      size_t meta_data_size,
      int flags,
      SharedFileSet *fileset,
      const char *name)
{
 SharedTuplestoreAccessor *accessor;
 int   i;

 Assert(my_participant_number < participants);

 sts->nparticipants = participants;
 sts->meta_data_size = meta_data_size;
 sts->flags = flags;

 if (strlen(name) > sizeof(sts->name) - 1)
  elog(ERROR, "SharedTuplestore name too long");
 strcpy(sts->name, name);

 /*
  * Limit meta-data so it + tuple size always fits into a single chunk.
  * sts_puttuple() and sts_read_tuple() could be made to support scenarios
  * where that's not the case, but it's not currently required. If so,
  * meta-data size probably should be made variable, too.
 */
 if (meta_data_size + sizeof(uint32) >= STS_CHUNK_DATA_SIZE)
  elog(ERROR, "meta-data too long");

 for (i = 0; i < participants; ++i)
 {
  LWLockInitialize(&sts->participants[i].lock,
       LWTRANCHE_SHARED_TUPLESTORE);
  sts->participants[i].read_page = 0;
  sts->participants[i].npages = 0;
  sts->participants[i].writing = false;
 }

 accessor = palloc0(sizeof(SharedTuplestoreAccessor));
 accessor->participant = my_participant_number;
 accessor->sts = sts;
 accessor->fileset = fileset;
 accessor->context = CurrentMemoryContext;

 return accessor;
}

/*
 * Attach to a SharedTuplestore that has been initialized by another backend,
 * so that this backend can read and write tuples.
 */
SharedTuplestoreAccessor *
sts_attach(SharedTuplestore *sts,
     int my_participant_number,
     SharedFileSet *fileset)
{
 SharedTuplestoreAccessor *accessor;

 Assert(my_participant_number < sts->nparticipants);

 accessor = palloc0(sizeof(SharedTuplestoreAccessor));
 accessor->participant = my_participant_number;
 accessor->sts = sts;
 accessor->fileset = fileset;
 accessor->context = CurrentMemoryContext;

 return accessor;
}

static void
sts_flush_chunk(SharedTuplestoreAccessor *accessor)
{
 size_t  size;

 size = STS_CHUNK_PAGES * BLCKSZ;
 BufFileWrite(accessor->write_file, accessor->write_chunk, size);
 memset(accessor->write_chunk, 0, size);
 accessor->write_pointer = &accessor->write_chunk->data[0];
 accessor->sts->participants[accessor->participant].npages +=
  STS_CHUNK_PAGES;
}

/*
 * Finish writing tuples.  This must be called by all backends that have
 * written data before any backend begins reading it.
 */
void
sts_end_write(SharedTuplestoreAccessor *accessor)
{
 if (accessor->write_file != NULL)
 {
  sts_flush_chunk(accessor);
  BufFileClose(accessor->write_file);
  pfree(accessor->write_chunk);
  accessor->write_chunk = NULL;
  accessor->write_file = NULL;
  accessor->sts->participants[accessor->participant].writing = false;
 }
}

/*
 * Prepare to rescan.  Only one participant must call this.  After it returns,
 * all participants may call sts_begin_parallel_scan() and then loop over
 * sts_parallel_scan_next().  This function must not be called concurrently
 * with a scan, and synchronization to avoid that is the caller's
 * responsibility.
 */
void
sts_reinitialize(SharedTuplestoreAccessor *accessor)
{
 int   i;

 /*
  * Reset the shared read head for all participants' files.  Also set the
  * initial chunk size to the minimum (any increases from that size will be
  * recorded in chunk_expansion_log).
 */
 for (i = 0; i < accessor->sts->nparticipants; ++i)
 {
  accessor->sts->participants[i].read_page = 0;
 }
}

/*
 * Begin scanning the contents in parallel.
 */
void
sts_begin_parallel_scan(SharedTuplestoreAccessor *accessor)
{
 int   i PG_USED_FOR_ASSERTS_ONLY;

 /* End any existing scan that was in progress. */
 sts_end_parallel_scan(accessor);

 /*
  * Any backend that might have written into this shared tuplestore must
  * have called sts_end_write(), so that all buffers are flushed and the
  * files have stopped growing.
 */
 for (i = 0; i < accessor->sts->nparticipants; ++i)
  Assert(!accessor->sts->participants[i].writing);

 /*
  * We will start out reading the file that THIS backend wrote.  There may
  * be some caching locality advantage to that.
 */
 accessor->read_participant = accessor->participant;
 accessor->read_file = NULL;
 accessor->read_next_page = 0;
}

/*
 * Finish a parallel scan, freeing associated backend-local resources.
 */
void
sts_end_parallel_scan(SharedTuplestoreAccessor *accessor)
{
 /*
  * Here we could delete all files if SHARED_TUPLESTORE_SINGLE_PASS, but
  * we'd probably need a reference count of current parallel scanners so we
  * could safely do it only when the reference count reaches zero.
 */
 if (accessor->read_file != NULL)
 {
  BufFileClose(accessor->read_file);
  accessor->read_file = NULL;
 }
}

/*
 * Write a tuple.  If a meta-data size was provided to sts_initialize, then a
 * pointer to meta data of that size must be provided.
 */
void
sts_puttuple(SharedTuplestoreAccessor *accessor, void *meta_data,
    MinimalTuple tuple)
{
 size_t  size;

 /* Do we have our own file yet? */
 if (accessor->write_file == NULL)
 {
  SharedTuplestoreParticipant *participant;
  char  name[MAXPGPATH];
  MemoryContext oldcxt;

  /* Create one.  Only this backend will write into it. */
  sts_filename(name, accessor, accessor->participant);

  oldcxt = MemoryContextSwitchTo(accessor->context);
  accessor->write_file =
   BufFileCreateFileSet(&accessor->fileset->fs, name);
  MemoryContextSwitchTo(oldcxt);

  /* Set up the shared state for this backend's file. */
  participant = &accessor->sts->participants[accessor->participant];
  participant->writing = true; /* for assertions only */
 }

 /* Do we have space? */
 size = accessor->sts->meta_data_size + tuple->t_len;
 if (accessor->write_pointer == NULL ||
  accessor->write_pointer + size > accessor->write_end)
 {
  if (accessor->write_chunk == NULL)
  {
   /* First time through.  Allocate chunk. */
   accessor->write_chunk = (SharedTuplestoreChunk *)
    MemoryContextAllocZero(accessor->context,
            STS_CHUNK_PAGES * BLCKSZ);
   accessor->write_chunk->ntuples = 0;
   accessor->write_pointer = &accessor->write_chunk->data[0];
   accessor->write_end = (char *)
    accessor->write_chunk + STS_CHUNK_PAGES * BLCKSZ;
  }
  else
  {
   /* See if flushing helps. */
   sts_flush_chunk(accessor);
  }

  /* It may still not be enough in the case of a gigantic tuple. */
  if (accessor->write_pointer + size > accessor->write_end)
  {
   size_t  written;

   /*
    * We'll write the beginning of the oversized tuple, and then
    * write the rest in some number of 'overflow' chunks.
    *
    * sts_initialize() verifies that the size of the tuple +
    * meta-data always fits into a chunk. Because the chunk has been
    * flushed above, we can be sure to have all of a chunk's usable
    * space available.
 */
   Assert(accessor->write_pointer + accessor->sts->meta_data_size +
       sizeof(uint32) < accessor->write_end);

   /* Write the meta-data as one chunk. */
   if (accessor->sts->meta_data_size > 0)
    memcpy(accessor->write_pointer, meta_data,
        accessor->sts->meta_data_size);

   /*
    * Write as much of the tuple as we can fit. This includes the
    * tuple's size at the start.
 */
   written = accessor->write_end - accessor->write_pointer -
    accessor->sts->meta_data_size;
   memcpy(accessor->write_pointer + accessor->sts->meta_data_size,
       tuple, written);
   ++accessor->write_chunk->ntuples;
   size -= accessor->sts->meta_data_size;
   size -= written;
   /* Now write as many overflow chunks as we need for the rest. */
   while (size > 0)
   {
    size_t  written_this_chunk;

    sts_flush_chunk(accessor);

    /*
     * How many overflow chunks to go?  This will allow readers to
     * skip all of them at once instead of reading each one.
 */
    accessor->write_chunk->overflow = (size + STS_CHUNK_DATA_SIZE - 1) /
     STS_CHUNK_DATA_SIZE;
    written_this_chunk =
     Min(accessor->write_end - accessor->write_pointer, size);
    memcpy(accessor->write_pointer, (char *) tuple + written,
        written_this_chunk);
    accessor->write_pointer += written_this_chunk;
    size -= written_this_chunk;
    written += written_this_chunk;
   }
   return;
  }
 }

 /* Copy meta-data and tuple into buffer. */
 if (accessor->sts->meta_data_size > 0)
  memcpy(accessor->write_pointer, meta_data,
      accessor->sts->meta_data_size);
 memcpy(accessor->write_pointer + accessor->sts->meta_data_size, tuple,
     tuple->t_len);
 accessor->write_pointer += size;
 ++accessor->write_chunk->ntuples;
}

static MinimalTuple
sts_read_tuple(SharedTuplestoreAccessor *accessor, void *meta_data)
{
 MinimalTuple tuple;
 uint32  size;
 size_t  remaining_size;
 size_t  this_chunk_size;
 char    *destination;

 /*
  * We'll keep track of bytes read from this chunk so that we can detect an
  * overflowing tuple and switch to reading overflow pages.
 */
 if (accessor->sts->meta_data_size > 0)
 {
  BufFileReadExact(accessor->read_file, meta_data, accessor->sts->meta_data_size);
  accessor->read_bytes += accessor->sts->meta_data_size;
 }
 BufFileReadExact(accessor->read_file, &size, sizeof(size));
 accessor->read_bytes += sizeof(size);
 if (size > accessor->read_buffer_size)
 {
  size_t  new_read_buffer_size;

  if (accessor->read_buffer != NULL)
   pfree(accessor->read_buffer);
  new_read_buffer_size = Max(size, accessor->read_buffer_size * 2);
  accessor->read_buffer =
   MemoryContextAlloc(accessor->context, new_read_buffer_size);
  accessor->read_buffer_size = new_read_buffer_size;
 }
 remaining_size = size - sizeof(uint32);
 this_chunk_size = Min(remaining_size,
        BLCKSZ * STS_CHUNK_PAGES - accessor->read_bytes);
 destination = accessor->read_buffer + sizeof(uint32);
 BufFileReadExact(accessor->read_file, destination, this_chunk_size);
 accessor->read_bytes += this_chunk_size;
 remaining_size -= this_chunk_size;
 destination += this_chunk_size;
 ++accessor->read_ntuples;

 /* Check if we need to read any overflow chunks. */
 while (remaining_size > 0)
 {
  /* We are now positioned at the start of an overflow chunk. */
  SharedTuplestoreChunk chunk_header;

  BufFileReadExact(accessor->read_file, &chunk_header, STS_CHUNK_HEADER_SIZE);
  accessor->read_bytes = STS_CHUNK_HEADER_SIZE;
  if (chunk_header.overflow == 0)
   ereport(ERROR,
     (errcode_for_file_access(),
      errmsg("unexpected chunk in shared tuplestore temporary file"),
      errdetail_internal("Expected overflow chunk.")));
  accessor->read_next_page += STS_CHUNK_PAGES;
  this_chunk_size = Min(remaining_size,
         BLCKSZ * STS_CHUNK_PAGES -
         STS_CHUNK_HEADER_SIZE);
  BufFileReadExact(accessor->read_file, destination, this_chunk_size);
  accessor->read_bytes += this_chunk_size;
  remaining_size -= this_chunk_size;
  destination += this_chunk_size;

  /*
   * These will be used to count regular tuples following the oversized
   * tuple that spilled into this overflow chunk.
 */
  accessor->read_ntuples = 0;
  accessor->read_ntuples_available = chunk_header.ntuples;
 }

 tuple = (MinimalTuple) accessor->read_buffer;
 tuple->t_len = size;

 return tuple;
}

/*
 * Get the next tuple in the current parallel scan.
 */
MinimalTuple
sts_parallel_scan_next(SharedTuplestoreAccessor *accessor, void *meta_data)
{
 SharedTuplestoreParticipant *p;
 BlockNumber read_page;
 bool  eof;

 for (;;)
 {
  /* Can we read more tuples from the current chunk? */
  if (accessor->read_ntuples < accessor->read_ntuples_available)
   return sts_read_tuple(accessor, meta_data);

  /* Find the location of a new chunk to read. */
  p = &accessor->sts->participants[accessor->read_participant];

  LWLockAcquire(&p->lock, LW_EXCLUSIVE);
  /* We can skip directly past overflow pages we know about. */
  if (p->read_page < accessor->read_next_page)
   p->read_page = accessor->read_next_page;
  eof = p->read_page >= p->npages;
  if (!eof)
  {
   /* Claim the next chunk. */
   read_page = p->read_page;
   /* Advance the read head for the next reader. */
   p->read_page += STS_CHUNK_PAGES;
   accessor->read_next_page = p->read_page;
  }
  LWLockRelease(&p->lock);

  if (!eof)
  {
   SharedTuplestoreChunk chunk_header;

   /* Make sure we have the file open. */
   if (accessor->read_file == NULL)
   {
    char  name[MAXPGPATH];
    MemoryContext oldcxt;

    sts_filename(name, accessor, accessor->read_participant);

    oldcxt = MemoryContextSwitchTo(accessor->context);
    accessor->read_file =
     BufFileOpenFileSet(&accessor->fileset->fs, name, O_RDONLY,
            false);
    MemoryContextSwitchTo(oldcxt);
   }

   /* Seek and load the chunk header. */
   if (BufFileSeekBlock(accessor->read_file, read_page) != 0)
    ereport(ERROR,
      (errcode_for_file_access(),
       errmsg("could not seek to block %u in shared tuplestore temporary file",
        read_page)));
   BufFileReadExact(accessor->read_file, &chunk_header, STS_CHUNK_HEADER_SIZE);

   /*
    * If this is an overflow chunk, we skip it and any following
    * overflow chunks all at once.
 */
   if (chunk_header.overflow > 0)
   {
    accessor->read_next_page = read_page +
     chunk_header.overflow * STS_CHUNK_PAGES;
    continue;
   }

   accessor->read_ntuples = 0;
   accessor->read_ntuples_available = chunk_header.ntuples;
   accessor->read_bytes = STS_CHUNK_HEADER_SIZE;

   /* Go around again, so we can get a tuple from this chunk. */
  }
  else
  {
   if (accessor->read_file != NULL)
   {
    BufFileClose(accessor->read_file);
    accessor->read_file = NULL;
   }

   /*
    * Try the next participant's file.  If we've gone full circle,
    * we're done.
 */
   accessor->read_participant = (accessor->read_participant + 1) %
    accessor->sts->nparticipants;
   if (accessor->read_participant == accessor->participant)
    break;
   accessor->read_next_page = 0;

   /* Go around again, so we can get a chunk from this file. */
  }
 }

 return NULL;
}

/*
 * Create the name used for the BufFile that a given participant will write.
 */
static void
sts_filename(char *name, SharedTuplestoreAccessor *accessor, int participant)
{
 snprintf(name, MAXPGPATH, "%s.p%d", accessor->sts->name, participant);
}