LCOV - code coverage report
Current view: top level - contrib/pg_prewarm - pg_prewarm.c (source / functions) Hit Total Coverage
Test: PostgreSQL 19devel Lines: 64 76 84.2 %
Date: 2025-11-13 06:18:16 Functions: 3 3 100.0 %
Legend: Lines: hit not hit

          Line data    Source code
       1             : /*-------------------------------------------------------------------------
       2             :  *
       3             :  * pg_prewarm.c
       4             :  *        prewarming utilities
       5             :  *
       6             :  * Copyright (c) 2010-2025, PostgreSQL Global Development Group
       7             :  *
       8             :  * IDENTIFICATION
       9             :  *        contrib/pg_prewarm/pg_prewarm.c
      10             :  *
      11             :  *-------------------------------------------------------------------------
      12             :  */
      13             : #include "postgres.h"
      14             : 
      15             : #include <sys/stat.h>
      16             : #include <unistd.h>
      17             : 
      18             : #include "access/relation.h"
      19             : #include "catalog/index.h"
      20             : #include "fmgr.h"
      21             : #include "miscadmin.h"
      22             : #include "storage/bufmgr.h"
      23             : #include "storage/lmgr.h"
      24             : #include "storage/read_stream.h"
      25             : #include "storage/smgr.h"
      26             : #include "utils/acl.h"
      27             : #include "utils/builtins.h"
      28             : #include "utils/lsyscache.h"
      29             : #include "utils/rel.h"
      30             : 
      31          12 : PG_MODULE_MAGIC_EXT(
      32             :                     .name = "pg_prewarm",
      33             :                     .version = PG_VERSION
      34             : );
      35             : 
      36          26 : PG_FUNCTION_INFO_V1(pg_prewarm);
      37             : 
      38             : typedef enum
      39             : {
      40             :     PREWARM_PREFETCH,
      41             :     PREWARM_READ,
      42             :     PREWARM_BUFFER,
      43             : } PrewarmType;
      44             : 
      45             : static PGIOAlignedBlock blockbuffer;
      46             : 
      47             : /*
      48             :  * pg_prewarm(regclass, mode text, fork text,
      49             :  *            first_block int8, last_block int8)
      50             :  *
      51             :  * The first argument is the relation to be prewarmed; the second controls
      52             :  * how prewarming is done; legal options are 'prefetch', 'read', and 'buffer'.
      53             :  * The third is the name of the relation fork to be prewarmed.  The fourth
      54             :  * and fifth arguments specify the first and last block to be prewarmed.
      55             :  * If the fourth argument is NULL, it will be taken as 0; if the fifth argument
      56             :  * is NULL, it will be taken as the number of blocks in the relation.  The
      57             :  * return value is the number of blocks successfully prewarmed.
      58             :  */
      59             : Datum
      60        5476 : pg_prewarm(PG_FUNCTION_ARGS)
      61             : {
      62             :     Oid         relOid;
      63             :     text       *forkName;
      64             :     text       *type;
      65             :     int64       first_block;
      66             :     int64       last_block;
      67             :     int64       nblocks;
      68        5476 :     int64       blocks_done = 0;
      69             :     int64       block;
      70             :     Relation    rel;
      71             :     ForkNumber  forkNumber;
      72             :     char       *forkString;
      73             :     char       *ttype;
      74             :     PrewarmType ptype;
      75             :     AclResult   aclresult;
      76             :     char        relkind;
      77             :     Oid         privOid;
      78             : 
      79             :     /* Basic sanity checking. */
      80        5476 :     if (PG_ARGISNULL(0))
      81           0 :         ereport(ERROR,
      82             :                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
      83             :                  errmsg("relation cannot be null")));
      84        5476 :     relOid = PG_GETARG_OID(0);
      85        5476 :     if (PG_ARGISNULL(1))
      86           0 :         ereport(ERROR,
      87             :                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
      88             :                  errmsg("prewarm type cannot be null")));
      89        5476 :     type = PG_GETARG_TEXT_PP(1);
      90        5476 :     ttype = text_to_cstring(type);
      91        5476 :     if (strcmp(ttype, "prefetch") == 0)
      92           2 :         ptype = PREWARM_PREFETCH;
      93        5474 :     else if (strcmp(ttype, "read") == 0)
      94           2 :         ptype = PREWARM_READ;
      95        5472 :     else if (strcmp(ttype, "buffer") == 0)
      96        5472 :         ptype = PREWARM_BUFFER;
      97             :     else
      98             :     {
      99           0 :         ereport(ERROR,
     100             :                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
     101             :                  errmsg("invalid prewarm type"),
     102             :                  errhint("Valid prewarm types are \"prefetch\", \"read\", and \"buffer\".")));
     103             :         PG_RETURN_INT64(0);     /* Placate compiler. */
     104             :     }
     105        5476 :     if (PG_ARGISNULL(2))
     106           0 :         ereport(ERROR,
     107             :                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
     108             :                  errmsg("relation fork cannot be null")));
     109        5476 :     forkName = PG_GETARG_TEXT_PP(2);
     110        5476 :     forkString = text_to_cstring(forkName);
     111        5476 :     forkNumber = forkname_to_number(forkString);
     112             : 
     113             :     /*
     114             :      * Open relation and check privileges.  If the relation is an index, we
     115             :      * must check the privileges on its parent table instead.
     116             :      */
     117        5476 :     relkind = get_rel_relkind(relOid);
     118        5476 :     if (relkind == RELKIND_INDEX ||
     119             :         relkind == RELKIND_PARTITIONED_INDEX)
     120             :     {
     121        3284 :         privOid = IndexGetRelation(relOid, true);
     122             : 
     123             :         /* Lock table before index to avoid deadlock. */
     124        3284 :         if (OidIsValid(privOid))
     125        3284 :             LockRelationOid(privOid, AccessShareLock);
     126             :     }
     127             :     else
     128        2192 :         privOid = relOid;
     129             : 
     130        5476 :     rel = relation_open(relOid, AccessShareLock);
     131             : 
     132             :     /*
     133             :      * It's possible that the relation with OID "privOid" was dropped and the
     134             :      * OID was reused before we locked it.  If that happens, we could be left
     135             :      * with the wrong parent table OID, in which case we must ERROR.  It's
     136             :      * possible that such a race would change the outcome of
     137             :      * get_rel_relkind(), too, but the worst case scenario there is that we'll
     138             :      * check privileges on the index instead of its parent table, which isn't
     139             :      * too terrible.
     140             :      */
     141        5476 :     if (!OidIsValid(privOid) ||
     142        3284 :         (privOid != relOid &&
     143        3284 :          privOid != IndexGetRelation(relOid, true)))
     144           0 :         ereport(ERROR,
     145             :                 (errcode(ERRCODE_UNDEFINED_TABLE),
     146             :                  errmsg("could not find parent table of index \"%s\"",
     147             :                         RelationGetRelationName(rel))));
     148             : 
     149        5476 :     aclresult = pg_class_aclcheck(privOid, GetUserId(), ACL_SELECT);
     150        5476 :     if (aclresult != ACLCHECK_OK)
     151           4 :         aclcheck_error(aclresult, get_relkind_objtype(rel->rd_rel->relkind), get_rel_name(relOid));
     152             : 
     153             :     /* Check that the relation has storage. */
     154        5472 :     if (!RELKIND_HAS_STORAGE(rel->rd_rel->relkind))
     155           2 :         ereport(ERROR,
     156             :                 (errcode(ERRCODE_WRONG_OBJECT_TYPE),
     157             :                  errmsg("relation \"%s\" does not have storage",
     158             :                         RelationGetRelationName(rel)),
     159             :                  errdetail_relkind_not_supported(rel->rd_rel->relkind)));
     160             : 
     161             :     /* Check that the fork exists. */
     162        5470 :     if (!smgrexists(RelationGetSmgr(rel), forkNumber))
     163           0 :         ereport(ERROR,
     164             :                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
     165             :                  errmsg("fork \"%s\" does not exist for this relation",
     166             :                         forkString)));
     167             : 
     168             :     /* Validate block numbers, or handle nulls. */
     169        5470 :     nblocks = RelationGetNumberOfBlocksInFork(rel, forkNumber);
     170        5470 :     if (PG_ARGISNULL(3))
     171        5470 :         first_block = 0;
     172             :     else
     173             :     {
     174           0 :         first_block = PG_GETARG_INT64(3);
     175           0 :         if (first_block < 0 || first_block >= nblocks)
     176           0 :             ereport(ERROR,
     177             :                     (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
     178             :                      errmsg("starting block number must be between 0 and %" PRId64,
     179             :                             (nblocks - 1))));
     180             :     }
     181        5470 :     if (PG_ARGISNULL(4))
     182        5470 :         last_block = nblocks - 1;
     183             :     else
     184             :     {
     185           0 :         last_block = PG_GETARG_INT64(4);
     186           0 :         if (last_block < 0 || last_block >= nblocks)
     187           0 :             ereport(ERROR,
     188             :                     (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
     189             :                      errmsg("ending block number must be between 0 and %" PRId64,
     190             :                             (nblocks - 1))));
     191             :     }
     192             : 
     193             :     /* Now we're ready to do the real work. */
     194        5470 :     if (ptype == PREWARM_PREFETCH)
     195             :     {
     196             : #ifdef USE_PREFETCH
     197             : 
     198             :         /*
     199             :          * In prefetch mode, we just hint the OS to read the blocks, but we
     200             :          * don't know whether it really does it, and we don't wait for it to
     201             :          * finish.
     202             :          *
     203             :          * It would probably be better to pass our prefetch requests in chunks
     204             :          * of a megabyte or maybe even a whole segment at a time, but there's
     205             :          * no practical way to do that at present without a gross modularity
     206             :          * violation, so we just do this.
     207             :          */
     208           4 :         for (block = first_block; block <= last_block; ++block)
     209             :         {
     210           2 :             CHECK_FOR_INTERRUPTS();
     211           2 :             PrefetchBuffer(rel, forkNumber, block);
     212           2 :             ++blocks_done;
     213             :         }
     214             : #else
     215             :         ereport(ERROR,
     216             :                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
     217             :                  errmsg("prefetch is not supported by this build")));
     218             : #endif
     219             :     }
     220        5468 :     else if (ptype == PREWARM_READ)
     221             :     {
     222             :         /*
     223             :          * In read mode, we actually read the blocks, but not into shared
     224             :          * buffers.  This is more portable than prefetch mode (it works
     225             :          * everywhere) and is synchronous.
     226             :          */
     227           4 :         for (block = first_block; block <= last_block; ++block)
     228             :         {
     229           2 :             CHECK_FOR_INTERRUPTS();
     230           2 :             smgrread(RelationGetSmgr(rel), forkNumber, block, blockbuffer.data);
     231           2 :             ++blocks_done;
     232             :         }
     233             :     }
     234        5466 :     else if (ptype == PREWARM_BUFFER)
     235             :     {
     236             :         BlockRangeReadStreamPrivate p;
     237             :         ReadStream *stream;
     238             : 
     239             :         /*
     240             :          * In buffer mode, we actually pull the data into shared_buffers.
     241             :          */
     242             : 
     243             :         /* Set up the private state for our streaming buffer read callback. */
     244        5466 :         p.current_blocknum = first_block;
     245        5466 :         p.last_exclusive = last_block + 1;
     246             : 
     247             :         /*
     248             :          * It is safe to use batchmode as block_range_read_stream_cb takes no
     249             :          * locks.
     250             :          */
     251        5466 :         stream = read_stream_begin_relation(READ_STREAM_MAINTENANCE |
     252             :                                             READ_STREAM_FULL |
     253             :                                             READ_STREAM_USE_BATCHING,
     254             :                                             NULL,
     255             :                                             rel,
     256             :                                             forkNumber,
     257             :                                             block_range_read_stream_cb,
     258             :                                             &p,
     259             :                                             0);
     260             : 
     261       24134 :         for (block = first_block; block <= last_block; ++block)
     262             :         {
     263             :             Buffer      buf;
     264             : 
     265       18668 :             CHECK_FOR_INTERRUPTS();
     266       18668 :             buf = read_stream_next_buffer(stream, NULL);
     267       18668 :             ReleaseBuffer(buf);
     268       18668 :             ++blocks_done;
     269             :         }
     270             :         Assert(read_stream_next_buffer(stream, NULL) == InvalidBuffer);
     271        5466 :         read_stream_end(stream);
     272             :     }
     273             : 
     274             :     /* Close relation, release locks. */
     275        5470 :     relation_close(rel, AccessShareLock);
     276             : 
     277        5470 :     if (privOid != relOid)
     278        3282 :         UnlockRelationOid(privOid, AccessShareLock);
     279             : 
     280        5470 :     PG_RETURN_INT64(blocks_done);
     281             : }

Generated by: LCOV version 1.16