Fixup: Adjust SLRU download hook API

Author: mmeent-databricks

src/backend/access/transam/slru.c

@@ -725,64 +725,29 @@ SlruInternalWritePage(SlruCtl ctl, int slotno, SlruWriteAll fdata)
 }
 
 /*
- * NEON: we do not want to include large pg_xact/multixact files in basebackup and prefer
- * to download them on demand to reduce startup time.
- * If SLRU segment is not found, we try to download it from page server
+ * NEON: we do not want to include large pg_xact/multixact files in the
+ * basebackup and prefer to download them on demand to reduce startup time.
+ *
+ * If SLRU segment is not found, we try to download it from the pageserver.
+ *
+ * Returns:
+ *   true if the file was successfully downloaded
+ *   false if the file was not found in pageserver
+ *   ereports if some other error happened
  */
-static int
-SimpleLruDownloadSegment(SlruCtl ctl, int pageno, char const* path)
+static bool
+SimpleLruDownloadSegment(SlruCtl ctl, int pageno, char const *path)
 {
 	SlruShared	shared = ctl->shared;
-	int segno;
-	int fd = -1;
-	int n_blocks;
-	char* buffer;
+	int			segno;
 
 	/* If page is greater than latest written page, then do not try to download segment from server */
 	if (ctl->PagePrecedes(pg_atomic_read_u64(&shared->latest_page_number), pageno))
-		return -1;
+		return false;
 
 	segno = pageno / SLRU_PAGES_PER_SEGMENT;
 
-	if (neon_use_communicator_worker) {
-		buffer = NULL;
-	} else {
-		buffer = palloc(BLCKSZ * SLRU_PAGES_PER_SEGMENT);
-	}
-
-	n_blocks = read_slru_segment(path, segno, buffer);
-	if (n_blocks > 0)
-	{
-		fd = OpenTransientFile(path, O_RDWR | O_CREAT | PG_BINARY);
-		if (fd < 0)
-		{
-			slru_errcause = SLRU_OPEN_FAILED;
-			slru_errno = errno;
-			pfree(buffer);
-			return -1;
-		}
-
-		if (!neon_use_communicator_worker) {
-			errno = 0;
-			pgstat_report_wait_start(WAIT_EVENT_SLRU_WRITE);
-			if (pg_pwrite(fd, buffer, n_blocks*BLCKSZ, 0) != n_blocks*BLCKSZ)
-			{
-				pgstat_report_wait_end();
-				/* if write didn't set errno, assume problem is no disk space */
-				if (errno == 0)
-					errno = ENOSPC;
-				slru_errcause = SLRU_WRITE_FAILED;
-				slru_errno = errno;
-
-				CloseTransientFile(fd);
-				pfree(buffer);
-				return -1;
-			}
-			pgstat_report_wait_end();
-		}
-	}
-	pfree(buffer);
-	return fd;
+	return smgr_read_slru_segment(path, segno);
 }
 
 /*
@@ -813,29 +778,34 @@ SimpleLruDoesPhysicalPageExist(SlruCtl ctl, int64 pageno)
 	int			fd;
 	bool		result;
 	off_t		endpos;
+	bool		attempted_download = false;
 
 	/* update the stats counter of checked pages */
 	pgstat_count_slru_page_exists(ctl->shared->slru_stats_idx);
 
 	SlruFileName(ctl, path, segno);
 
+ retry_after_download:
 	fd = OpenTransientFile(path, O_RDONLY | PG_BINARY);
 	if (fd < 0)
 	{
-		/* expected: file doesn't exist */
-		if (errno == ENOENT)
+		if (errno == ENOENT && !attempted_download)
 		{
-			fd = SimpleLruDownloadSegment(ctl, pageno, path);
-			if (fd < 0)
-				return false;
-		}
-		else
-		{
-			/* report error normally */
-			slru_errcause = SLRU_OPEN_FAILED;
-			slru_errno = errno;
-			SlruReportIOError(ctl, pageno, 0);
+			/* Try to download the file from the pageserver */
+			attempted_download = true;
+			if (SimpleLruDownloadSegment(ctl, pageno, path))
+				goto retry_after_download;
+			errno = ENOENT;
 		}
+
+		/* expected: file doesn't exist */
+		if (errno == ENOENT)
+			return false;
+
+		/* report error normally */
+		slru_errcause = SLRU_OPEN_FAILED;
+		slru_errno = errno;
+		SlruReportIOError(ctl, pageno, 0);
 	}
 
 	if ((endpos = lseek(fd, 0, SEEK_END)) < 0)
@@ -876,6 +846,7 @@ SlruPhysicalReadPage(SlruCtl ctl, int64 pageno, int slotno)
 	off_t		offset = rpageno * BLCKSZ;
 	char		path[MAXPGPATH];
 	int			fd;
+	bool		attempted_download = false;
 
 	SlruFileName(ctl, path, segno);
 
@@ -886,34 +857,31 @@ SlruPhysicalReadPage(SlruCtl ctl, int64 pageno, int slotno)
 	 * SlruPhysicalWritePage).  Hence, if we are InRecovery, allow the case
 	 * where the file doesn't exist, and return zeroes instead.
 	 */
