Logical replication timeout

diff --git a/src/backend/replication/logical/reorderbuffer.c b/src/backend/replication/logical/reorderbuffer.c
index 3a68a393d2..2f5077c0b6 100644
--- a/src/backend/replication/logical/reorderbuffer.c
+++ b/src/backend/replication/logical/reorderbuffer.c
@@ -4409,27 +4409,38 @@ ReorderBufferRestoreChange(ReorderBuffer *rb, ReorderBufferTXN *txn,
 static void
 ReorderBufferRestoreCleanup(ReorderBuffer *rb, ReorderBufferTXN *txn)
 {
-	XLogSegNo	first;
-	XLogSegNo	cur;
-	XLogSegNo	last;
+	DIR		   *spill_dir;
+	struct dirent *spill_de;
+	struct stat statbuf;
+	char		path[MAXPGPATH * 2 + sizeof(PG_REPLSLOT_DIR)];
+	char		file_filter[MAXPGPATH];
 
-	Assert(txn->first_lsn != InvalidXLogRecPtr);
-	Assert(txn->final_lsn != InvalidXLogRecPtr);
+	sprintf(path, "%s/%s", PG_REPLSLOT_DIR,
+			NameStr(MyReplicationSlot->data.name));
 
-	XLByteToSeg(txn->first_lsn, first, wal_segment_size);
-	XLByteToSeg(txn->final_lsn, last, wal_segment_size);
+	/* we're only handling directories here, skip if it's not ours */
+	if (lstat(path, &statbuf) == 0 && !S_ISDIR(statbuf.st_mode))
+		return;
+
+	snprintf(file_filter, MAXPGPATH, "xid-%u-", txn->xid);
 
-	/* iterate over all possible filenames, and delete them */
-	for (cur = first; cur <= last; cur++)
+	spill_dir = AllocateDir(path);
+	while ((spill_de = ReadDirExtended(spill_dir, path, INFO)) != NULL)
 	{
-		char		path[MAXPGPATH];
+		/* only look at names that can be ours */
+		if (strstr(spill_de->d_name, file_filter))
+		{
+			snprintf(path, sizeof(path),
+					 "%s/%s/%s", PG_REPLSLOT_DIR,
+					 NameStr(MyReplicationSlot->data.name), spill_de->d_name);
 
-		ReorderBufferSerializedPath(path, MyReplicationSlot, txn->xid, cur);
-		if (unlink(path) != 0 && errno != ENOENT)
-			ereport(ERROR,
-					(errcode_for_file_access(),
-					 errmsg("could not remove file \"%s\": %m", path)));
+			if (unlink(path) != 0)
+				ereport(ERROR,
+						(errcode_for_file_access(),
+						 errmsg("could not remove file \"%s\": %m", path)));
+		}
 	}
+	FreeDir(spill_dir);
 }
 
 /*

--- a/src/backend/replication/logical/reorderbuffer.c	2024-12-12 10:21:12.962202250 +0100
+++ b/src/backend/replication/logical/reorderbuffer.c	2024-12-12 11:40:24.217634527 +0100
@@ -4567,27 +4567,35 @@
 static void
 ReorderBufferRestoreCleanup(ReorderBuffer *rb, ReorderBufferTXN *txn)
 {
-	XLogSegNo	first;
-	XLogSegNo	cur;
-	XLogSegNo	last;
+	DIR *spill_dir;
+	struct dirent *spill_de;
+	char path[MAXPGPATH];
+	char spill_path[MAXPGPATH];
+	char file_filter[MAXPGPATH];
+	int len_filter;
+ 
+	snprintf(spill_path, MAXPGPATH, "pg_replslot/%s", 
+			NameStr(MyReplicationSlot->data.name));
 
-	Assert(txn->first_lsn != InvalidXLogRecPtr);
-	Assert(txn->final_lsn != InvalidXLogRecPtr);
-
-	XLByteToSeg(txn->first_lsn, first, wal_segment_size);
-	XLByteToSeg(txn->final_lsn, last, wal_segment_size);
-
-	/* iterate over all possible filenames, and delete them */
-	for (cur = first; cur <= last; cur++)
-	{
-		char		path[MAXPGPATH];
-
-		ReorderBufferSerializedPath(path, MyReplicationSlot, txn->xid, cur);
-		if (unlink(path) != 0 && errno != ENOENT)
-			ereport(ERROR,
-					(errcode_for_file_access(),
-					 errmsg("could not remove file \"%s\": %m", path)));
-	}
+ 	snprintf(file_filter, sizeof(file_filter), "xid-%u-", txn->xid);
+	len_filter = strlen(file_filter);
+ 
+	spill_dir = AllocateDir(spill_path);
+	while ((spill_de = ReadDirExtended(spill_dir, spill_path, INFO)) != NULL)
+ 	{
+		/* only look at names that can be ours */
+		if (strncmp(spill_de->d_name, file_filter, len_filter) == 0)
+		{
+			snprintf(path, sizeof(path), "%s/%s", spill_path,
+					spill_de->d_name);
+ 
+			if (unlink(path) != 0)
+				ereport(ERROR,
+						(errcode_for_file_access(),
+						 errmsg("could not remove file \"%s\": %m", path)));
+		}
+ 	}
+	FreeDir(spill_dir);
 }
 
 /*

diff --git a/src/backend/replication/logical/reorderbuffer.c b/src/backend/replication/logical/reorderbuffer.c
index e3a5c7b660..61943e72d8 100644
--- a/src/backend/replication/logical/reorderbuffer.c
+++ b/src/backend/replication/logical/reorderbuffer.c
@@ -149,7 +149,7 @@ typedef struct ReorderBufferIterTXNEntry
 	ReorderBufferChange *change;
 	ReorderBufferTXN *txn;
 	TXNEntryFile file;
-	XLogSegNo	segno;
+	int	restore_from;
 } ReorderBufferIterTXNEntry;
 
 typedef struct ReorderBufferIterTXNState
@@ -254,7 +254,7 @@ static void ReorderBufferSerializeTXN(ReorderBuffer *rb, ReorderBufferTXN *txn);
 static void ReorderBufferSerializeChange(ReorderBuffer *rb, ReorderBufferTXN *txn,
 										 int fd, ReorderBufferChange *change);
 static Size ReorderBufferRestoreChanges(ReorderBuffer *rb, ReorderBufferTXN *txn,
-										TXNEntryFile *file, XLogSegNo *segno);
+										TXNEntryFile *file, int *restore_from);
 static void ReorderBufferRestoreChange(ReorderBuffer *rb, ReorderBufferTXN *txn,
 									   char *data);
 static void ReorderBufferRestoreCleanup(ReorderBuffer *rb, ReorderBufferTXN *txn);
@@ -432,6 +432,7 @@ ReorderBufferGetTXN(ReorderBuffer *rb)
 	/* InvalidCommandId is not zero, so set it explicitly */
 	txn->command_id = InvalidCommandId;
 	txn->output_plugin_private = NULL;
