From 570bdcc5387dc829eb891a58ef5fa4533a5982e1 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 v5 3/3] 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.

Author: Nazir Bilal Yavuz <byavuz81@gmail.com>
Reviewed-by: Andres Freund <andres@anarazel.de>
Reviewed-by: Aidar Imamov <a.imamov@postgrespro.ru>
Reviewed-by: Joseph Koshakow <koshy44@gmail.com>
Discussion: https://postgr.es/m/CAN55FZ0h_YoSqqutxV6DES1RW8ig6wcA8CR9rJk358YRMxZFmw%40mail.gmail.com
---
 .../pg_buffercache--1.5--1.6.sql              | 10 ++++
 contrib/pg_buffercache/pg_buffercache_pages.c | 43 +++++++++++++++
 doc/src/sgml/pgbuffercache.sgml               | 54 +++++++++++++++++--
 src/backend/storage/buffer/bufmgr.c           | 54 +++++++++++++++++++
 src/include/storage/bufmgr.h                  |  1 +
 5 files changed, 159 insertions(+), 3 deletions(-)

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 d54bb1fd6f8..b40ee2599a4 100644
--- a/contrib/pg_buffercache/pg_buffercache--1.5--1.6.sql
+++ b/contrib/pg_buffercache/pg_buffercache--1.5--1.6.sql
@@ -20,3 +20,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 2f275354c56..23415016a8f 100644
--- a/contrib/pg_buffercache/pg_buffercache_pages.c
+++ b/contrib/pg_buffercache/pg_buffercache_pages.c
@@ -75,6 +75,8 @@ PG_FUNCTION_INFO_V1(pg_buffercache_usage_counts);
 PG_FUNCTION_INFO_V1(pg_buffercache_evict);
 PG_FUNCTION_INFO_V1(pg_buffercache_evict_relation);
 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)
@@ -490,3 +492,44 @@ pg_buffercache_evict_all(PG_FUNCTION_ARGS)
 
 	PG_RETURN_DATUM(result);
 }
+
+/*
+ * Try to mark a shared buffer as dirty.
+ */
+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 all the shared buffers as dirty.
+ */
+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);
+}
diff --git a/doc/src/sgml/pgbuffercache.sgml b/doc/src/sgml/pgbuffercache.sgml
index 681c74251d4..baf64e07b4e 100644
--- a/doc/src/sgml/pgbuffercache.sgml
+++ b/doc/src/sgml/pgbuffercache.sgml
@@ -35,14 +35,24 @@
   <primary>pg_buffercache_evict_all</primary>
  </indexterm>
 
+ <indexterm>
+  <primary>pg_buffercache_mark_dirty</primary>
+ </indexterm>
+
+ <indexterm>
+  <primary>pg_buffercache_mark_dirty_all</primary>
+ </indexterm>
+
  <para>
   This module provides the <function>pg_buffercache_pages()</function>
   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>, the
-  <function>pg_buffercache_evict_relation()</function> function and the
-  <function>pg_buffercache_evict_all()</function> function.
+  <function>pg_buffercache_evict()</function> function, the
+  <function>pg_buffercache_evict_relation()</function> function, the
+  <function>pg_buffercache_evict_all()</function> function, the
+  <function>pg_buffercache_mark_dirty()</function> function and the
+  <function>pg_buffercache_mark_dirty_all()</function> function.
  </para>
 
  <para>
@@ -87,6 +97,18 @@
   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.  Use of this function is
+  restricted to superusers only.
+ </para>
+
  <sect2 id="pgbuffercache-pg-buffercache">
   <title>The <structname>pg_buffercache</structname> View</title>
 
@@ -433,6 +455,32 @@
   </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 result is immediately out of
+   date upon return, as the buffer might become valid again at any time due to
+   concurrent activity.  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 the <function>pg_buffercache_mark_dirty()</function> function.
+   The difference is, the <function>pg_buffercache_mark_dirty_all()</function>
+   function does not take an argument; instead it tries to mark all buffers
+   dirty in the buffer pool.  The result is immediately out of date upon
+   return, as the buffer might become valid again at any time due to
+   concurrent activity.  The function is intended for developer testing only.
+  </para>
+ </sect2>
+
 <sect2 id="pgbuffercache-sample-output">
   <title>Sample Output</title>
 
diff --git a/src/backend/storage/buffer/bufmgr.c b/src/backend/storage/buffer/bufmgr.c
index 5d82e3fa297..1f902c4b54f 100644
--- a/src/backend/storage/buffer/bufmgr.c
+++ b/src/backend/storage/buffer/bufmgr.c
@@ -6643,6 +6643,60 @@ EvictRelUnpinnedBuffers(Relation rel, int32 *buffers_evicted, int32 *buffers_flu
 	}
 }
 
+/*
+ * Try to mark the provided shared buffer as dirty.
+ *
+ * This function is intended for testing/development use only!
+ *
+ * Same as EvictUnpinnedBuffer() but with MarkBufferDirty() call inside.
+ *
+ * Returns true if the buffer was already dirty or it has successfully been
+ * marked as dirty.
+ */
+bool
+MarkUnpinnedBufferDirty(Buffer buf)
+{
+	BufferDesc *desc;
+	uint32		buf_state;
+
+	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));
+	}
+
+	UnpinBuffer(desc);
+
+	return true;
+}
+
 /*
  * Generic implementation of the AIO handle staging callback for readv/writev
  * on local/shared buffers.
diff --git a/src/include/storage/bufmgr.h b/src/include/storage/bufmgr.h
index ac6a2d521f3..a2e65a1d918 100644
--- a/src/include/storage/bufmgr.h
+++ b/src/include/storage/bufmgr.h
@@ -307,6 +307,7 @@ extern void LimitAdditionalLocalPins(uint32 *additional_pins);
 extern bool EvictUnpinnedBuffer(Buffer buf, uint32 buf_state, bool *flushed);
 extern void EvictRelUnpinnedBuffers(Relation rel, int32 *buffers_evicted,
 									int32 *buffers_flushed);
+extern bool MarkUnpinnedBufferDirty(Buffer buf);
 
 /* in buf_init.c */
 extern void BufferManagerShmemInit(void);
-- 
2.43.0

