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