LCOV - code coverage report
Current view: top level - contrib/pg_prewarm - pg_prewarm.c (source / functions) Hit Total Coverage
Test: PostgreSQL 17beta1 Lines: 55 67 82.1 %
Date: 2024-06-14 17:13:31 Functions: 4 4 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-2024, 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 "fmgr.h"
      20             : #include "miscadmin.h"
      21             : #include "storage/bufmgr.h"
      22             : #include "storage/read_stream.h"
      23             : #include "storage/smgr.h"
      24             : #include "utils/acl.h"
      25             : #include "utils/builtins.h"
      26             : #include "utils/lsyscache.h"
      27             : #include "utils/rel.h"
      28             : 
      29          10 : PG_MODULE_MAGIC;
      30             : 
      31          14 : PG_FUNCTION_INFO_V1(pg_prewarm);
      32             : 
      33             : typedef enum
      34             : {
      35             :     PREWARM_PREFETCH,
      36             :     PREWARM_READ,
      37             :     PREWARM_BUFFER,
      38             : } PrewarmType;
      39             : 
      40             : static PGIOAlignedBlock blockbuffer;
      41             : 
      42             : struct pg_prewarm_read_stream_private
      43             : {
      44             :     BlockNumber blocknum;
      45             :     int64       last_block;
      46             : };
      47             : 
      48             : static BlockNumber
      49       22384 : pg_prewarm_read_stream_next_block(ReadStream *stream,
      50             :                                   void *callback_private_data,
      51             :                                   void *per_buffer_data)
      52             : {
      53       22384 :     struct pg_prewarm_read_stream_private *p = callback_private_data;
      54             : 
      55       22384 :     if (p->blocknum <= p->last_block)
      56       18222 :         return p->blocknum++;
      57             : 
      58        4162 :     return InvalidBlockNumber;
      59             : }
      60             : 
      61             : /*
      62             :  * pg_prewarm(regclass, mode text, fork text,
      63             :  *            first_block int8, last_block int8)
      64             :  *
      65             :  * The first argument is the relation to be prewarmed; the second controls
      66             :  * how prewarming is done; legal options are 'prefetch', 'read', and 'buffer'.
      67             :  * The third is the name of the relation fork to be prewarmed.  The fourth
      68             :  * and fifth arguments specify the first and last block to be prewarmed.
      69             :  * If the fourth argument is NULL, it will be taken as 0; if the fifth argument
      70             :  * is NULL, it will be taken as the number of blocks in the relation.  The
      71             :  * return value is the number of blocks successfully prewarmed.
      72             :  */
      73             : Datum
      74        5506 : pg_prewarm(PG_FUNCTION_ARGS)
      75             : {
      76             :     Oid         relOid;
      77             :     text       *forkName;
      78             :     text       *type;
      79             :     int64       first_block;
      80             :     int64       last_block;
      81             :     int64       nblocks;
      82        5506 :     int64       blocks_done = 0;
      83             :     int64       block;
      84             :     Relation    rel;
      85             :     ForkNumber  forkNumber;
      86             :     char       *forkString;
      87             :     char       *ttype;
      88             :     PrewarmType ptype;
      89             :     AclResult   aclresult;
      90             : 
      91             :     /* Basic sanity checking. */
      92        5506 :     if (PG_ARGISNULL(0))
      93           0 :         ereport(ERROR,
      94             :                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
      95             :                  errmsg("relation cannot be null")));
      96        5506 :     relOid = PG_GETARG_OID(0);
      97        5506 :     if (PG_ARGISNULL(1))
      98           0 :         ereport(ERROR,
      99             :                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
     100             :                  errmsg("prewarm type cannot be null")));
     101        5506 :     type = PG_GETARG_TEXT_PP(1);
     102        5506 :     ttype = text_to_cstring(type);
     103        5506 :     if (strcmp(ttype, "prefetch") == 0)
     104           2 :         ptype = PREWARM_PREFETCH;
     105        5504 :     else if (strcmp(ttype, "read") == 0)
     106           2 :         ptype = PREWARM_READ;
     107        5502 :     else if (strcmp(ttype, "buffer") == 0)
     108        5502 :         ptype = PREWARM_BUFFER;
     109             :     else
     110             :     {
     111           0 :         ereport(ERROR,
     112             :                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
     113             :                  errmsg("invalid prewarm type"),
     114             :                  errhint("Valid prewarm types are \"prefetch\", \"read\", and \"buffer\".")));
     115             :         PG_RETURN_INT64(0);     /* Placate compiler. */
     116             :     }
     117        5506 :     if (PG_ARGISNULL(2))
     118           0 :         ereport(ERROR,
     119             :                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
     120             :                  errmsg("relation fork cannot be null")));
     121        5506 :     forkName = PG_GETARG_TEXT_PP(2);
     122        5506 :     forkString = text_to_cstring(forkName);
     123        5506 :     forkNumber = forkname_to_number(forkString);
     124             : 
     125             :     /* Open relation and check privileges. */
     126        5506 :     rel = relation_open(relOid, AccessShareLock);
     127        5506 :     aclresult = pg_class_aclcheck(relOid, GetUserId(), ACL_SELECT);
     128        5506 :     if (aclresult != ACLCHECK_OK)
     129           0 :         aclcheck_error(aclresult, get_relkind_objtype(rel->rd_rel->relkind), get_rel_name(relOid));
     130             : 
     131             :     /* Check that the fork exists. */
     132        5506 :     if (!smgrexists(RelationGetSmgr(rel), forkNumber))
     133           0 :         ereport(ERROR,
     134             :                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
     135             :                  errmsg("fork \"%s\" does not exist for this relation",
     136             :                         forkString)));
     137             : 
     138             :     /* Validate block numbers, or handle nulls. */
     139        5506 :     nblocks = RelationGetNumberOfBlocksInFork(rel, forkNumber);
     140        5506 :     if (PG_ARGISNULL(3))
     141        5506 :         first_block = 0;
     142             :     else
     143             :     {
     144           0 :         first_block = PG_GETARG_INT64(3);
     145           0 :         if (first_block < 0 || first_block >= nblocks)
     146           0 :             ereport(ERROR,
     147             :                     (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
     148             :                      errmsg("starting block number must be between 0 and %lld",
     149             :                             (long long) (nblocks - 1))));
     150             :     }
     151        5506 :     if (PG_ARGISNULL(4))
     152        5506 :         last_block = nblocks - 1;
     153             :     else
     154             :     {
     155           0 :         last_block = PG_GETARG_INT64(4);
     156           0 :         if (last_block < 0 || last_block >= nblocks)
     157           0 :             ereport(ERROR,
     158             :                     (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
     159             :                      errmsg("ending block number must be between 0 and %lld",
     160             :                             (long long) (nblocks - 1))));
     161             :     }
     162             : 
     163             :     /* Now we're ready to do the real work. */
     164        5506 :     if (ptype == PREWARM_PREFETCH)
     165             :     {
     166             : #ifdef USE_PREFETCH
     167             : 
     168             :         /*
     169             :          * In prefetch mode, we just hint the OS to read the blocks, but we
     170             :          * don't know whether it really does it, and we don't wait for it to
     171             :          * finish.
     172             :          *
     173             :          * It would probably be better to pass our prefetch requests in chunks
     174             :          * of a megabyte or maybe even a whole segment at a time, but there's
     175             :          * no practical way to do that at present without a gross modularity
     176             :          * violation, so we just do this.
     177             :          */
     178           4 :         for (block = first_block; block <= last_block; ++block)
     179             :         {
     180           2 :             CHECK_FOR_INTERRUPTS();
     181           2 :             PrefetchBuffer(rel, forkNumber, block);
     182           2 :             ++blocks_done;
     183             :         }
     184             : #else
     185             :         ereport(ERROR,
     186             :                 (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
     187             :                  errmsg("prefetch is not supported by this build")));
     188             : #endif
     189             :     }
     190        5504 :     else if (ptype == PREWARM_READ)
     191             :     {
     192             :         /*
     193             :          * In read mode, we actually read the blocks, but not into shared
     194             :          * buffers.  This is more portable than prefetch mode (it works
     195             :          * everywhere) and is synchronous.
     196             :          */
     197           4 :         for (block = first_block; block <= last_block; ++block)
     198             :         {
     199           2 :             CHECK_FOR_INTERRUPTS();
     200           2 :             smgrread(RelationGetSmgr(rel), forkNumber, block, blockbuffer.data);
     201           2 :             ++blocks_done;
     202             :         }
     203             :     }
     204        5502 :     else if (ptype == PREWARM_BUFFER)
     205             :     {
     206             :         struct pg_prewarm_read_stream_private p;
     207             :         ReadStream *stream;
     208             : 
     209             :         /*
     210             :          * In buffer mode, we actually pull the data into shared_buffers.
     211             :          */
     212             : 
     213             :         /* Set up the private state for our streaming buffer read callback. */
     214        5502 :         p.blocknum = first_block;
     215        5502 :         p.last_block = last_block;
     216             : 
     217        5502 :         stream = read_stream_begin_relation(READ_STREAM_FULL,
     218             :                                             NULL,
     219             :                                             rel,
     220             :                                             forkNumber,
     221             :                                             pg_prewarm_read_stream_next_block,
     222             :                                             &p,
     223             :                                             0);
     224             : 
     225       23724 :         for (block = first_block; block <= last_block; ++block)
     226             :         {
     227             :             Buffer      buf;
     228             : 
     229       18222 :             CHECK_FOR_INTERRUPTS();
     230       18222 :             buf = read_stream_next_buffer(stream, NULL);
     231       18222 :             ReleaseBuffer(buf);
     232       18222 :             ++blocks_done;
     233             :         }
     234             :         Assert(read_stream_next_buffer(stream, NULL) == InvalidBuffer);
     235        5502 :         read_stream_end(stream);
     236             :     }
     237             : 
     238             :     /* Close relation, release lock. */
     239        5506 :     relation_close(rel, AccessShareLock);
     240             : 
     241        5506 :     PG_RETURN_INT64(blocks_done);
     242             : }

Generated by: LCOV version 1.14