LCOV - code coverage report
Current view: top level - contrib/pg_stash_advice - stashpersist.c (source / functions) Coverage Total Hit
Test: PostgreSQL 20devel Lines: 77.3 % 269 208
Test Date: 2026-07-03 19:57:34 Functions: 91.7 % 12 11
Legend: Lines:     hit not hit
Branches: + taken - not taken # not executed
Branches: 47.4 % 171 81

             Branch data     Line data    Source code
       1                 :             : /*-------------------------------------------------------------------------
       2                 :             :  *
       3                 :             :  * stashpersist.c
       4                 :             :  *    Persistence support for pg_stash_advice.
       5                 :             :  *
       6                 :             :  * Copyright (c) 2016-2026, PostgreSQL Global Development Group
       7                 :             :  *
       8                 :             :  *    contrib/pg_stash_advice/stashpersist.c
       9                 :             :  *
      10                 :             :  *-------------------------------------------------------------------------
      11                 :             :  */
      12                 :             : #include "postgres.h"
      13                 :             : 
      14                 :             : #include <sys/stat.h>
      15                 :             : 
      16                 :             : #include "common/hashfn.h"
      17                 :             : #include "miscadmin.h"
      18                 :             : #include "pg_stash_advice.h"
      19                 :             : #include "postmaster/bgworker.h"
      20                 :             : #include "postmaster/interrupt.h"
      21                 :             : #include "storage/fd.h"
      22                 :             : #include "storage/ipc.h"
      23                 :             : #include "storage/latch.h"
      24                 :             : #include "storage/proc.h"
      25                 :             : #include "storage/procsignal.h"
      26                 :             : #include "utils/backend_status.h"
      27                 :             : #include "utils/guc.h"
      28                 :             : #include "utils/memutils.h"
      29                 :             : #include "utils/timestamp.h"
      30                 :             : 
      31                 :             : typedef struct pgsa_writer_context
      32                 :             : {
      33                 :             :     char        pathname[MAXPGPATH];
      34                 :             :     FILE       *file;
      35                 :             :     pgsa_stash_name_table_hash *nhash;
      36                 :             :     StringInfoData buf;
      37                 :             :     int         entries_written;
      38                 :             : } pgsa_writer_context;
      39                 :             : 
      40                 :             : /*
      41                 :             :  * A parsed entry line, with pointers into the slurp buffer.
      42                 :             :  */
      43                 :             : typedef struct pgsa_saved_entry
      44                 :             : {
      45                 :             :     char       *stash_name;
      46                 :             :     int64       queryId;
      47                 :             :     char       *advice_string;
      48                 :             : } pgsa_saved_entry;
      49                 :             : 
      50                 :             : /*
      51                 :             :  * simplehash for detecting duplicate stash names during parsing.
      52                 :             :  * Keyed by stash name (char *), pointing into the slurp buffer.
      53                 :             :  */
      54                 :             : typedef struct pgsa_saved_stash
      55                 :             : {
      56                 :             :     uint32      status;
      57                 :             :     char       *name;
      58                 :             : } pgsa_saved_stash;
      59                 :             : 
      60                 :             : #define SH_PREFIX pgsa_saved_stash_table
      61                 :             : #define SH_ELEMENT_TYPE pgsa_saved_stash
      62                 :             : #define SH_KEY_TYPE char *
      63                 :             : #define SH_KEY name
      64                 :             : #define SH_HASH_KEY(tb, key) hash_bytes((const unsigned char *) (key), strlen(key))
      65                 :             : #define SH_EQUAL(tb, a, b) (strcmp(a, b) == 0)
      66                 :             : #define SH_SCOPE static inline
      67                 :             : #define SH_DEFINE
      68                 :             : #define SH_DECLARE
      69                 :             : #include "lib/simplehash.h"
      70                 :             : 
      71                 :             : extern PGDLLEXPORT void pg_stash_advice_worker_main(Datum main_arg);
      72                 :             : static void pgsa_append_tsv_escaped_string(StringInfo buf, const char *str);
      73                 :             : static void pgsa_detach_shmem(int code, Datum arg);
      74                 :             : static char *pgsa_next_tsv_field(char **cursor);
      75                 :             : static void pgsa_read_from_disk(void);
      76                 :             : static void pgsa_restore_entries(pgsa_saved_entry *entries, int num_entries);
      77                 :             : static void pgsa_restore_stashes(pgsa_saved_stash_table_hash *saved_stashes);
      78                 :             : static void pgsa_unescape_tsv_field(char *str, const char *filename,
      79                 :             :                                     unsigned lineno);
      80                 :             : static void pgsa_write_entries(pgsa_writer_context *wctx);
      81                 :             : pg_noreturn static void pgsa_write_error(pgsa_writer_context *wctx);
      82                 :             : static void pgsa_write_stashes(pgsa_writer_context *wctx);
      83                 :             : static void pgsa_write_to_disk(void);
      84                 :             : 
      85                 :             : /*
      86                 :             :  * Background worker entry point for pg_stash_advice persistence.
      87                 :             :  *
      88                 :             :  * On startup, if stashes_ready is set, we load previously saved
      89                 :             :  * stash data from disk.  Then we enter a loop, periodically checking whether
      90                 :             :  * any changes have been made (via the change_count atomic counter) and
      91                 :             :  * writing them to disk.  On shutdown, we perform a final write.
      92                 :             :  */
      93                 :             : PGDLLEXPORT void
      94                 :           4 : pg_stash_advice_worker_main(Datum main_arg)
      95                 :             : {
      96                 :             :     uint64      last_change_count;
      97                 :           4 :     TimestampTz last_write_time = 0;
      98                 :             : 
      99                 :             :     /* Establish signal handlers; once that's done, unblock signals. */
     100                 :           4 :     pqsignal(SIGTERM, SignalHandlerForShutdownRequest);
     101                 :           4 :     pqsignal(SIGHUP, SignalHandlerForConfigReload);
     102                 :           4 :     pqsignal(SIGUSR1, procsignal_sigusr1_handler);
     103                 :           4 :     BackgroundWorkerUnblockSignals();
     104                 :             : 
     105                 :             :     /* Log a debug message */
     106         [ -  + ]:           4 :     ereport(DEBUG1,
     107                 :             :             errmsg("pg_stash_advice worker started"));
     108                 :             : 
     109                 :             :     /* Set up session user so pgstat can report it. */
     110                 :           4 :     InitializeSessionUserIdStandalone();
     111                 :             : 
     112                 :             :     /* Report this worker in pg_stat_activity. */
     113                 :           4 :     pgstat_beinit();
     114                 :           4 :     pgstat_bestart_initial();
     115                 :           4 :     pgstat_bestart_final();
     116                 :             : 
     117                 :             :     /* Attach to shared memory structures. */
     118                 :           4 :     pgsa_attach();
     119                 :             : 
     120                 :             :     /* Set on-detach hook so that our PID will be cleared on exit. */
     121                 :           4 :     before_shmem_exit(pgsa_detach_shmem, 0);
     122                 :             : 
     123                 :             :     /*
     124                 :             :      * Store our PID in shared memory, unless there's already another worker
     125                 :             :      * running, in which case just exit.
     126                 :             :      */
     127                 :           4 :     LWLockAcquire(&pgsa_state->lock, LW_EXCLUSIVE);
     128         [ -  + ]:           4 :     if (pgsa_state->bgworker_pid != InvalidPid)
     129                 :             :     {
     130                 :           0 :         LWLockRelease(&pgsa_state->lock);
     131         [ #  # ]:           0 :         ereport(LOG,
     132                 :             :                 (errmsg("pg_stash_advice worker is already running under PID %d",
     133                 :             :                         (int) pgsa_state->bgworker_pid)));
     134                 :           0 :         return;
     135                 :             :     }
     136                 :           4 :     pgsa_state->bgworker_pid = MyProcPid;
     137                 :           4 :     LWLockRelease(&pgsa_state->lock);
     138                 :             : 
     139                 :             :     /*
     140                 :             :      * If pg_stash_advice.persist was set to true during
     141                 :             :      * process_shared_preload_libraries() and the data has not yet been
     142                 :             :      * successfully loaded, load it now.
     143                 :             :      */
     144         [ +  - ]:           4 :     if (pg_atomic_unlocked_test_flag(&pgsa_state->stashes_ready))
     145                 :             :     {
     146                 :           4 :         pgsa_read_from_disk();
     147                 :           4 :         pg_atomic_test_set_flag(&pgsa_state->stashes_ready);
     148                 :             :     }
     149                 :             : 
     150                 :             :     /* Note the current change count so we can detect future changes. */
     151                 :           4 :     last_change_count = pg_atomic_read_u64(&pgsa_state->change_count);
     152                 :             : 
     153                 :             :     /* Periodically write to disk until terminated. */
     154         [ +  + ]:          12 :     while (!ShutdownRequestPending)
     155                 :             :     {
     156                 :             :         /* In case of a SIGHUP, just reload the configuration. */
     157         [ -  + ]:           8 :         if (ConfigReloadPending)
     158                 :             :         {
     159                 :           0 :             ConfigReloadPending = false;
     160                 :           0 :             ProcessConfigFile(PGC_SIGHUP);
     161                 :             :         }
     162                 :             : 
     163         [ +  - ]:           8 :         if (pg_stash_advice_persist_interval <= 0)
     164                 :             :         {
     165                 :             :             /* Only writing at shutdown, so just wait forever. */
     166                 :           8 :             (void) WaitLatch(MyLatch,
     167                 :             :                              WL_LATCH_SET | WL_EXIT_ON_PM_DEATH,
     168                 :             :                              -1L,
     169                 :             :                              PG_WAIT_EXTENSION);
     170                 :             :         }
     171                 :             :         else
     172                 :             :         {
     173                 :             :             TimestampTz next_write_time;
     174                 :             :             long        delay_in_ms;
     175                 :             :             uint64      current_change_count;
     176                 :             : 
     177                 :             :             /* Compute when the next write should happen. */
     178                 :           0 :             next_write_time =
     179                 :           0 :                 TimestampTzPlusMilliseconds(last_write_time,
     180                 :             :                                             pg_stash_advice_persist_interval * 1000);
     181                 :             :             delay_in_ms =
     182                 :           0 :                 TimestampDifferenceMilliseconds(GetCurrentTimestamp(),
     183                 :             :                                                 next_write_time);
     184                 :             : 
     185                 :             :             /*
     186                 :             :              * When we reach next_write_time, we always update last_write_time
     187                 :             :              * (which is really the time at which we last considered writing),
     188                 :             :              * but we only actually write to disk if something has changed.
     189                 :             :              */
     190         [ #  # ]:           0 :             if (delay_in_ms <= 0)
     191                 :             :             {
     192                 :             :                 current_change_count =
     193                 :           0 :                     pg_atomic_read_u64(&pgsa_state->change_count);
     194         [ #  # ]:           0 :                 if (current_change_count != last_change_count)
     195                 :             :                 {
     196                 :           0 :                     pgsa_write_to_disk();
     197                 :           0 :                     last_change_count = current_change_count;
     198                 :             :                 }
     199                 :           0 :                 last_write_time = GetCurrentTimestamp();
     200                 :           0 :                 continue;
     201                 :             :             }
     202                 :             : 
     203                 :             :             /* Sleep until the next write time. */
     204                 :           0 :             (void) WaitLatch(MyLatch,
     205                 :             :                              WL_LATCH_SET | WL_TIMEOUT | WL_EXIT_ON_PM_DEATH,
     206                 :             :                              delay_in_ms,
     207                 :             :                              PG_WAIT_EXTENSION);
     208                 :             :         }
     209                 :             : 
     210                 :           8 :         ResetLatch(MyLatch);
     211                 :             :     }
     212                 :             : 
     213                 :             :     /* Write one last time before exiting. */
     214                 :           4 :     pgsa_write_to_disk();
     215                 :             : }
     216                 :             : 
     217                 :             : /*
     218                 :             :  * Clear our PID from shared memory on exit.
     219                 :             :  */
     220                 :             : static void
     221                 :           4 : pgsa_detach_shmem(int code, Datum arg)
     222                 :             : {
     223                 :           4 :     LWLockAcquire(&pgsa_state->lock, LW_EXCLUSIVE);
     224         [ +  - ]:           4 :     if (pgsa_state->bgworker_pid == MyProcPid)
     225                 :           4 :         pgsa_state->bgworker_pid = InvalidPid;
     226                 :           4 :     LWLockRelease(&pgsa_state->lock);
     227                 :           4 : }
     228                 :             : 
     229                 :             : /*
     230                 :             :  * Load advice stash data from a dump file on disk, if there is one.
     231                 :             :  */
     232                 :             : static void
     233                 :           4 : pgsa_read_from_disk(void)
     234                 :             : {
     235                 :             :     struct stat statbuf;
     236                 :             :     FILE       *file;
     237                 :             :     char       *filebuf;
     238                 :             :     size_t      nread;
     239                 :             :     char       *p;
     240                 :             :     unsigned    lineno;
     241                 :             :     pgsa_saved_stash_table_hash *saved_stashes;
     242                 :           4 :     int         num_stashes = 0;
     243                 :             :     pgsa_saved_entry *entries;
     244                 :           4 :     int         num_entries = 0;
     245                 :           4 :     int         max_entries = 64;
     246                 :             :     MemoryContext tmpcxt;
     247                 :             :     MemoryContext oldcxt;
     248                 :             : 
     249                 :             :     Assert(pgsa_entry_dshash != NULL);
     250                 :             : 
     251                 :             :     /*
     252                 :             :      * Clear any existing shared memory state.
     253                 :             :      *
     254                 :             :      * Normally, there won't be any, but if this function was called before
     255                 :             :      * and failed after beginning to apply changes to shared memory, then we
     256                 :             :      * need to get rid of any entries created at that time before trying
     257                 :             :      * again.
     258                 :             :      */
     259                 :           4 :     LWLockAcquire(&pgsa_state->lock, LW_EXCLUSIVE);
     260                 :           4 :     pgsa_reset_all_stashes();
     261                 :           4 :     LWLockRelease(&pgsa_state->lock);
     262                 :             : 
     263                 :             :     /* Open the dump file. If it doesn't exist, we're done. */
     264                 :           4 :     file = AllocateFile(PGSA_DUMP_FILE, "r");
     265         [ +  + ]:           4 :     if (!file)
     266                 :             :     {
     267         [ +  - ]:           2 :         if (errno == ENOENT)
     268                 :           2 :             return;
     269         [ #  # ]:           0 :         ereport(ERROR,
     270                 :             :                 (errcode_for_file_access(),
     271                 :             :                  errmsg("could not open file \"%s\": %m", PGSA_DUMP_FILE)));
     272                 :             :     }
     273                 :             : 
     274                 :             :     /* Use a temporary context for all parse-phase allocations. */
     275                 :           2 :     tmpcxt = AllocSetContextCreate(CurrentMemoryContext,
     276                 :             :                                    "pg_stash_advice load",
     277                 :             :                                    ALLOCSET_DEFAULT_SIZES);
     278                 :           2 :     oldcxt = MemoryContextSwitchTo(tmpcxt);
     279                 :             : 
     280                 :             :     /* Figure out how long the file is. */
     281         [ -  + ]:           2 :     if (fstat(fileno(file), &statbuf) != 0)
     282         [ #  # ]:           0 :         ereport(ERROR,
     283                 :             :                 (errcode_for_file_access(),
     284                 :             :                  errmsg("could not stat file \"%s\": %m", PGSA_DUMP_FILE)));
     285                 :             : 
     286                 :             :     /*
     287                 :             :      * Slurp the entire file into memory all at once.
     288                 :             :      *
     289                 :             :      * We could avoid this by reading the file incrementally and applying
     290                 :             :      * changes to pgsa_stash_dshash and pgsa_entry_dshash as we go. Given the
     291                 :             :      * lockout mechanism implemented by stashes_ready, that shouldn't have any
     292                 :             :      * user-visible behavioral consequences, but it would consume shared
     293                 :             :      * memory to no benefit. It seems better to buffer everything in private
     294                 :             :      * memory first, and then only apply the changes once the file has been
     295                 :             :      * successfully parsed in its entirety.
     296                 :             :      *
     297                 :             :      * That also has the advantage of possibly being more future-proof: if we
     298                 :             :      * decide to remove the stashes_ready mechanism in the future, or say
     299                 :             :      * allow for multiple save files, fully validating the file before
     300                 :             :      * applying any changes will become much more important.
     301                 :             :      *
     302                 :             :      * Of course, this approach does have one major disadvantage, which is
     303                 :             :      * that we'll temporarily use about twice as much memory as we're
     304                 :             :      * ultimately going to need, but that seems like it shouldn't be a problem
     305                 :             :      * in practice. If there's so much stashed advice that parsing the disk
     306                 :             :      * file runs us out of memory, something has gone terribly wrong. In that
     307                 :             :      * situation, there probably also isn't enough free memory for the
     308                 :             :      * workload that the advice is attempting to manipulate to run
     309                 :             :      * successfully.
     310                 :             :      */
     311                 :           2 :     filebuf = palloc_extended(statbuf.st_size + 1, MCXT_ALLOC_HUGE);
     312                 :           2 :     nread = fread(filebuf, 1, statbuf.st_size, file);
     313         [ -  + ]:           2 :     if (ferror(file))
     314         [ #  # ]:           0 :         ereport(ERROR,
     315                 :             :                 (errcode_for_file_access(),
     316                 :             :                  errmsg("could not read file \"%s\": %m", PGSA_DUMP_FILE)));
     317                 :           2 :     FreeFile(file);
     318                 :           2 :     filebuf[nread] = '\0';
     319                 :             : 
     320                 :             :     /* Initial memory allocations. */
     321                 :           2 :     saved_stashes = pgsa_saved_stash_table_create(tmpcxt, 64, NULL);
     322                 :           2 :     entries = palloc(max_entries * sizeof(pgsa_saved_entry));
     323                 :             : 
     324                 :             :     /*
     325                 :             :      * For memory and CPU efficiency, we parse the file in place. The end of
     326                 :             :      * each line gets replaced with a NUL byte, and then the end of each field
     327                 :             :      * within a line gets the same treatment. The advice string is unescaped
     328                 :             :      * in place, and stash names and query IDs can't contain any special
     329                 :             :      * characters. All of the resulting pointers point right back into the
     330                 :             :      * buffer; we only need additional memory to grow the 'entries' array and
     331                 :             :      * the 'saved_stashes' hash table.
     332                 :             :      */
     333         [ +  + ]:          13 :     for (p = filebuf, lineno = 1; *p != '\0'; lineno++)
     334                 :             :     {
     335                 :          11 :         char       *cursor = p;
     336                 :             :         char       *eol;
     337                 :             :         char       *line_type;
     338                 :             : 
     339                 :             :         /* Find end of line and NUL-terminate. */
     340                 :          11 :         eol = strchr(p, '\n');
     341         [ +  - ]:          11 :         if (eol != NULL)
     342                 :             :         {
     343                 :          11 :             *eol = '\0';
     344                 :          11 :             p = eol + 1;
     345   [ +  -  -  + ]:          11 :             if (eol > cursor && eol[-1] == '\r')
     346                 :           0 :                 eol[-1] = '\0';
     347                 :             :         }
     348                 :             :         else
     349                 :           0 :             p += strlen(p);
     350                 :             : 
     351                 :             :         /* Skip empty lines. */
     352         [ -  + ]:          11 :         if (*cursor == '\0')
     353                 :           0 :             continue;
     354                 :             : 
     355                 :             :         /* First field is the type of line, either "stash" or "entry". */
     356                 :          11 :         line_type = pgsa_next_tsv_field(&cursor);
     357         [ +  + ]:          11 :         if (strcmp(line_type, "stash") == 0)
     358                 :             :         {
     359                 :             :             char       *name;
     360                 :             :             bool        found;
     361                 :             : 
     362                 :             :             /* Second field should be the stash name. */
     363                 :           5 :             name = pgsa_next_tsv_field(&cursor);
     364   [ +  -  -  + ]:           5 :             if (name == NULL || *name == '\0')
     365         [ #  # ]:           0 :                 ereport(ERROR,
     366                 :             :                         (errcode(ERRCODE_DATA_CORRUPTED),
     367                 :             :                          errmsg("syntax error in file \"%s\" line %u: expected stash name",
     368                 :             :                                 PGSA_DUMP_FILE, lineno)));
     369                 :             : 
     370                 :             :             /* No further fields are expected. */
     371         [ -  + ]:           5 :             if (*cursor != '\0')
     372         [ #  # ]:           0 :                 ereport(ERROR,
     373                 :             :                         (errcode(ERRCODE_DATA_CORRUPTED),
     374                 :             :                          errmsg("syntax error in file \"%s\" line %u: expected end of line",
     375                 :             :                                 PGSA_DUMP_FILE, lineno)));
     376                 :             : 
     377                 :             :             /* Reject overlong stash names. */
     378         [ -  + ]:           5 :             if (strlen(name) >= NAMEDATALEN)
     379         [ #  # ]:           0 :                 ereport(ERROR,
     380                 :             :                         (errcode(ERRCODE_DATA_CORRUPTED),
     381                 :             :                          errmsg("syntax error in file \"%s\" line %u: stash name too long",
     382                 :             :                                 PGSA_DUMP_FILE, lineno)));
     383                 :             : 
     384                 :             :             /* Duplicate check. */
     385                 :           5 :             (void) pgsa_saved_stash_table_insert(saved_stashes, name, &found);
     386         [ -  + ]:           5 :             if (found)
     387         [ #  # ]:           0 :                 ereport(ERROR,
     388                 :             :                         (errcode(ERRCODE_DATA_CORRUPTED),
     389                 :             :                          errmsg("syntax error in file \"%s\" line %u: duplicate stash name \"%s\"",
     390                 :             :                                 PGSA_DUMP_FILE, lineno, name)));
     391                 :           5 :             num_stashes++;
     392                 :             :         }
     393         [ +  - ]:           6 :         else if (strcmp(line_type, "entry") == 0)
     394                 :             :         {
     395                 :             :             char       *stash_name;
     396                 :             :             char       *queryid_str;
     397                 :             :             char       *advice_str;
     398                 :             :             char       *endptr;
     399                 :             :             int64       queryId;
     400                 :             : 
     401                 :             :             /* Second field should be the stash name. */
     402                 :           6 :             stash_name = pgsa_next_tsv_field(&cursor);
     403         [ -  + ]:           6 :             if (stash_name == NULL)
     404         [ #  # ]:           0 :                 ereport(ERROR,
     405                 :             :                         (errcode(ERRCODE_DATA_CORRUPTED),
     406                 :             :                          errmsg("syntax error in file \"%s\" line %u: expected stash name",
     407                 :             :                                 PGSA_DUMP_FILE, lineno)));
     408                 :             : 
     409                 :             :             /* Third field should be the query ID. */
     410                 :           6 :             queryid_str = pgsa_next_tsv_field(&cursor);
     411         [ -  + ]:           6 :             if (queryid_str == NULL)
     412         [ #  # ]:           0 :                 ereport(ERROR,
     413                 :             :                         (errcode(ERRCODE_DATA_CORRUPTED),
     414                 :             :                          errmsg("syntax error in file \"%s\" line %u: expected query ID",
     415                 :             :                                 PGSA_DUMP_FILE, lineno)));
     416                 :             : 
     417                 :             :             /* Fourth field should be the advice string. */
     418                 :           6 :             advice_str = pgsa_next_tsv_field(&cursor);
     419         [ -  + ]:           6 :             if (advice_str == NULL)
     420         [ #  # ]:           0 :                 ereport(ERROR,
     421                 :             :                         (errcode(ERRCODE_DATA_CORRUPTED),
     422                 :             :                          errmsg("syntax error in file \"%s\" line %u: expected advice string",
     423                 :             :                                 PGSA_DUMP_FILE, lineno)));
     424                 :             : 
     425                 :             :             /* No further fields are expected. */
     426         [ -  + ]:           6 :             if (*cursor != '\0')
     427         [ #  # ]:           0 :                 ereport(ERROR,
     428                 :             :                         (errcode(ERRCODE_DATA_CORRUPTED),
     429                 :             :                          errmsg("syntax error in file \"%s\" line %u: expected end of line",
     430                 :             :                                 PGSA_DUMP_FILE, lineno)));
     431                 :             : 
     432                 :             :             /* Make sure the stash is one we've actually seen. */
     433         [ -  + ]:           6 :             if (pgsa_saved_stash_table_lookup(saved_stashes,
     434                 :             :                                               stash_name) == NULL)
     435         [ #  # ]:           0 :                 ereport(ERROR,
     436                 :             :                         (errcode(ERRCODE_DATA_CORRUPTED),
     437                 :             :                          errmsg("syntax error in file \"%s\" line %u: unknown stash \"%s\"",
     438                 :             :                                 PGSA_DUMP_FILE, lineno, stash_name)));
     439                 :             : 
     440                 :             :             /* Parse the query ID. */
     441                 :           6 :             errno = 0;
     442                 :           6 :             queryId = strtoi64(queryid_str, &endptr, 10);
     443   [ +  -  +  -  :           6 :             if (*endptr != '\0' || errno != 0 || queryid_str == endptr ||
             +  -  -  + ]
     444                 :             :                 queryId == 0)
     445         [ #  # ]:           0 :                 ereport(ERROR,
     446                 :             :                         (errcode(ERRCODE_DATA_CORRUPTED),
     447                 :             :                          errmsg("syntax error in file \"%s\" line %u: invalid query ID \"%s\"",
     448                 :             :                                 PGSA_DUMP_FILE, lineno, queryid_str)));
     449                 :             : 
     450                 :             :             /* Unescape the advice string. */
     451                 :           6 :             pgsa_unescape_tsv_field(advice_str, PGSA_DUMP_FILE, lineno);
     452                 :             : 
     453                 :             :             /* Append to the entry array. */
     454         [ -  + ]:           6 :             if (num_entries >= max_entries)
     455                 :             :             {
     456                 :           0 :                 max_entries *= 2;
     457                 :           0 :                 entries = repalloc(entries,
     458                 :             :                                    max_entries * sizeof(pgsa_saved_entry));
     459                 :             :             }
     460                 :           6 :             entries[num_entries].stash_name = stash_name;
     461                 :           6 :             entries[num_entries].queryId = queryId;
     462                 :           6 :             entries[num_entries].advice_string = advice_str;
     463                 :           6 :             num_entries++;
     464                 :             :         }
     465                 :             :         else
     466                 :             :         {
     467         [ #  # ]:           0 :             ereport(ERROR,
     468                 :             :                     (errcode(ERRCODE_DATA_CORRUPTED),
     469                 :             :                      errmsg("syntax error in file \"%s\" line %u: unrecognized line type",
     470                 :             :                             PGSA_DUMP_FILE, lineno)));
     471                 :             :         }
     472                 :             :     }
     473                 :             : 
     474                 :             :     /*
     475                 :             :      * Parsing succeeded. Apply everything to shared memory.
     476                 :             :      *
     477                 :             :      * At this point, we know that the file we just read is fully valid, but
     478                 :             :      * it's still possible for this to fail if, for example, DSA memory cannot
     479                 :             :      * be allocated. If that happens, the worker will die, the postmaster will
     480                 :             :      * eventually restart it, and we'll try again after clearing any data that
     481                 :             :      * we did manage to put into shared memory. (Note that we call
     482                 :             :      * pgsa_reset_all_stashes() at the top of this function.)
     483                 :             :      */
     484                 :           2 :     pgsa_restore_stashes(saved_stashes);
     485                 :           2 :     pgsa_restore_entries(entries, num_entries);
     486                 :             : 
     487                 :             :     /* Hooray, it worked! Notify the user. */
     488         [ +  - ]:           2 :     ereport(LOG,
     489                 :             :             (errmsg("loaded %d advice stashes and %d entries from \"%s\"",
     490                 :             :                     num_stashes, num_entries, PGSA_DUMP_FILE)));
     491                 :             : 
     492                 :             :     /* Clean up. */
     493                 :           2 :     MemoryContextSwitchTo(oldcxt);
     494                 :           2 :     MemoryContextDelete(tmpcxt);
     495                 :             : }
     496                 :             : 
     497                 :             : /*
     498                 :             :  * Write all advice stash data to disk.
     499                 :             :  *
     500                 :             :  * The file format is a simple TSV with a line-type prefix:
     501                 :             :  *   stash\tstash_name
     502                 :             :  *   entry\tstash_name\tquery_id\tadvice_string
     503                 :             :  */
     504                 :             : static void
     505                 :           4 : pgsa_write_to_disk(void)
     506                 :             : {
     507                 :           4 :     pgsa_writer_context wctx = {0};
     508                 :             :     MemoryContext tmpcxt;
     509                 :             :     MemoryContext oldcxt;
     510                 :             : 
     511                 :             :     Assert(pgsa_entry_dshash != NULL);
     512                 :             : 
     513                 :             :     /* Use a temporary context so all allocations are freed at the end. */
     514                 :           4 :     tmpcxt = AllocSetContextCreate(CurrentMemoryContext,
     515                 :             :                                    "pg_stash_advice dump",
     516                 :             :                                    ALLOCSET_DEFAULT_SIZES);
     517                 :           4 :     oldcxt = MemoryContextSwitchTo(tmpcxt);
     518                 :             : 
     519                 :             :     /* Set up the writer context. */
     520                 :           4 :     snprintf(wctx.pathname, MAXPGPATH, "%s.tmp", PGSA_DUMP_FILE);
     521                 :           4 :     wctx.file = AllocateFile(wctx.pathname, "w");
     522         [ -  + ]:           4 :     if (!wctx.file)
     523         [ #  # ]:           0 :         ereport(ERROR,
     524                 :             :                 (errcode_for_file_access(),
     525                 :             :                  errmsg("could not open file \"%s\": %m", wctx.pathname)));
     526                 :           4 :     wctx.nhash = pgsa_stash_name_table_create(tmpcxt, 64, NULL);
     527                 :           4 :     initStringInfo(&wctx.buf);
     528                 :             : 
     529                 :             :     /* Write stash lines, then entry lines. */
     530                 :           4 :     pgsa_write_stashes(&wctx);
     531                 :           4 :     pgsa_write_entries(&wctx);
     532                 :             : 
     533                 :             :     /*
     534                 :             :      * If nothing was written, remove both the temp file and any existing dump
     535                 :             :      * file rather than installing a zero-length file.
     536                 :             :      */
     537         [ +  + ]:           4 :     if (wctx.nhash->members == 0)
     538                 :             :     {
     539         [ -  + ]:           2 :         ereport(DEBUG1,
     540                 :             :                 errmsg("there are no advice stashes to save"));
     541                 :           2 :         FreeFile(wctx.file);
     542                 :           2 :         unlink(wctx.pathname);
     543         [ +  + ]:           2 :         if (unlink(PGSA_DUMP_FILE) == 0)
     544         [ -  + ]:           1 :             ereport(DEBUG1,
     545                 :             :                     errmsg("removed \"%s\"", PGSA_DUMP_FILE));
     546                 :             :     }
     547                 :             :     else
     548                 :             :     {
     549         [ -  + ]:           2 :         if (FreeFile(wctx.file) != 0)
     550                 :             :         {
     551                 :           0 :             int         save_errno = errno;
     552                 :             : 
     553                 :           0 :             unlink(wctx.pathname);
     554                 :           0 :             errno = save_errno;
     555         [ #  # ]:           0 :             ereport(ERROR,
     556                 :             :                     (errcode_for_file_access(),
     557                 :             :                      errmsg("could not close file \"%s\": %m",
     558                 :             :                             wctx.pathname)));
     559                 :             :         }
     560                 :           2 :         (void) durable_rename(wctx.pathname, PGSA_DUMP_FILE, ERROR);
     561                 :             : 
     562         [ +  - ]:           2 :         ereport(LOG,
     563                 :             :                 errmsg("saved %d advice stashes and %d entries to \"%s\"",
     564                 :             :                        (int) wctx.nhash->members, wctx.entries_written,
     565                 :             :                        PGSA_DUMP_FILE));
     566                 :             :     }
     567                 :             : 
     568                 :           4 :     MemoryContextSwitchTo(oldcxt);
     569                 :           4 :     MemoryContextDelete(tmpcxt);
     570                 :           4 : }
     571                 :             : 
     572                 :             : /*
     573                 :             :  * Append the TSV-escaped form of str to buf.
     574                 :             :  *
     575                 :             :  * Backslash, tab, newline, and carriage return are escaped with backslash
     576                 :             :  * sequences.  All other characters are passed through unchanged.
     577                 :             :  */
     578                 :             : static void
     579                 :           6 : pgsa_append_tsv_escaped_string(StringInfo buf, const char *str)
     580                 :             : {
     581         [ +  + ]:         100 :     for (const char *p = str; *p != '\0'; p++)
     582                 :             :     {
     583   [ +  +  +  -  :          94 :         switch (*p)
                      + ]
     584                 :             :         {
     585                 :           2 :             case '\\':
     586                 :           2 :                 appendStringInfoString(buf, "\\\\");
     587                 :           2 :                 break;
     588                 :           2 :             case '\t':
     589                 :           2 :                 appendStringInfoString(buf, "\\t");
     590                 :           2 :                 break;
     591                 :           2 :             case '\n':
     592                 :           2 :                 appendStringInfoString(buf, "\\n");
     593                 :           2 :                 break;
     594                 :           0 :             case '\r':
     595                 :           0 :                 appendStringInfoString(buf, "\\r");
     596                 :           0 :                 break;
     597                 :          88 :             default:
     598                 :          88 :                 appendStringInfoChar(buf, *p);
     599                 :          88 :                 break;
     600                 :             :         }
     601                 :             :     }
     602                 :           6 : }
     603                 :             : 
     604                 :             : /*
     605                 :             :  * Extract the next tab-delimited field from *cursor.
     606                 :             :  *
     607                 :             :  * The tab delimiter is replaced with '\0' and *cursor is advanced past it.
     608                 :             :  * If *cursor already points to '\0' (no more fields), returns NULL.
     609                 :             :  */
     610                 :             : static char *
     611                 :          34 : pgsa_next_tsv_field(char **cursor)
     612                 :             : {
     613                 :          34 :     char       *start = *cursor;
     614                 :          34 :     char       *p = start;
     615                 :             : 
     616         [ -  + ]:          34 :     if (*p == '\0')
     617                 :           0 :         return NULL;
     618                 :             : 
     619   [ +  +  +  + ]:         290 :     while (*p != '\0' && *p != '\t')
     620                 :         256 :         p++;
     621                 :             : 
     622         [ +  + ]:          34 :     if (*p == '\t')
     623                 :          23 :         *p++ = '\0';
     624                 :             : 
     625                 :          34 :     *cursor = p;
     626                 :          34 :     return start;
     627                 :             : }
     628                 :             : 
     629                 :             : /*
     630                 :             :  * Insert entries into shared memory from the parsed entry array.
     631                 :             :  */
     632                 :             : static void
     633                 :           2 : pgsa_restore_entries(pgsa_saved_entry *entries, int num_entries)
     634                 :             : {
     635                 :           2 :     LWLockAcquire(&pgsa_state->lock, LW_SHARED);
     636         [ +  + ]:           8 :     for (int i = 0; i < num_entries; i++)
     637                 :             :     {
     638         [ -  + ]:           6 :         ereport(DEBUG2,
     639                 :             :                 errmsg("restoring advice stash entry for \"%s\", query ID %" PRId64,
     640                 :             :                        entries[i].stash_name, entries[i].queryId));
     641                 :           6 :         pgsa_set_advice_string(entries[i].stash_name,
     642                 :           6 :                                entries[i].queryId,
     643                 :           6 :                                entries[i].advice_string);
     644                 :             :     }
     645                 :           2 :     LWLockRelease(&pgsa_state->lock);
     646                 :           2 : }
     647                 :             : 
     648                 :             : /*
     649                 :             :  * Create stashes in shared memory from the parsed stash hash table.
     650                 :             :  */
     651                 :             : static void
     652                 :           2 : pgsa_restore_stashes(pgsa_saved_stash_table_hash *saved_stashes)
     653                 :             : {
     654                 :             :     pgsa_saved_stash_table_iterator iter;
     655                 :             :     pgsa_saved_stash *s;
     656                 :             : 
     657                 :           2 :     LWLockAcquire(&pgsa_state->lock, LW_EXCLUSIVE);
     658                 :           2 :     pgsa_saved_stash_table_start_iterate(saved_stashes, &iter);
     659                 :           7 :     while ((s = pgsa_saved_stash_table_iterate(saved_stashes,
     660         [ +  + ]:           7 :                                                &iter)) != NULL)
     661                 :             :     {
     662         [ -  + ]:           5 :         ereport(DEBUG2,
     663                 :             :                 errmsg("restoring advice stash \"%s\"", s->name));
     664                 :           5 :         pgsa_create_stash(s->name);
     665                 :             :     }
     666                 :           2 :     LWLockRelease(&pgsa_state->lock);
     667                 :           2 : }
     668                 :             : 
     669                 :             : /*
     670                 :             :  * Unescape a TSV field in place.
     671                 :             :  *
     672                 :             :  * Recognized escape sequences are \\, \t, \n, and \r.  A trailing backslash
     673                 :             :  * or an unrecognized escape sequence is a syntax error.
     674                 :             :  */
     675                 :             : static void
     676                 :           6 : pgsa_unescape_tsv_field(char *str, const char *filename, unsigned lineno)
     677                 :             : {
     678                 :           6 :     char       *src = str;
     679                 :           6 :     char       *dst = str;
     680                 :             : 
     681         [ +  + ]:         100 :     while (*src != '\0')
     682                 :             :     {
     683                 :             :         /* Just pass through anything that's not a backslash-escape. */
     684         [ +  + ]:          94 :         if (likely(*src != '\\'))
     685                 :             :         {
     686                 :          88 :             *dst++ = *src++;
     687                 :          88 :             continue;
     688                 :             :         }
     689                 :             : 
     690                 :             :         /* Check what sort of escape we've got. */
     691   [ +  +  +  -  :           6 :         switch (src[1])
                   -  - ]
     692                 :             :         {
     693                 :           2 :             case '\\':
     694                 :           2 :                 *dst++ = '\\';
     695                 :           2 :                 break;
     696                 :           2 :             case 't':
     697                 :           2 :                 *dst++ = '\t';
     698                 :           2 :                 break;
     699                 :           2 :             case 'n':
     700                 :           2 :                 *dst++ = '\n';
     701                 :           2 :                 break;
     702                 :           0 :             case 'r':
     703                 :           0 :                 *dst++ = '\r';
     704                 :           0 :                 break;
     705                 :           0 :             case '\0':
     706         [ #  # ]:           0 :                 ereport(ERROR,
     707                 :             :                         (errcode(ERRCODE_DATA_CORRUPTED),
     708                 :             :                          errmsg("syntax error in file \"%s\" line %u: trailing backslash",
     709                 :             :                                 filename, lineno)));
     710                 :             :                 break;
     711                 :           0 :             default:
     712         [ #  # ]:           0 :                 ereport(ERROR,
     713                 :             :                         (errcode(ERRCODE_DATA_CORRUPTED),
     714                 :             :                          errmsg("syntax error in file \"%s\" line %u: unrecognized escape \"\\%c\"",
     715                 :             :                                 filename, lineno, src[1])));
     716                 :             :                 break;
     717                 :             :         }
     718                 :             : 
     719                 :             :         /* We consumed the backslash and the following character. */
     720                 :           6 :         src += 2;
     721                 :             :     }
     722                 :           6 :     *dst = '\0';
     723                 :           6 : }
     724                 :             : 
     725                 :             : /*
     726                 :             :  * Write an entry line for each advice entry.
     727                 :             :  */
     728                 :             : static void
     729                 :           4 : pgsa_write_entries(pgsa_writer_context *wctx)
     730                 :             : {
     731                 :             :     dshash_seq_status iter;
     732                 :             :     pgsa_entry *entry;
     733                 :             : 
     734                 :           4 :     dshash_seq_init(&iter, pgsa_entry_dshash, false);
     735         [ +  + ]:          10 :     while ((entry = dshash_seq_next(&iter)) != NULL)
     736                 :             :     {
     737                 :             :         pgsa_stash_name *n;
     738                 :             :         char       *advice_string;
     739                 :             : 
     740         [ -  + ]:           6 :         if (entry->advice_string == InvalidDsaPointer)
     741                 :           0 :             continue;
     742                 :             : 
     743                 :           6 :         n = pgsa_stash_name_table_lookup(wctx->nhash,
     744                 :             :                                          entry->key.pgsa_stash_id);
     745         [ -  + ]:           6 :         if (n == NULL)
     746                 :           0 :             continue;           /* orphan entry, skip */
     747                 :             : 
     748                 :           6 :         advice_string = dsa_get_address(pgsa_dsa_area, entry->advice_string);
     749                 :             : 
     750                 :           6 :         resetStringInfo(&wctx->buf);
     751                 :           6 :         appendStringInfo(&wctx->buf, "entry\t%s\t%" PRId64 "\t",
     752                 :             :                          n->name, entry->key.queryId);
     753                 :           6 :         pgsa_append_tsv_escaped_string(&wctx->buf, advice_string);
     754                 :           6 :         appendStringInfoChar(&wctx->buf, '\n');
     755                 :           6 :         fwrite(wctx->buf.data, 1, wctx->buf.len, wctx->file);
     756         [ -  + ]:           6 :         if (ferror(wctx->file))
     757                 :           0 :             pgsa_write_error(wctx);
     758                 :           6 :         wctx->entries_written++;
     759                 :             :     }
     760                 :           4 :     dshash_seq_term(&iter);
     761                 :           4 : }
     762                 :             : 
     763                 :             : /*
     764                 :             :  * Clean up and report a write error.  Does not return.
     765                 :             :  */
     766                 :             : static void
     767                 :           0 : pgsa_write_error(pgsa_writer_context *wctx)
     768                 :             : {
     769                 :           0 :     int         save_errno = errno;
     770                 :             : 
     771                 :           0 :     FreeFile(wctx->file);
     772                 :           0 :     unlink(wctx->pathname);
     773                 :           0 :     errno = save_errno;
     774         [ #  # ]:           0 :     ereport(ERROR,
     775                 :             :             (errcode_for_file_access(),
     776                 :             :              errmsg("could not write to file \"%s\": %m", wctx->pathname)));
     777                 :             : }
     778                 :             : 
     779                 :             : /*
     780                 :             :  * Write a stash line for each advice stash, and populate the ID-to-name
     781                 :             :  * hash table for use by pgsa_write_entries.
     782                 :             :  */
     783                 :             : static void
     784                 :           4 : pgsa_write_stashes(pgsa_writer_context *wctx)
     785                 :             : {
     786                 :             :     dshash_seq_status iter;
     787                 :             :     pgsa_stash *stash;
     788                 :             : 
     789                 :           4 :     dshash_seq_init(&iter, pgsa_stash_dshash, false);
     790         [ +  + ]:           9 :     while ((stash = dshash_seq_next(&iter)) != NULL)
     791                 :             :     {
     792                 :             :         pgsa_stash_name *n;
     793                 :             :         bool        found;
     794                 :             : 
     795                 :           5 :         n = pgsa_stash_name_table_insert(wctx->nhash, stash->pgsa_stash_id,
     796                 :             :                                          &found);
     797                 :             :         Assert(!found);
     798                 :           5 :         n->name = pstrdup(stash->name);
     799                 :             : 
     800                 :           5 :         resetStringInfo(&wctx->buf);
     801                 :           5 :         appendStringInfo(&wctx->buf, "stash\t%s\n", n->name);
     802                 :           5 :         fwrite(wctx->buf.data, 1, wctx->buf.len, wctx->file);
     803         [ -  + ]:           5 :         if (ferror(wctx->file))
     804                 :           0 :             pgsa_write_error(wctx);
     805                 :             :     }
     806                 :           4 :     dshash_seq_term(&iter);
     807                 :           4 : }
        

Generated by: LCOV version 2.0-1