From 1953fc7a93967211c467e8fd61c143d93fdb56a0 Mon Sep 17 00:00:00 2001 From: Jeremy Schneider Date: Mon, 6 Jul 2026 08:40:06 -0700 Subject: [PATCH 1/2] fix(restorer): scope end-of-WAL markers by WAL timeline Add timeline-aware end-of-WAL marker handling to `pkg/restorer.WALRestorer`. The previous global `end-of-wal-stream` marker could be set by a prefetch miss on one timeline and consumed by a later restore request on another timeline. In CNPG restore flows, that can make PostgreSQL skip a valid child-timeline WAL lookup and replay wrong parent-timeline bytes. Signed-off-by: Jeremy Schneider --- pkg/restorer/restorer.go | 124 ++++++++++++++++++++++++++++++-- pkg/restorer/restorer_test.go | 129 ++++++++++++++++++++++++++++++++++ 2 files changed, 246 insertions(+), 7 deletions(-) diff --git a/pkg/restorer/restorer.go b/pkg/restorer/restorer.go index 46f17a7b..4d038609 100644 --- a/pkg/restorer/restorer.go +++ b/pkg/restorer/restorer.go @@ -25,6 +25,9 @@ import ( "fmt" "math" "os/exec" + "path" + "regexp" + "strconv" "sync" "time" @@ -37,8 +40,11 @@ import ( const ( endOfWALStreamFlagFilename = "end-of-wal-stream" + endOfWALStreamFlagPrefix = endOfWALStreamFlagFilename + "-" ) +var walSegmentRe = regexp.MustCompile(`^([0-9A-F]{8})([0-9A-F]{8})([0-9A-F]{8})$`) + // ErrWALNotFound is returned when the WAL is not found in the cloud archive var ErrWALNotFound = errors.New("WAL not found") @@ -123,9 +129,45 @@ func (restorer *WALRestorer) RestoreFromSpool(walName, destinationPath string) ( } } -// SetEndOfWALStream add end-of-wal-stream in the spool directory +// SetEndOfWALStream add end-of-wal-stream in the spool directory. +// +// Deprecated: use SetEndOfWALStreamFromResults to avoid sharing the same marker +// across WAL timelines. func (restorer *WALRestorer) SetEndOfWALStream() error { - contains, err := restorer.IsEndOfWALStream() + return restorer.setEndOfWALStream(endOfWALStreamFlagFilename) +} + +// SetEndOfWALStreamFromResults sets timeline-scoped end-of-wal-stream markers +// for missing regular WAL segment files in the passed restore results. It +// returns true if any marker was set or already present. +func (restorer *WALRestorer) SetEndOfWALStreamFromResults(results []Result) (markerSet bool, err error) { + timelines := make(map[uint32]struct{}) + for _, result := range results { + if !errors.Is(result.Err, ErrWALNotFound) { + continue + } + + timeline, err := timelineFromWALName(result.WalName) + if errors.Is(err, ErrInvalidWALName) { + continue + } + if err != nil { + return false, err + } + timelines[timeline] = struct{}{} + } + + for timeline := range timelines { + if err := restorer.setEndOfWALStream(endOfWALStreamFlagName(timeline)); err != nil { + return false, err + } + } + + return len(timelines) > 0, nil +} + +func (restorer *WALRestorer) setEndOfWALStream(flagName string) error { + contains, err := restorer.isEndOfWALStream(flagName) if err != nil { return err } @@ -134,7 +176,7 @@ func (restorer *WALRestorer) SetEndOfWALStream() error { return nil } - err = restorer.spool.Touch(endOfWALStreamFlagFilename) + err = restorer.spool.Touch(flagName) if err != nil { return err } @@ -142,9 +184,52 @@ func (restorer *WALRestorer) SetEndOfWALStream() error { return nil } -// IsEndOfWALStream check whether end-of-wal-stream flag is presents in the spool directory +// IsEndOfWALStream check whether end-of-wal-stream flag is presents in the spool directory. +// +// Deprecated: use ConsumeEndOfWALStreamForWAL to avoid sharing the same marker +// across WAL timelines. func (restorer *WALRestorer) IsEndOfWALStream() (bool, error) { - isEOS, err := restorer.spool.Contains(endOfWALStreamFlagFilename) + return restorer.isEndOfWALStream(endOfWALStreamFlagFilename) +} + +// ConsumeEndOfWALStreamForWAL removes and returns whether an end-of-wal-stream +// marker is present for the requested WAL segment's timeline. Non-WAL names do +// not consume markers. +func (restorer *WALRestorer) ConsumeEndOfWALStreamForWAL(walName string) (markerConsumed bool, err error) { + timeline, err := timelineFromWALName(walName) + if errors.Is(err, ErrInvalidWALName) { + return false, nil + } + if err != nil { + return false, err + } + + containsLegacy, err := restorer.IsEndOfWALStream() + if err != nil { + return false, err + } + if containsLegacy { + if err := restorer.ResetEndOfWalStream(); err != nil { + return false, err + } + } + + markerConsumed, err = restorer.isEndOfWALStream(endOfWALStreamFlagName(timeline)) + if err != nil { + return false, err + } + if !markerConsumed { + return false, nil + } + + if err := restorer.resetEndOfWalStream(endOfWALStreamFlagName(timeline)); err != nil { + return false, err + } + return true, nil +} + +func (restorer *WALRestorer) isEndOfWALStream(flagName string) (bool, error) { + isEOS, err := restorer.spool.Contains(flagName) if err != nil { return false, fmt.Errorf("failed to check end-of-wal-stream flag: %w", err) } @@ -152,9 +237,16 @@ func (restorer *WALRestorer) IsEndOfWALStream() (bool, error) { return isEOS, nil } -// ResetEndOfWalStream remove end-of-wal-stream flag from the spool directory +// ResetEndOfWalStream remove end-of-wal-stream flag from the spool directory. +// +// Deprecated: use ConsumeEndOfWALStreamForWAL to avoid sharing the same +// marker across WAL timelines. func (restorer *WALRestorer) ResetEndOfWalStream() error { - err := restorer.spool.Remove(endOfWALStreamFlagFilename) + return restorer.resetEndOfWalStream(endOfWALStreamFlagFilename) +} + +func (restorer *WALRestorer) resetEndOfWalStream(flagName string) error { + err := restorer.spool.Remove(flagName) if err != nil { return fmt.Errorf("failed to remove end-of-wal-stream flag: %w", err) } @@ -162,6 +254,24 @@ func (restorer *WALRestorer) ResetEndOfWalStream() error { return nil } +func endOfWALStreamFlagName(timeline uint32) string { + return fmt.Sprintf("%s%08X", endOfWALStreamFlagPrefix, timeline) +} + +func timelineFromWALName(walName string) (uint32, error) { + subMatches := walSegmentRe.FindStringSubmatch(path.Base(walName)) + if len(subMatches) != 4 { + return 0, ErrInvalidWALName + } + + timeline, err := strconv.ParseUint(subMatches[1], 16, 32) + if err != nil { + return 0, ErrInvalidWALName + } + + return uint32(timeline), nil +} + // RestoreList restores a list of WALs. The first WAL of the list will go directly into the // destination path, the others will be adopted by the spool func (restorer *WALRestorer) RestoreList( diff --git a/pkg/restorer/restorer_test.go b/pkg/restorer/restorer_test.go index 32effad0..90befb84 100644 --- a/pkg/restorer/restorer_test.go +++ b/pkg/restorer/restorer_test.go @@ -20,13 +20,142 @@ SPDX-License-Identifier: Apache-2.0 package restorer import ( + "context" "errors" "fmt" + "os" + "path/filepath" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" ) +var _ = Describe("SetEndOfWALStreamFromResults", func() { + var ( + spoolDirectory string + restorer *WALRestorer + ) + + BeforeEach(func() { + spoolDirectory = GinkgoT().TempDir() + var err error + restorer, err = New(context.Background(), nil, spoolDirectory) + Expect(err).ToNot(HaveOccurred()) + }) + + It("sets a timeline-scoped marker for missing regular WAL restore results", func() { + contains, err := restorer.SetEndOfWALStreamFromResults([]Result{ + {WalName: "000000010000000000000002", Err: ErrWALNotFound}, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(contains).To(BeTrue()) + + Expect(filepath.Join(spoolDirectory, "end-of-wal-stream-00000001")).To(BeAnExistingFile()) + }) + + It("is idempotent when setting a timeline marker", func() { + for range 2 { + contains, err := restorer.SetEndOfWALStreamFromResults([]Result{ + {WalName: "000000010000000000000002", Err: ErrWALNotFound}, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(contains).To(BeTrue()) + } + + entries, err := os.ReadDir(spoolDirectory) + Expect(err).ToNot(HaveOccurred()) + Expect(entries).To(HaveLen(1)) + Expect(entries[0].Name()).To(Equal("end-of-wal-stream-00000001")) + }) + + It("sets markers from missing regular WAL restore results only", func() { + contains, err := restorer.SetEndOfWALStreamFromResults([]Result{ + {WalName: "000000010000000000000001"}, + {WalName: "000000010000000000000002", Err: ErrWALNotFound}, + {WalName: "00000002.history", Err: ErrWALNotFound}, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(contains).To(BeTrue()) + + Expect(filepath.Join(spoolDirectory, "end-of-wal-stream-00000001")).To(BeAnExistingFile()) + entries, err := os.ReadDir(spoolDirectory) + Expect(err).ToNot(HaveOccurred()) + Expect(entries).To(HaveLen(1)) + }) + + It("does not set markers for missing non-WAL restore results", func() { + contains, err := restorer.SetEndOfWALStreamFromResults([]Result{ + {WalName: "00000002.history", Err: ErrWALNotFound}, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(contains).To(BeFalse()) + + entries, err := os.ReadDir(spoolDirectory) + Expect(err).ToNot(HaveOccurred()) + Expect(entries).To(BeEmpty()) + }) +}) + +var _ = Describe("ConsumeEndOfWALStreamForWAL", func() { + var ( + spoolDirectory string + restorer *WALRestorer + ) + + BeforeEach(func() { + spoolDirectory = GinkgoT().TempDir() + var err error + restorer, err = New(context.Background(), nil, spoolDirectory) + Expect(err).ToNot(HaveOccurred()) + }) + + It("consumes only the requested WAL timeline marker", func() { + contains, err := restorer.SetEndOfWALStreamFromResults([]Result{ + {WalName: "000000010000000000000002", Err: ErrWALNotFound}, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(contains).To(BeTrue()) + + contains, err = restorer.ConsumeEndOfWALStreamForWAL("000000020000000000000003") + Expect(err).ToNot(HaveOccurred()) + Expect(contains).To(BeFalse()) + Expect(filepath.Join(spoolDirectory, "end-of-wal-stream-00000001")).To(BeAnExistingFile()) + + contains, err = restorer.ConsumeEndOfWALStreamForWAL("000000010000000000000003") + Expect(err).ToNot(HaveOccurred()) + Expect(contains).To(BeTrue()) + Expect(filepath.Join(spoolDirectory, "end-of-wal-stream-00000001")).ToNot(BeAnExistingFile()) + }) + + It("removes legacy global markers without treating them as authoritative", func() { + Expect(restorer.SetEndOfWALStream()).To(Succeed()) + Expect(filepath.Join(spoolDirectory, "end-of-wal-stream")).To(BeAnExistingFile()) + + contains, err := restorer.ConsumeEndOfWALStreamForWAL("000000010000000000000003") + Expect(err).ToNot(HaveOccurred()) + Expect(contains).To(BeFalse()) + Expect(filepath.Join(spoolDirectory, "end-of-wal-stream")).ToNot(BeAnExistingFile()) + }) + + It("consumes markers for WAL names but not non-WAL names", func() { + contains, err := restorer.SetEndOfWALStreamFromResults([]Result{ + {WalName: "000000010000000000000002", Err: ErrWALNotFound}, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(contains).To(BeTrue()) + + contains, err = restorer.ConsumeEndOfWALStreamForWAL("00000002.history") + Expect(err).ToNot(HaveOccurred()) + Expect(contains).To(BeFalse()) + Expect(filepath.Join(spoolDirectory, "end-of-wal-stream-00000001")).To(BeAnExistingFile()) + + contains, err = restorer.ConsumeEndOfWALStreamForWAL("000000010000000000000003") + Expect(err).ToNot(HaveOccurred()) + Expect(contains).To(BeTrue()) + Expect(filepath.Join(spoolDirectory, "end-of-wal-stream-00000001")).ToNot(BeAnExistingFile()) + }) +}) + var _ = Describe("errorForExitCode", func() { const walName = "000000010000000000000001" From a40d41ab4e97dea825fc6fe4e089fcc01e784235 Mon Sep 17 00:00:00 2001 From: Armando Ruocco Date: Wed, 8 Jul 2026 15:02:29 +0200 Subject: [PATCH 2/2] refactor(restorer): simplify end-of-WAL marker internals Route the legacy marker cleanup in ConsumeEndOfWALStreamForWAL through the private marker helpers instead of the deprecated exported methods, so the deprecated API can later be removed without touching this path. timelineFromWALName can only fail with ErrInvalidWALName, so collapse the callers' two-step error handling into a single check and document that contract on the function itself. Signed-off-by: Armando Ruocco --- pkg/restorer/restorer.go | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/pkg/restorer/restorer.go b/pkg/restorer/restorer.go index 4d038609..f84fbc57 100644 --- a/pkg/restorer/restorer.go +++ b/pkg/restorer/restorer.go @@ -148,11 +148,8 @@ func (restorer *WALRestorer) SetEndOfWALStreamFromResults(results []Result) (mar } timeline, err := timelineFromWALName(result.WalName) - if errors.Is(err, ErrInvalidWALName) { - continue - } if err != nil { - return false, err + continue } timelines[timeline] = struct{}{} } @@ -197,19 +194,16 @@ func (restorer *WALRestorer) IsEndOfWALStream() (bool, error) { // not consume markers. func (restorer *WALRestorer) ConsumeEndOfWALStreamForWAL(walName string) (markerConsumed bool, err error) { timeline, err := timelineFromWALName(walName) - if errors.Is(err, ErrInvalidWALName) { - return false, nil - } if err != nil { - return false, err + return false, nil } - containsLegacy, err := restorer.IsEndOfWALStream() + containsLegacy, err := restorer.isEndOfWALStream(endOfWALStreamFlagFilename) if err != nil { return false, err } if containsLegacy { - if err := restorer.ResetEndOfWalStream(); err != nil { + if err := restorer.resetEndOfWalStream(endOfWALStreamFlagFilename); err != nil { return false, err } } @@ -258,6 +252,9 @@ func endOfWALStreamFlagName(timeline uint32) string { return fmt.Sprintf("%s%08X", endOfWALStreamFlagPrefix, timeline) } +// timelineFromWALName extracts the timeline from a regular WAL segment file +// name. It fails with ErrInvalidWALName for any other name (e.g. timeline +// history or backup label files); no other error is possible. func timelineFromWALName(walName string) (uint32, error) { subMatches := walSegmentRe.FindStringSubmatch(path.Base(walName)) if len(subMatches) != 4 {