From 19b0a56e72525ee1c9754fb06ab5f31d41e43ca7 Mon Sep 17 00:00:00 2001
From: Nazir Bilal Yavuz <byavuz81@gmail.com>
Date: Wed, 25 Dec 2024 15:46:10 +0300
Subject: [PATCH v1 2/2] Add pg_buffercache_mark_dirty[_all]() functions for
 testing

This commit introduces two new functions for marking shared buffers as
dirty:

pg_buffercache_mark_dirty(): Marks a specific shared buffer as dirty.
pg_buffercache_mark_dirty_all(): Marks all shared buffers as dirty in a
single operation.

The pg_buffercache_mark_dirty_all() function provides an efficient
way to dirty the entire buffer pool (e.g., ~550ms vs. ~70ms for 16GB of
shared buffers), complementing pg_buffercache_mark_dirty() for more
granular control.

These functions are intended for developer testing and debugging
scenarios, enabling users to simulate various buffer pool states and
test write-back behavior. Both functions are superuser-only.
---
 src/include/storage/bufmgr.h                  |  1 +
 src/backend/storage/buffer/bufmgr.c           | 62 +++++++++++++++++++
 doc/src/sgml/pgbuffercache.sgml               | 40 +++++++++++-
 .../pg_buffercache--1.5--1.6.sql              | 10 +++
 contrib/pg_buffercache/pg_buffercache_pages.c | 43 +++++++++++++
 5 files changed, 154 insertions(+), 2 deletions(-)

diff --git a/src/include/storage/bufmgr.h b/src/include/storage/bufmgr.h
index 7f4eeca4afd..ac077402dd9 100644
--- a/src/include/storage/bufmgr.h
+++ b/src/include/storage/bufmgr.h
@@ -298,6 +298,7 @@ extern void LimitAdditionalPins(uint32 *additional_pins);
 extern void LimitAdditionalLocalPins(uint32 *additional_pins);
 
 extern bool EvictUnpinnedBuffer(Buffer buf, bool *flushed);
+extern bool MarkUnpinnedBufferDirty(Buffer buf);
 
 /* in buf_init.c */
 extern void BufferManagerShmemInit(void);
diff --git a/src/backend/storage/buffer/bufmgr.c b/src/backend/storage/buffer/bufmgr.c
index d2a93cf7cc0..7ab8bcaa228 100644
--- a/src/backend/storage/buffer/bufmgr.c
+++ b/src/backend/storage/buffer/bufmgr.c
@@ -6146,3 +6146,65 @@ EvictUnpinnedBuffer(Buffer buf, bool *flushed)
 
 	return result;
 }
+
+/*
+ * MarkUnpinnedBufferDirty
+ *
+ * This function is intended for testing/development use only!
+ *
+ * To succeed, the buffer must not be pinned on entry, so if the caller had a
+ * particular block in mind, it might already have been replaced by some other
+ * block by the time this function runs.  It's also unpinned on return, so the
+ * buffer might be occupied and flushed by the time control is returned.  This
+ * inherent raciness without other interlocking makes the function unsuitable
+ * for non-testing usage.
+ *
+ * Returns true if the buffer was not dirty and it has now been marked as
+ * dirty.  Returns false if it wasn't valid, if it couldn't be marked as dirty
+ * due to a pin, or if the buffer was already dirty.
+ */
+bool
+MarkUnpinnedBufferDirty(Buffer buf)
+{
+	BufferDesc *desc;
+	uint32		buf_state;
+	bool		result = false;
+
+	Assert(!BufferIsLocal(buf));
+
+	/* Make sure we can pin the buffer. */
+	ResourceOwnerEnlarge(CurrentResourceOwner);
+	ReservePrivateRefCountEntry();
+
+	desc = GetBufferDescriptor(buf - 1);
+
+	/* Lock the header and check if it's valid. */
+	buf_state = LockBufHdr(desc);
+	if ((buf_state & BM_VALID) == 0)
+	{
+		UnlockBufHdr(desc, buf_state);
+		return false;
+	}
+
+	/* Check that it's not pinned already. */
+	if (BUF_STATE_GET_REFCOUNT(buf_state) > 0)
+	{
+		UnlockBufHdr(desc, buf_state);
+		return false;
+	}
+
+	PinBuffer_Locked(desc);		/* releases spinlock */
+
+	/* If it was not already dirty, mark it as dirty. */
+	if (!(buf_state & BM_DIRTY))
+	{
+		LWLockAcquire(BufferDescriptorGetContentLock(desc), LW_EXCLUSIVE);
+		MarkBufferDirty(buf);
+		LWLockRelease(BufferDescriptorGetContentLock(desc));
+		result = true;
+	}
+
+	UnpinBuffer(desc);
+
+	return result;
+}
diff --git a/doc/src/sgml/pgbuffercache.sgml b/doc/src/sgml/pgbuffercache.sgml
index df4d90a650a..a1c5036a383 100644
--- a/doc/src/sgml/pgbuffercache.sgml
+++ b/doc/src/sgml/pgbuffercache.sgml
@@ -32,8 +32,10 @@
   function (wrapped in the <structname>pg_buffercache</structname> view),
   the <function>pg_buffercache_summary()</function> function, the
   <function>pg_buffercache_usage_counts()</function> function, the
-  <function>pg_buffercache_evict()</function> and the
-  <function>pg_buffercache_evict_all()</function> function.
+  <function>pg_buffercache_evict()</function> the
+  <function>pg_buffercache_evict_all()</function>, the
+  <function>pg_buffercache_mark_dirty()</function> and the
+  <function>pg_buffercache_mark_dirty_all()</function> function.
  </para>
 
  <para>