+	txn->walsgmts = NIL;
 
 	return txn;
 }
@@ -1305,7 +1306,7 @@ ReorderBufferIterTXNInit(ReorderBuffer *rb, ReorderBufferTXN *txn,
 	for (off = 0; off < state->nr_txns; off++)
 	{
 		state->entries[off].file.vfd = -1;
-		state->entries[off].segno = 0;
+		state->entries[off].restore_from = 0;
 	}
 
 	/* allocate heap */
@@ -1333,7 +1334,7 @@ ReorderBufferIterTXNInit(ReorderBuffer *rb, ReorderBufferTXN *txn,
 			/* serialize remaining changes */
 			ReorderBufferSerializeTXN(rb, txn);
 			ReorderBufferRestoreChanges(rb, txn, &state->entries[off].file,
-										&state->entries[off].segno);
+										&state->entries[off].restore_from);
 		}
 
 		cur_change = dlist_head_element(ReorderBufferChange, node,
@@ -1363,7 +1364,7 @@ ReorderBufferIterTXNInit(ReorderBuffer *rb, ReorderBufferTXN *txn,
 				ReorderBufferSerializeTXN(rb, cur_txn);
 				ReorderBufferRestoreChanges(rb, cur_txn,
 											&state->entries[off].file,
-											&state->entries[off].segno);
+											&state->entries[off].restore_from);
 			}
 			cur_change = dlist_head_element(ReorderBufferChange, node,
 											&cur_txn->changes);