+ retry_after_download:
 	fd = OpenTransientFile(path, O_RDONLY | PG_BINARY);
 	if (fd < 0)
 	{
-		if (errno != ENOENT)
+		if (errno == ENOENT && !attempted_download)
+		{
+			/* Try to download the file from the pageserver */
+			attempted_download = true;
+			if (SimpleLruDownloadSegment(ctl, pageno, path))
+				goto retry_after_download;
+			errno = ENOENT;
+		}
+
+		if (errno != ENOENT || !InRecovery)
 		{
 			slru_errcause = SLRU_OPEN_FAILED;
 			slru_errno = errno;
 			return false;
 		}
 
-		fd = SimpleLruDownloadSegment(ctl, pageno, path);
-		if (fd < 0)
-		{
-			if (!InRecovery)
-			{
-				slru_errcause = SLRU_OPEN_FAILED;
-				slru_errno = errno;
-				return false;
-			}
-			else
-			{
-				ereport(LOG,
-						(errmsg("file \"%s\" doesn't exist, reading as zeroes",
-								path)));
-				MemSet(shared->page_buffer[slotno], 0, BLCKSZ);
-				return true;
-			}
-		}
+		ereport(LOG,
+				(errmsg("file \"%s\" doesn't exist, reading as zeroes",
+						path)));
+		MemSet(shared->page_buffer[slotno], 0, BLCKSZ);
+		return true;
 	}
 
 	errno = 0;

src/backend/storage/smgr/smgr.c

@@ -106,6 +106,12 @@ static int NSmgr = 1;
 start_unlogged_build_hook_type start_unlogged_build_hook;
 finish_unlogged_build_phase_1_hook_type finish_unlogged_build_phase_1_hook;
 end_unlogged_build_hook_type end_unlogged_build_hook;
+
+/*
+ * SLRU download isn't really part of the smgr API, as SLRUs are not
+ * relations. But we define this here anyway, to keep it close to the smgr
+ * hooks in Neon.
+ */
 read_slru_segment_hook_type read_slru_segment_hook;
 
 /*
@@ -1089,23 +1095,22 @@ smgr_end_unlogged_build(SMgrRelation reln)
 }
 
 /*
- * NEON: we do not want to include large pg_xact/multixact files in basebackup and prefer
- * to download them on demand to reduce startup time.
- * If SLRU segment is not found, we try to download it from page server
+ * NEON: Attempt to download an SLRU file from remote storage.
  *
- * This function returns number of blocks in segment. Usually it should be SLRU_PAGES_PER_SEGMENT but in case
- * of partial segment, it can be smaller. Zero value means that segment doesn't exist.
- * From Postgres point of view empty segment is the same as absent segment.
+ * To reduce startup time, we don't want to include large pg_xact/multixact
+ * files in the basebackup. Instead, we have this hook to download them on
+ * demand. If an SLRU segment is not found, the code in slru.c calls this to
+ * check if it can be downloaded from the pageserver.
  *
- * This should really be a separate hook, not something that's in the smgr API, but
- * oh well.
+ * If the segment is found in remote storage, the hook writes it to the local
+ * file and returns 'true'. If the file is not found, returns 'false'.
  */
-int
-read_slru_segment(const char* path, int segno, void* buffer)
+bool
+smgr_read_slru_segment(const char *path, int segno)
 {
 	if (read_slru_segment_hook)
-		return read_slru_segment_hook(path, segno, buffer);
-	return 0;
+		return read_slru_segment_hook(path, segno);
+	return false;
 }
 
 /*

src/backend/utils/init/globals.c

@@ -119,7 +119,6 @@ pid_t		PostmasterPid = 0;
 bool		IsPostmasterEnvironment = false;
 bool		IsUnderPostmaster = false;
 bool		IsBinaryUpgrade = false;
-bool		neon_use_communicator_worker = false;
 
 bool		ExitOnAnyError = false;
 

src/include/miscadmin.h

@@ -170,12 +170,6 @@ extern PGDLLIMPORT pid_t PostmasterPid;
 extern PGDLLIMPORT bool IsPostmasterEnvironment;
 extern PGDLLIMPORT bool IsUnderPostmaster;
 extern PGDLLIMPORT bool IsBinaryUpgrade;
-/* Whether the communicator worker is used or not. Defined here so that
- * it is also accessible from the main postgres code easily without
- * having to look up the value using strings and chain of other
- * functions.
- */
-extern PGDLLIMPORT bool neon_use_communicator_worker;
 
 extern PGDLLIMPORT bool ExitOnAnyError;
 

src/include/storage/smgr.h

@@ -155,7 +155,7 @@ extern finish_unlogged_build_phase_1_hook_type finish_unlogged_build_phase_1_hoo
 extern end_unlogged_build_hook_type end_unlogged_build_hook;
 
 /* NEON: Hook for reading an SLRU segment from e.g. remote storage */
-typedef int		(*read_slru_segment_hook_type) (const char *path, int segno, void* buffer);
+typedef bool (*read_slru_segment_hook_type) (const char *path, int segno);
 extern read_slru_segment_hook_type read_slru_segment_hook;
 
 /* NEON: Alternative implementation of calculate_database_size(), to make it O(1) */
@@ -217,7 +217,7 @@ extern void	smgr_finish_unlogged_build_phase_1(SMgrRelation reln);
 extern void smgr_end_unlogged_build(SMgrRelation reln);
 
 /* Neon: Allow on-demand download of SLRU segment data */
-extern int  read_slru_segment(const char *path, int segno, void* buffer);
+extern bool smgr_read_slru_segment(const char *path, int segno);
 
 static inline void
 smgrread(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,