@@ -72,6 +74,18 @@
   evicted and flushed. Use of this function is restricted to superusers only.
  </para>
 
+ <para>
+  The <function>pg_buffercache_mark_dirty()</function> function allows a block
+  to be marked as dirty from the buffer pool given a buffer identifier.  Use of
+  this function is restricted to superusers only.
+ </para>
+
+ <para>
+  The <function>pg_buffercache_mark_dirty_all()</function> function tries to
+  mark all buffers dirty in the buffer pool. It returns how many buffers are
+  dirtied. Use of this function is restricted to superusers only.
+ </para>
+
  <sect2 id="pgbuffercache-pg-buffercache">
   <title>The <structname>pg_buffercache</structname> View</title>
 
@@ -396,6 +410,28 @@
   </para>
  </sect2>
 
+ <sect2 id="pgbuffercache-pg-buffercache-mark-dirty">
+  <title>The <structname>pg_buffercache_mark_dirty</structname> Function</title>
+  <para>
+   The <function>pg_buffercache_mark_dirty()</function> function takes a buffer
+   identifier, as shown in the <structfield>bufferid</structfield> column of
+   the <structname>pg_buffercache</structname> view.  It returns true on success,
+   and false if the buffer wasn't valid or if it couldn't be marked as dirty
+   because it was pinned.  The function is intended for developer testing only.
+  </para>
+ </sect2>
+
+ <sect2 id="pgbuffercache-pg-buffercache-mark-dirty-all">
+  <title>The <structname>pg_buffercache_mark_dirty_all</structname> Function</title>
+  <para>
+   The <function>pg_buffercache_mark_dirty_all()</function> function is very similar
+   to <function>pg_buffercache_mark_dirty()</function> function.  The difference is,
+   the <function>pg_buffercache_mark_dirty_all()</function> does not take argument;
+   instead it tries to mark all buffers dirty in the buffer pool.  The function is
+   intended for developer testing only.
+  </para>
+ </sect2>
+
 <sect2 id="pgbuffercache-sample-output">
   <title>Sample Output</title>
 
diff --git a/contrib/pg_buffercache/pg_buffercache--1.5--1.6.sql b/contrib/pg_buffercache/pg_buffercache--1.5--1.6.sql
index c08b799be94..e725fef1247 100644
--- a/contrib/pg_buffercache/pg_buffercache--1.5--1.6.sql
+++ b/contrib/pg_buffercache/pg_buffercache--1.5--1.6.sql
@@ -5,3 +5,13 @@ CREATE FUNCTION pg_buffercache_evict_all(
     OUT buffers_flushed int4)
 AS 'MODULE_PATHNAME', 'pg_buffercache_evict_all'
 LANGUAGE C PARALLEL SAFE VOLATILE;
+
+CREATE FUNCTION pg_buffercache_mark_dirty(IN int)
+RETURNS bool
+AS 'MODULE_PATHNAME', 'pg_buffercache_mark_dirty'
+LANGUAGE C PARALLEL SAFE VOLATILE STRICT;
+
+CREATE FUNCTION pg_buffercache_mark_dirty_all()
+RETURNS INT4
+AS 'MODULE_PATHNAME', 'pg_buffercache_mark_dirty_all'
+LANGUAGE C PARALLEL SAFE VOLATILE;
diff --git a/contrib/pg_buffercache/pg_buffercache_pages.c b/contrib/pg_buffercache/pg_buffercache_pages.c
index ea7c0e6934c..1a0992467b2 100644
--- a/contrib/pg_buffercache/pg_buffercache_pages.c
+++ b/contrib/pg_buffercache/pg_buffercache_pages.c
@@ -66,6 +66,8 @@ PG_FUNCTION_INFO_V1(pg_buffercache_summary);
 PG_FUNCTION_INFO_V1(pg_buffercache_usage_counts);
 PG_FUNCTION_INFO_V1(pg_buffercache_evict);
 PG_FUNCTION_INFO_V1(pg_buffercache_evict_all);
+PG_FUNCTION_INFO_V1(pg_buffercache_mark_dirty);
+PG_FUNCTION_INFO_V1(pg_buffercache_mark_dirty_all);
 
 Datum
 pg_buffercache_pages(PG_FUNCTION_ARGS)
@@ -412,3 +414,44 @@ pg_buffercache_evict_all(PG_FUNCTION_ARGS)
 
 	PG_RETURN_DATUM(result);
 }
+
+/*
+ * Try to mark dirty a shared buffer.
+ */
+Datum
+pg_buffercache_mark_dirty(PG_FUNCTION_ARGS)
+{
+	Buffer		buf = PG_GETARG_INT32(0);
+
+	if (!superuser())
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
+				 errmsg("must be superuser to use pg_buffercache_mark_dirty function")));
+
+	if (buf < 1 || buf > NBuffers)
+		elog(ERROR, "bad buffer ID: %d", buf);
+
+	PG_RETURN_BOOL(MarkUnpinnedBufferDirty(buf));
+}
+
+/*
+ * Try to mark dirty all shared buffers.
+ */
+Datum
+pg_buffercache_mark_dirty_all(PG_FUNCTION_ARGS)
+{
+	int32		buffers_dirtied = 0;
+
+	if (!superuser())
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
+				 errmsg("must be superuser to use pg_buffercache_mark_dirty_all function")));
+
+	for (int buf = 1; buf < NBuffers; buf++)
+	{
+		if (MarkUnpinnedBufferDirty(buf))
+			buffers_dirtied++;
+	}
+
+	PG_RETURN_INT32(buffers_dirtied);
+}
-- 
2.45.2

