From 9e2d06739724554ad3cf7759c519449c8f9e2633 Mon Sep 17 00:00:00 2001
From: Toni Spets <toni.spets@beeper.com>
Date: Fri, 28 Apr 2023 09:28:03 +0300
Subject: [PATCH] Copy WAL frames through temp file to shadow

This implementation replicates what the new behavior in the legacy
branch did using a temporary file instead of memory to buffer WAL
frames between commits.

Fixes excess memory use if the WAL includes massive commits due to
big inserts or migrations.
---
 db.go | 51 +++++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 41 insertions(+), 10 deletions(-)

diff --git a/db.go b/db.go
index 773315ca..756b73db 100644
--- a/db.go
+++ b/db.go
@@ -1033,6 +1033,16 @@ func (db *DB) copyToShadowWAL(filename string) (origWalSize int64, newSize int64
 		return 0, 0, fmt.Errorf("last checksum: %w", err)
 	}
 
+	// Write to a temporary shadow file.
+	tempFilename := filename + ".tmp"
+	defer os.Remove(tempFilename)
+
+	f, err := internal.CreateFile(tempFilename, db.fileInfo)
+	if err != nil {
+		return 0, 0, fmt.Errorf("create temp file: %w", err)
+	}
+	defer f.Close()
+
 	// Seek to correct position on real wal.
 	if _, err := r.Seek(origSize, io.SeekStart); err != nil {
 		return 0, 0, fmt.Errorf("real wal seek: %w", err)
@@ -1043,7 +1053,6 @@ func (db *DB) copyToShadowWAL(filename string) (origWalSize int64, newSize int64
 	// Read through WAL from last position to find the page of the last
 	// committed transaction.
 	frame := make([]byte, db.pageSize+WALFrameHeaderSize)
-	var buf bytes.Buffer
 	offset := origSize
 	lastCommitSize := origSize
 	for {
@@ -1073,24 +1082,46 @@ func (db *DB) copyToShadowWAL(filename string) (origWalSize int64, newSize int64
 			break
 		}
 
-		// Add page to the new size of the shadow WAL.
-		buf.Write(frame)
+		// Write page to temporary WAL file.
+		if _, err := f.Write(frame); err != nil {
+			return 0, 0, fmt.Errorf("write temp shadow wal: %w", err)
+		}
 
 		Tracef("%s: copy-shadow: ok %s offset=%d salt=%x %x", db.path, filename, offset, salt0, salt1)
 		offset += int64(len(frame))
 
-		// Flush to shadow WAL if commit record.
+		// Update new size if written frame was a commit record.
 		newDBSize := binary.BigEndian.Uint32(frame[4:])
 		if newDBSize != 0 {
-			if _, err := buf.WriteTo(w); err != nil {
-				return 0, 0, fmt.Errorf("write shadow wal: %w", err)
-			}
-			buf.Reset()
 			lastCommitSize = offset
 		}
 	}
 
-	// Sync & close.
+	// If no WAL writes found, exit.
+	if origSize == lastCommitSize {
+		return origWalSize, origSize, nil
+	}
+
+	walByteN := lastCommitSize - origSize
+
+	// Move to beginning of temporary file.
+	if _, err := f.Seek(0, io.SeekStart); err != nil {
+		return 0, 0, fmt.Errorf("temp file seek: %w", err)
+	}
+
+	// Copy from temporary file to shadow WAL.
+	if _, err := io.Copy(w, &io.LimitedReader{R: f, N: walByteN}); err != nil {
+		return 0, 0, fmt.Errorf("write shadow file: %w", err)
+	}
+
+	// Close & remove temporary file.
+	if err := f.Close(); err != nil {
+		return 0, 0, err
+	} else if err := os.Remove(tempFilename); err != nil {
+		return 0, 0, err
+	}
+
+	// Sync & close shadow WAL.
 	if err := w.Sync(); err != nil {
 		return 0, 0, err
 	} else if err := w.Close(); err != nil {
@@ -1098,7 +1129,7 @@ func (db *DB) copyToShadowWAL(filename string) (origWalSize int64, newSize int64
 	}
 
 	// Track total number of bytes written to WAL.
-	db.totalWALBytesCounter.Add(float64(lastCommitSize - origSize))
+	db.totalWALBytesCounter.Add(float64(walByteN))
 
 	return origWalSize, lastCommitSize, nil
 }