@@ -1448,7 +1449,7 @@ ReorderBufferIterTXNNext(ReorderBuffer *rb, ReorderBufferIterTXNState *state)
 		 */
 		rb->totalBytes += entry->txn->size;
 		if (ReorderBufferRestoreChanges(rb, entry->txn, &entry->file,
-										&state->entries[off].segno))
+										&state->entries[off].restore_from))
 		{
 			/* successfully restored changes from disk */
 			ReorderBufferChange *next_change =
@@ -3715,6 +3716,7 @@ ReorderBufferSerializeTXN(ReorderBuffer *rb, ReorderBufferTXN *txn)
 	XLogSegNo	curOpenSegNo = 0;
 	Size		spilled = 0;
 	Size		size = txn->size;
+	MemoryContext oldcontext;
 
 	elog(DEBUG2, "spill %u changes in XID %u to disk",
 		 (uint32) txn->nentries_mem, txn->xid);
@@ -3758,7 +3760,17 @@ ReorderBufferSerializeTXN(ReorderBuffer *rb, ReorderBufferTXN *txn)
 
 			/* open segment, create it if necessary */
 			fd = OpenTransientFile(path,
-								   O_CREAT | O_WRONLY | O_APPEND | PG_BINARY);
+								   O_CREAT | O_EXCL | O_WRONLY | O_APPEND | PG_BINARY);
+
+			if (fd < 0)
+				fd = OpenTransientFile(path,
+									   O_CREAT | O_WRONLY | O_APPEND | PG_BINARY);
+			else
+			{
+				oldcontext = MemoryContextSwitchTo(rb->context);
+				txn->walsgmts = lappend(txn->walsgmts, curOpenSegNo);
+				MemoryContextSwitchTo(oldcontext);
+			}
 
 			if (fd < 0)
 				ereport(ERROR,
@@ -4255,16 +4267,12 @@ ReorderBufferChangeSize(ReorderBufferChange *change)
  */
 static Size
 ReorderBufferRestoreChanges(ReorderBuffer *rb, ReorderBufferTXN *txn,
-							TXNEntryFile *file, XLogSegNo *segno)
+							TXNEntryFile *file,	int *restore_from)
 {
 	Size		restored = 0;
-	XLogSegNo	last_segno;
 	dlist_mutable_iter cleanup_iter;
 	File	   *fd = &file->vfd;
 
-	Assert(txn->first_lsn != InvalidXLogRecPtr);
-	Assert(txn->final_lsn != InvalidXLogRecPtr);
-
 	/* free current entries, so we have memory for more */
 	dlist_foreach_modify(cleanup_iter, &txn->changes)
 	{
@@ -4277,9 +4285,8 @@ ReorderBufferRestoreChanges(ReorderBuffer *rb, ReorderBufferTXN *txn,
 	txn->nentries_mem = 0;
 	Assert(dlist_is_empty(&txn->changes));
 
-	XLByteToSeg(txn->final_lsn, last_segno, wal_segment_size);
-
-	while (restored < max_changes_in_memory && *segno <= last_segno)
+	while (restored < max_changes_in_memory &&
+		   (*restore_from) < txn->walsgmts->length)
 	{
 		int			readBytes;
 		ReorderBufferDiskChange *ondisk;
@@ -4289,19 +4296,21 @@ ReorderBufferRestoreChanges(ReorderBuffer *rb, ReorderBufferTXN *txn,
 		if (*fd == -1)
 		{
 			char		path[MAXPGPATH];
+			ListCell *lc;
+			XLogSegNo segno;
 
-			/* first time in */
-			if (*segno == 0)
-				XLByteToSeg(txn->first_lsn, *segno, wal_segment_size);
+			/* next wal segment for the transaction */
+			lc = list_nth_cell(txn->walsgmts, *restore_from);
+			segno = lfirst(lc);
 
-			Assert(*segno != 0 || dlist_is_empty(&txn->changes));
+			Assert(segno != 0 || dlist_is_empty(&txn->changes));
 
 			/*
 			 * No need to care about TLIs here, only used during a single run,
 			 * so each LSN only maps to a specific WAL record.
 			 */
 			ReorderBufferSerializedPath(path, MyReplicationSlot, txn->xid,
-										*segno);
+										segno);
 
 			*fd = PathNameOpenFile(path, O_RDONLY | PG_BINARY);
 
@@ -4311,7 +4320,7 @@ ReorderBufferRestoreChanges(ReorderBuffer *rb, ReorderBufferTXN *txn,
 			if (*fd < 0 && errno == ENOENT)
 			{
 				*fd = -1;
-				(*segno)++;
+				(*restore_from)++;
 				continue;
 			}
 			else if (*fd < 0)
@@ -4336,7 +4345,7 @@ ReorderBufferRestoreChanges(ReorderBuffer *rb, ReorderBufferTXN *txn,
 		{
 			FileClose(*fd);
 			*fd = -1;
-			(*segno)++;
+			(*restore_from)++;
 			continue;
 		}
 		else if (readBytes < 0)
@@ -4567,27 +4576,26 @@ ReorderBufferRestoreChange(ReorderBuffer *rb, ReorderBufferTXN *txn,
 static void
 ReorderBufferRestoreCleanup(ReorderBuffer *rb, ReorderBufferTXN *txn)
 {
-	XLogSegNo	first;
-	XLogSegNo	cur;
-	XLogSegNo	last;
-
-	Assert(txn->first_lsn != InvalidXLogRecPtr);
-	Assert(txn->final_lsn != InvalidXLogRecPtr);
-
-	XLByteToSeg(txn->first_lsn, first, wal_segment_size);
-	XLByteToSeg(txn->final_lsn, last, wal_segment_size);
+	ListCell *cell;
 
 	/* iterate over all possible filenames, and delete them */
-	for (cur = first; cur <= last; cur++)
+	foreach(cell, txn->walsgmts)
 	{
+		XLogSegNo curr_segno = (XLogSegNo) lfirst(cell);
 		char		path[MAXPGPATH];
 
-		ReorderBufferSerializedPath(path, MyReplicationSlot, txn->xid, cur);
+		ReorderBufferSerializedPath(path, MyReplicationSlot, txn->xid, curr_segno);
 		if (unlink(path) != 0 && errno != ENOENT)
 			ereport(ERROR,
 					(errcode_for_file_access(),
 					 errmsg("could not remove file \"%s\": %m", path)));
 	}
+
+	if(txn->walsgmts != NIL)
+	{
+		pfree(txn->walsgmts);
+		txn->walsgmts = NIL;
+	}
 }
 
 /*
diff --git a/src/include/replication/reorderbuffer.h b/src/include/replication/reorderbuffer.h
index 3bc365a7b0..c23fa877dc 100644
--- a/src/include/replication/reorderbuffer.h
+++ b/src/include/replication/reorderbuffer.h
@@ -426,6 +426,11 @@ typedef struct ReorderBufferTXN
 	 * Private data pointer of the output plugin.
 	 */
 	void	   *output_plugin_private;
+
+	/*
+	 * List of wal segments this txn is part of
+	 */
+	List	*walsgmts;
 } ReorderBufferTXN;
 
 /* so we can define the callbacks used inside struct ReorderBuffer itself */

From 99ffc05998a157c2632711aaa8b5d7c4e22ebd70 Mon Sep 17 00:00:00 2001
From: Hayato Kuroda <kuroda.hayato@fujitsu.com>
Date: Wed, 25 Dec 2024 08:02:33 +0000
Subject: [PATCH] WIP: track wal segments

---
 .../replication/logical/reorderbuffer.c       | 84 +++++++++++--------
 src/include/replication/reorderbuffer.h       |  7 ++
 2 files changed, 58 insertions(+), 33 deletions(-)

diff --git a/src/backend/replication/logical/reorderbuffer.c b/src/backend/replication/logical/reorderbuffer.c
index 9c742e96eb..b2bb4423b8 100644
--- a/src/backend/replication/logical/reorderbuffer.c
+++ b/src/backend/replication/logical/reorderbuffer.c
@@ -149,7 +149,7 @@ typedef struct ReorderBufferIterTXNEntry
 	ReorderBufferChange *change;
 	ReorderBufferTXN *txn;
 	TXNEntryFile file;
-	XLogSegNo	segno;
+	int	restore_from;
 } ReorderBufferIterTXNEntry;
 
 typedef struct ReorderBufferIterTXNState
@@ -215,6 +215,11 @@ static const Size max_changes_in_memory = 4096; /* XXX for restore only */
 /* GUC variable */
 int			debug_logical_replication_streaming = DEBUG_LOGICAL_REP_STREAMING_BUFFERED;
 
+typedef struct WalSgmtsEntry
+{
+	XLogSegNo segno;
+} WalSgmtsEntry;
+
 /* ---------------------------------------
  * primary reorderbuffer support routines
  * ---------------------------------------
@@ -254,7 +259,7 @@ static void ReorderBufferSerializeTXN(ReorderBuffer *rb, ReorderBufferTXN *txn);
 static void ReorderBufferSerializeChange(ReorderBuffer *rb, ReorderBufferTXN *txn,
 										 int fd, ReorderBufferChange *change);
 static Size ReorderBufferRestoreChanges(ReorderBuffer *rb, ReorderBufferTXN *txn,
-										TXNEntryFile *file, XLogSegNo *segno);
+										TXNEntryFile *file, int *restore_from);
 static void ReorderBufferRestoreChange(ReorderBuffer *rb, ReorderBufferTXN *txn,
 									   char *data);
 static void ReorderBufferRestoreCleanup(ReorderBuffer *rb, ReorderBufferTXN *txn);
@@ -432,6 +437,7 @@ ReorderBufferGetTXN(ReorderBuffer *rb)
 	/* InvalidCommandId is not zero, so set it explicitly */
 	txn->command_id = InvalidCommandId;
 	txn->output_plugin_private = NULL;
+	txn->walsgmts = NIL;
 
 	return txn;
 }
@@ -1305,7 +1311,7 @@ ReorderBufferIterTXNInit(ReorderBuffer *rb, ReorderBufferTXN *txn,
 	for (off = 0; off < state->nr_txns; off++)
 	{
 		state->entries[off].file.vfd = -1;
-		state->entries[off].segno = 0;
+		state->entries[off].restore_from = 0;
 	}
 
 	/* allocate heap */
@@ -1333,7 +1339,7 @@ ReorderBufferIterTXNInit(ReorderBuffer *rb, ReorderBufferTXN *txn,
 			/* serialize remaining changes */
 			ReorderBufferSerializeTXN(rb, txn);
 			ReorderBufferRestoreChanges(rb, txn, &state->entries[off].file,
-										&state->entries[off].segno);
+										&state->entries[off].restore_from);
 		}
 
 		cur_change = dlist_head_element(ReorderBufferChange, node,
@@ -1363,7 +1369,7 @@ ReorderBufferIterTXNInit(ReorderBuffer *rb, ReorderBufferTXN *txn,
 				ReorderBufferSerializeTXN(rb, cur_txn);
 				ReorderBufferRestoreChanges(rb, cur_txn,
 											&state->entries[off].file,
-											&state->entries[off].segno);
+											&state->entries[off].restore_from);
 			}
 			cur_change = dlist_head_element(ReorderBufferChange, node,
 											&cur_txn->changes);
@@ -1448,7 +1454,7 @@ ReorderBufferIterTXNNext(ReorderBuffer *rb, ReorderBufferIterTXNState *state)
 		 */
 		rb->totalBytes += entry->txn->size;
 		if (ReorderBufferRestoreChanges(rb, entry->txn, &entry->file,
-										&state->entries[off].segno))
+										&state->entries[off].restore_from))
 		{
 			/* successfully restored changes from disk */
 			ReorderBufferChange *next_change =
@@ -3715,6 +3721,7 @@ ReorderBufferSerializeTXN(ReorderBuffer *rb, ReorderBufferTXN *txn)
 	XLogSegNo	curOpenSegNo = 0;
 	Size		spilled = 0;
 	Size		size = txn->size;
+	MemoryContext oldcontext;
 
 	elog(DEBUG2, "spill %u changes in XID %u to disk",
 		 (uint32) txn->nentries_mem, txn->xid);
@@ -3758,7 +3765,23 @@ ReorderBufferSerializeTXN(ReorderBuffer *rb, ReorderBufferTXN *txn)
 
 			/* open segment, create it if necessary */
 			fd = OpenTransientFile(path,
-								   O_CREAT | O_WRONLY | O_APPEND | PG_BINARY);
+								   O_CREAT | O_EXCL | O_WRONLY | O_APPEND | PG_BINARY);
+
+			if (fd < 0)
+				fd = OpenTransientFile(path,
+									   O_CREAT | O_WRONLY | O_APPEND | PG_BINARY);
+			else
+			{
+				WalSgmtsEntry *entry;
+
+				oldcontext = MemoryContextSwitchTo(rb->context);
+
+				entry = palloc(sizeof(WalSgmtsEntry));
+				entry->segno = curOpenSegNo;
+
+				txn->walsgmts = lappend(txn->walsgmts, entry);
+				MemoryContextSwitchTo(oldcontext);
+			}
 
 			if (fd < 0)
 				ereport(ERROR,
@@ -4255,16 +4278,12 @@ ReorderBufferChangeSize(ReorderBufferChange *change)
  */
 static Size
 ReorderBufferRestoreChanges(ReorderBuffer *rb, ReorderBufferTXN *txn,
-							TXNEntryFile *file, XLogSegNo *segno)
+							TXNEntryFile *file,	int *restore_from)
 {
 	Size		restored = 0;
-	XLogSegNo	last_segno;
 	dlist_mutable_iter cleanup_iter;
 	File	   *fd = &file->vfd;
 
-	Assert(txn->first_lsn != InvalidXLogRecPtr);
-	Assert(txn->final_lsn != InvalidXLogRecPtr);
-
 	/* free current entries, so we have memory for more */
 	dlist_foreach_modify(cleanup_iter, &txn->changes)
 	{
@@ -4277,9 +4296,8 @@ ReorderBufferRestoreChanges(ReorderBuffer *rb, ReorderBufferTXN *txn,
 	txn->nentries_mem = 0;
 	Assert(dlist_is_empty(&txn->changes));
 
-	XLByteToSeg(txn->final_lsn, last_segno, wal_segment_size);
-
-	while (restored < max_changes_in_memory && *segno <= last_segno)
+	while (restored < max_changes_in_memory &&
+		   (*restore_from) < txn->walsgmts->length)
 	{
 		int			readBytes;
 		ReorderBufferDiskChange *ondisk;
@@ -4289,19 +4307,23 @@ ReorderBufferRestoreChanges(ReorderBuffer *rb, ReorderBufferTXN *txn,
 		if (*fd == -1)
 		{
 			char		path[MAXPGPATH];
+			ListCell *lc;
+			WalSgmtsEntry *entry;
+			XLogSegNo segno;
 
-			/* first time in */
-			if (*segno == 0)
-				XLByteToSeg(txn->first_lsn, *segno, wal_segment_size);
+			/* Next wal segment for the transaction */
+			lc = list_nth_cell(txn->walsgmts, *restore_from);
+			entry = (WalSgmtsEntry *) lfirst(lc);
+			segno = entry->segno;
 
-			Assert(*segno != 0 || dlist_is_empty(&txn->changes));
+			Assert(segno != 0 || dlist_is_empty(&txn->changes));
 
 			/*
 			 * No need to care about TLIs here, only used during a single run,
 			 * so each LSN only maps to a specific WAL record.
 			 */
 			ReorderBufferSerializedPath(path, MyReplicationSlot, txn->xid,
-										*segno);
+										segno);
 
 			*fd = PathNameOpenFile(path, O_RDONLY | PG_BINARY);
 
@@ -4311,7 +4333,7 @@ ReorderBufferRestoreChanges(ReorderBuffer *rb, ReorderBufferTXN *txn,
 			if (*fd < 0 && errno == ENOENT)
 			{
 				*fd = -1;
-				(*segno)++;
+				(*restore_from)++;
 				continue;
 			}
 			else if (*fd < 0)
@@ -4336,7 +4358,7 @@ ReorderBufferRestoreChanges(ReorderBuffer *rb, ReorderBufferTXN *txn,
 		{
 			FileClose(*fd);
 			*fd = -1;
-			(*segno)++;
+			(*restore_from)++;
 			continue;
 		}
 		else if (readBytes < 0)
@@ -4567,26 +4589,22 @@ ReorderBufferRestoreChange(ReorderBuffer *rb, ReorderBufferTXN *txn,
 static void
 ReorderBufferRestoreCleanup(ReorderBuffer *rb, ReorderBufferTXN *txn)
 {
-	XLogSegNo	first;
-	XLogSegNo	cur;
-	XLogSegNo	last;
-
-	Assert(txn->first_lsn != InvalidXLogRecPtr);
-	Assert(txn->final_lsn != InvalidXLogRecPtr);
-
-	XLByteToSeg(txn->first_lsn, first, wal_segment_size);
-	XLByteToSeg(txn->final_lsn, last, wal_segment_size);
+	ListCell *cell;
 
 	/* iterate over all possible filenames, and delete them */
-	for (cur = first; cur <= last; cur++)
+	foreach(cell, txn->walsgmts)
 	{
+		WalSgmtsEntry *entry = (WalSgmtsEntry *)lfirst(cell);
+		XLogSegNo curr_segno = entry->segno;
 		char		path[MAXPGPATH];
 
-		ReorderBufferSerializedPath(path, MyReplicationSlot, txn->xid, cur);
+		ReorderBufferSerializedPath(path, MyReplicationSlot, txn->xid, curr_segno);
 		if (unlink(path) != 0 && errno != ENOENT)
 			ereport(ERROR,
 					(errcode_for_file_access(),
 					 errmsg("could not remove file \"%s\": %m", path)));
+
+		txn->walsgmts = foreach_delete_current(txn->walsgmts, cell);
 	}
 }
 
diff --git a/src/include/replication/reorderbuffer.h b/src/include/replication/reorderbuffer.h
index 7de50462dc..b2ef2640aa 100644
--- a/src/include/replication/reorderbuffer.h
+++ b/src/include/replication/reorderbuffer.h
@@ -422,6 +422,13 @@ typedef struct ReorderBufferTXN
 	 * Private data pointer of the output plugin.
 	 */
 	void	   *output_plugin_private;
+
+	/*
+	 * List of wal segments this txn is part of.
+	 *
+	 * XXX: check whether the attribute doesn't break ABI 
+	 */
+	List	*walsgmts;
 } ReorderBufferTXN;
 
 /* so we can define the callbacks used inside struct ReorderBuffer itself */
-- 
2.43.0

diff --git a/src/backend/replication/logical/reorderbuffer.c b/src/backend/replication/logical/reorderbuffer.c
index 9313d2b01e..de4a6f7baa 100644
--- a/src/backend/replication/logical/reorderbuffer.c
+++ b/src/backend/replication/logical/reorderbuffer.c
@@ -88,6 +88,7 @@
 
 #include <unistd.h>
 #include <sys/stat.h>
+#include <sys/time.h>
 
 #include "access/detoast.h"
 #include "access/heapam.h"
@@ -1515,6 +1516,9 @@ ReorderBufferCleanupTXN(ReorderBuffer *rb, ReorderBufferTXN *txn)
 	bool		found;
 	dlist_mutable_iter iter;
 	Size		mem_freed = 0;
+	ListCell *cell;
+	struct timeval stop, start;
+	gettimeofday(&start, NULL);
 
 	/* cleanup subtransactions & their changes */
 	dlist_foreach_modify(iter, &txn->subtxns)
@@ -1614,6 +1618,9 @@ ReorderBufferCleanupTXN(ReorderBuffer *rb, ReorderBufferTXN *txn)
 	if (rbtxn_is_serialized(txn))
 		ReorderBufferRestoreCleanup(rb, txn);
 
+	gettimeofday(&stop, NULL);
+	elog(LOG, "Time to Cleanup txn: %d is : %lu", txn->xid, (stop.tv_sec - start.tv_sec) * 1000000 + stop.tv_usec - start.tv_usec);
+
 	/* deallocate */
 	ReorderBufferReturnTXN(rb, txn);
 }

From 1f78955508dd31887f681d802c44b77eae921a30 Mon Sep 17 00:00:00 2001
From: Hayato Kuroda <kuroda.hayato@fujitsu.com>
Date: Wed, 25 Dec 2024 08:02:33 +0000
Subject: [PATCH v2] WIP: track wal segments

---
 .../replication/logical/reorderbuffer.c       | 84 +++++++++++--------
 src/include/replication/reorderbuffer.h       |  7 ++
 2 files changed, 58 insertions(+), 33 deletions(-)

diff --git a/src/backend/replication/logical/reorderbuffer.c b/src/backend/replication/logical/reorderbuffer.c
index 5186ad2a39..5775b5022c 100644
--- a/src/backend/replication/logical/reorderbuffer.c
+++ b/src/backend/replication/logical/reorderbuffer.c
@@ -150,7 +150,7 @@ typedef struct ReorderBufferIterTXNEntry
 	ReorderBufferChange *change;
 	ReorderBufferTXN *txn;
 	TXNEntryFile file;
-	XLogSegNo	segno;
+	int	restore_from;
 } ReorderBufferIterTXNEntry;
 
 typedef struct ReorderBufferIterTXNState
@@ -216,6 +216,11 @@ static const Size max_changes_in_memory = 4096; /* XXX for restore only */
 /* GUC variable */
 int			debug_logical_replication_streaming = DEBUG_LOGICAL_REP_STREAMING_BUFFERED;
 
+typedef struct WalSgmtsEntry
+{
+	XLogSegNo segno;
+} WalSgmtsEntry;
+
 /* ---------------------------------------
  * primary reorderbuffer support routines
  * ---------------------------------------
@@ -255,7 +260,7 @@ static void ReorderBufferSerializeTXN(ReorderBuffer *rb, ReorderBufferTXN *txn);
 static void ReorderBufferSerializeChange(ReorderBuffer *rb, ReorderBufferTXN *txn,
 										 int fd, ReorderBufferChange *change);
 static Size ReorderBufferRestoreChanges(ReorderBuffer *rb, ReorderBufferTXN *txn,
-										TXNEntryFile *file, XLogSegNo *segno);
+										TXNEntryFile *file, int *restore_from);
 static void ReorderBufferRestoreChange(ReorderBuffer *rb, ReorderBufferTXN *txn,
 									   char *data);
 static void ReorderBufferRestoreCleanup(ReorderBuffer *rb, ReorderBufferTXN *txn);
@@ -435,6 +440,7 @@ ReorderBufferGetTXN(ReorderBuffer *rb)
 	/* InvalidCommandId is not zero, so set it explicitly */
 	txn->command_id = InvalidCommandId;
 	txn->output_plugin_private = NULL;
+	txn->walsgmts = NIL;
 
 	return txn;
 }
@@ -1308,7 +1314,7 @@ ReorderBufferIterTXNInit(ReorderBuffer *rb, ReorderBufferTXN *txn,
 	for (off = 0; off < state->nr_txns; off++)
 	{
 		state->entries[off].file.vfd = -1;
-		state->entries[off].segno = 0;
+		state->entries[off].restore_from = 0;
 	}
 
 	/* allocate heap */
@@ -1336,7 +1342,7 @@ ReorderBufferIterTXNInit(ReorderBuffer *rb, ReorderBufferTXN *txn,
 			/* serialize remaining changes */
 			ReorderBufferSerializeTXN(rb, txn);
 			ReorderBufferRestoreChanges(rb, txn, &state->entries[off].file,
-										&state->entries[off].segno);
+										&state->entries[off].restore_from);
 		}
 
 		cur_change = dlist_head_element(ReorderBufferChange, node,
@@ -1366,7 +1372,7 @@ ReorderBufferIterTXNInit(ReorderBuffer *rb, ReorderBufferTXN *txn,
 				ReorderBufferSerializeTXN(rb, cur_txn);
 				ReorderBufferRestoreChanges(rb, cur_txn,
 											&state->entries[off].file,
-											&state->entries[off].segno);
+											&state->entries[off].restore_from);
 			}
 			cur_change = dlist_head_element(ReorderBufferChange, node,
 											&cur_txn->changes);
@@ -1451,7 +1457,7 @@ ReorderBufferIterTXNNext(ReorderBuffer *rb, ReorderBufferIterTXNState *state)
 		 */
 		rb->totalBytes += entry->txn->size;
 		if (ReorderBufferRestoreChanges(rb, entry->txn, &entry->file,
-										&state->entries[off].segno))
+										&state->entries[off].restore_from))
 		{
 			/* successfully restored changes from disk */
 			ReorderBufferChange *next_change =
@@ -3838,6 +3844,7 @@ ReorderBufferSerializeTXN(ReorderBuffer *rb, ReorderBufferTXN *txn)
 	XLogSegNo	curOpenSegNo = 0;
 	Size		spilled = 0;
 	Size		size = txn->size;
+	MemoryContext oldcontext;
 
 	elog(DEBUG2, "spill %u changes in XID %u to disk",
 		 (uint32) txn->nentries_mem, txn->xid);
@@ -3881,7 +3888,23 @@ ReorderBufferSerializeTXN(ReorderBuffer *rb, ReorderBufferTXN *txn)
 
 			/* open segment, create it if necessary */
 			fd = OpenTransientFile(path,
-								   O_CREAT | O_WRONLY | O_APPEND | PG_BINARY);
+								   O_CREAT | O_EXCL | O_WRONLY | O_APPEND | PG_BINARY);
+
+			if (fd < 0)
+				fd = OpenTransientFile(path,
+									   O_CREAT | O_WRONLY | O_APPEND | PG_BINARY);
+			else
+			{
+				WalSgmtsEntry *entry;
+
+				oldcontext = MemoryContextSwitchTo(rb->context);
+
+				entry = palloc(sizeof(WalSgmtsEntry));
+				entry->segno = curOpenSegNo;
+
+				txn->walsgmts = lappend(txn->walsgmts, entry);
+				MemoryContextSwitchTo(oldcontext);
+			}
 
 			if (fd < 0)
 				ereport(ERROR,
@@ -4378,16 +4401,12 @@ ReorderBufferChangeSize(ReorderBufferChange *change)
  */
 static Size
 ReorderBufferRestoreChanges(ReorderBuffer *rb, ReorderBufferTXN *txn,
-							TXNEntryFile *file, XLogSegNo *segno)
+							TXNEntryFile *file,	int *restore_from)
 {
 	Size		restored = 0;
-	XLogSegNo	last_segno;
 	dlist_mutable_iter cleanup_iter;
 	File	   *fd = &file->vfd;
 
-	Assert(txn->first_lsn != InvalidXLogRecPtr);
-	Assert(txn->final_lsn != InvalidXLogRecPtr);
-
 	/* free current entries, so we have memory for more */
 	dlist_foreach_modify(cleanup_iter, &txn->changes)
 	{
@@ -4400,9 +4419,8 @@ ReorderBufferRestoreChanges(ReorderBuffer *rb, ReorderBufferTXN *txn,
 	txn->nentries_mem = 0;
 	Assert(dlist_is_empty(&txn->changes));
 
-	XLByteToSeg(txn->final_lsn, last_segno, wal_segment_size);
-
-	while (restored < max_changes_in_memory && *segno <= last_segno)
+	while (restored < max_changes_in_memory &&
+		   (*restore_from) < txn->walsgmts->length)
 	{
 		int			readBytes;
 		ReorderBufferDiskChange *ondisk;
@@ -4412,19 +4430,23 @@ ReorderBufferRestoreChanges(ReorderBuffer *rb, ReorderBufferTXN *txn,
 		if (*fd == -1)
 		{
 			char		path[MAXPGPATH];
+			ListCell *lc;
+			WalSgmtsEntry *entry;
+			XLogSegNo segno;
 
-			/* first time in */
-			if (*segno == 0)
-				XLByteToSeg(txn->first_lsn, *segno, wal_segment_size);
+			/* Next wal segment for the transaction */
+			lc = list_nth_cell(txn->walsgmts, *restore_from);
+			entry = (WalSgmtsEntry *) lfirst(lc);
+			segno = entry->segno;
 
-			Assert(*segno != 0 || dlist_is_empty(&txn->changes));
+			Assert(segno != 0 || dlist_is_empty(&txn->changes));
 
 			/*
 			 * No need to care about TLIs here, only used during a single run,
 			 * so each LSN only maps to a specific WAL record.
 			 */
 			ReorderBufferSerializedPath(path, MyReplicationSlot, txn->xid,
-										*segno);
+										segno);
 
 			*fd = PathNameOpenFile(path, O_RDONLY | PG_BINARY);
 
@@ -4434,7 +4456,7 @@ ReorderBufferRestoreChanges(ReorderBuffer *rb, ReorderBufferTXN *txn,
 			if (*fd < 0 && errno == ENOENT)
 			{
 				*fd = -1;
-				(*segno)++;
+				(*restore_from)++;
 				continue;
 			}
 			else if (*fd < 0)
@@ -4459,7 +4481,7 @@ ReorderBufferRestoreChanges(ReorderBuffer *rb, ReorderBufferTXN *txn,
 		{
 			FileClose(*fd);
 			*fd = -1;
-			(*segno)++;
+			(*restore_from)++;
 			continue;
 		}
 		else if (readBytes < 0)
@@ -4690,26 +4712,22 @@ ReorderBufferRestoreChange(ReorderBuffer *rb, ReorderBufferTXN *txn,
 static void
 ReorderBufferRestoreCleanup(ReorderBuffer *rb, ReorderBufferTXN *txn)
 {
-	XLogSegNo	first;
-	XLogSegNo	cur;
-	XLogSegNo	last;
-
-	Assert(txn->first_lsn != InvalidXLogRecPtr);
-	Assert(txn->final_lsn != InvalidXLogRecPtr);
-
-	XLByteToSeg(txn->first_lsn, first, wal_segment_size);
-	XLByteToSeg(txn->final_lsn, last, wal_segment_size);
+	ListCell *cell;
 
 	/* iterate over all possible filenames, and delete them */
-	for (cur = first; cur <= last; cur++)
+	foreach(cell, txn->walsgmts)
 	{
+		WalSgmtsEntry *entry = (WalSgmtsEntry *)lfirst(cell);
+		XLogSegNo curr_segno = entry->segno;
 		char		path[MAXPGPATH];
 
-		ReorderBufferSerializedPath(path, MyReplicationSlot, txn->xid, cur);
+		ReorderBufferSerializedPath(path, MyReplicationSlot, txn->xid, curr_segno);
 		if (unlink(path) != 0 && errno != ENOENT)
 			ereport(ERROR,
 					(errcode_for_file_access(),
 					 errmsg("could not remove file \"%s\": %m", path)));
+
+		txn->walsgmts = foreach_delete_current(txn->walsgmts, cell);
 	}
 }
 
diff --git a/src/include/replication/reorderbuffer.h b/src/include/replication/reorderbuffer.h
index 517a8e3634..fa72b040fb 100644
--- a/src/include/replication/reorderbuffer.h
+++ b/src/include/replication/reorderbuffer.h
@@ -452,6 +452,13 @@ typedef struct ReorderBufferTXN
 	 * Private data pointer of the output plugin.
 	 */
 	void	   *output_plugin_private;
+
+	/*
+	 * List of wal segments this txn is part of.
+	 *
+	 * XXX: check whether the attribute doesn't break ABI.
+	 */
+	List	*walsgmts;
 } ReorderBufferTXN;
 
 /* so we can define the callbacks used inside struct ReorderBuffer itself */
-- 
2.43.5