amadan

carlos / platform Public

dbdoor: resume a download that drops part way

Downloading each file whole lost litestream's resume: a connection that
reset mid-file deleted the bytes kept so far, and litestream reopened
from the requested offset, so a large LTX file over a flaky link never
finished. download now keeps what it has and reopens from offset plus
bytes written, for the remaining size. Reopens that deliver nothing are
bounded (five in a row); any progress resets the count. A file that
vanishes on a reopen still marks the restore for its one fresh plan.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
claude 838ea5c3038a9e3c53c2ada8f5b5dcf71f6db2a2 parent f300cd2
2 files changed, +120 −13
  • internal/dbdoor/restore.go +55 −13
  • internal/dbdoor/restore_test.go +65 −0
diff --git a/internal/dbdoor/restore.go b/internal/dbdoor/restore.go
index 824d9859b..dd47e4621 100644
--- a/internal/dbdoor/restore.go
+++ b/internal/dbdoor/restore.go
@@ -138,7 +138,7 @@ func (r *recorder) OpenLTXFile(ctx context.Context, level int, minTXID, maxTXID
}
return nil, err
}
- if rc, err = r.download(rc); err != nil {
+ if rc, err = r.download(ctx, rc, level, minTXID, maxTXID, offset, size); err != nil {
return nil, err
}
if offset != 0 {
@@ -154,34 +154,76 @@ func (r *recorder) OpenLTXFile(ctx context.Context, level int, minTXID, maxTXID
return &headerTee{ReadCloser: rc, buf: r.header}, nil
}
+// dbDownloadMaxStalls is how many reopens in a row may deliver no byte
+// before a download gives up. Progress resets the count, so a slow but
+// moving link always finishes.
+const dbDownloadMaxStalls = 5
+
// download reads one file whole into the private directory and closes the
// stream before the restore reads a byte of it. litestream's compactor
// holds every planned file open until the restore ends, so streaming
// straight through would hold one door stream per planned file, and the
// door bounds a person's concurrent streams (apidb.go). Downloaded, a
-// restore holds one stream at a time. The file goes when the restore closes
-// it; on failure, with the whole directory. A download that fails part way
-// is an ordinary error, which litestream's resumable reader retries.
-func (r *recorder) download(rc io.ReadCloser) (io.ReadCloser, error) {
- defer rc.Close()
+// restore holds one stream at a time.
+//
+// A connection that drops part way is resumed here, from the bytes already
+// kept, rather than handed to litestream: its resumable reader would start
+// the file again from the requested offset, and a large file over a flaky
+// link would never finish. The file goes when the restore closes it; on
+// failure, with the whole directory.
+func (r *recorder) download(ctx context.Context, rc io.ReadCloser, level int, minTXID, maxTXID ltx.TXID, offset, size int64) (io.ReadCloser, error) {
f, err := os.CreateTemp(r.dir, "ltx-*")
if err != nil {
+ rc.Close()
return nil, err
}
- _, err = io.Copy(f, rc)
- if err == nil {
- _, err = f.Seek(0, io.SeekStart)
- }
- if err != nil {
+ fail := func(err error) (io.ReadCloser, error) {
f.Close()
os.Remove(f.Name())
return nil, err
}
- return &downloaded{f}, nil
+ var written int64
+ for stalls := 0; ; {
+ n, cerr := io.Copy(f, rc)
+ rc.Close()
+ written += n
+ if cerr == nil {
+ break
+ }
+ if n > 0 {
+ stalls = 0
+ } else if stalls++; stalls >= dbDownloadMaxStalls {
+ return fail(cerr)
+ }
+ if ctx.Err() != nil {
+ return fail(ctx.Err())
+ }
+ var rest int64 // 0: to the end, as RangeHeader reads it
+ if size > 0 {
+ if rest = size - written; rest <= 0 {
+ break
+ }
+ }
+ if rc, err = r.Client.OpenLTXFile(ctx, level, minTXID, maxTXID, offset+written, rest); err != nil {
+ if errors.Is(err, os.ErrNotExist) {
+ r.mu.Lock()
+ r.vanished = true
+ r.mu.Unlock()
+ }
+ return fail(err)
+ }
+ }
+ if _, err := f.Seek(0, io.SeekStart); err != nil {
+ return fail(err)
+ }
+ return &downloaded{ctx: ctx, File: f}, nil
}
// downloaded is a file read once and removed on Close.
-type downloaded struct{ *os.File }
+type downloaded struct {
+ ctx context.Context
+ *os.File
+}
func (d *downloaded) Close() error {
err := d.File.Close()
diff --git a/internal/dbdoor/restore_test.go b/internal/dbdoor/restore_test.go
index bf20f749a..e169f1e14 100644
--- a/internal/dbdoor/restore_test.go
+++ b/internal/dbdoor/restore_test.go
@@ -206,3 +206,68 @@ func TestRestoreToHoldsOneStreamAtATime(t *testing.T) {
t.Fatalf("downloaded files left in the private directory: %v", ents)
}
}
+
+// flakyBody hands over at most limit bytes and then fails the way a reset
+// connection does.
+type flakyBody struct {
+ io.ReadCloser
+ left int
+}
+
+func (b *flakyBody) Read(p []byte) (int, error) {
+ if b.left <= 0 {
+ return 0, io.ErrUnexpectedEOF
+ }
+ if len(p) > b.left {
+ p = p[:b.left]
+ }
+ n, err := b.ReadCloser.Read(p)
+ b.left -= n
+ return n, err
+}
+
+// flaky wraps a door so every file response breaks after limit bytes.
+func flaky(door *testDoor, limit int, gets *atomic.Int32) Doer {
+ return func(ctx context.Context, tail string, h http.Header) (*http.Response, error) {
+ res, err := door.do(ctx, tail, h)
+ if err != nil || !strings.Contains(strings.TrimPrefix(tail, "/ltx/"), "/") {
+ return res, err
+ }
+ gets.Add(1)
+ res.Body = &flakyBody{ReadCloser: res.Body, left: limit}
+ return res, nil
+ }
+}
+
+// Over a link that resets every few hundred bytes, a download keeps what it
+// has and asks for the rest, so a file larger than any one connection
+// survives still arrives.
+func TestRestoreToResumesAFileAcrossResets(t *testing.T) {
+ src := t.TempDir()
+ dbdoortest.Replicate(t, src, 2)
+ var gets atomic.Int32
+ got, err := RestoreTo(context.Background(), NewClient(flaky(newTestDoor(t, src), 100, &gets)), t.TempDir())
+ if err != nil {
+ t.Fatalf("RestoreTo over a resetting link = %v, want it to resume and finish", err)
+ }
+ defer os.RemoveAll(got.Dir)
+ if rows := countRows(t, got.Path); rows != 2*dbdoortest.Rows {
+ t.Fatalf("restored %d rows, want %d", rows, 2*dbdoortest.Rows)
+ }
+}
+
+// A link that delivers nothing is given up on after a bounded number of
+// attempts, not retried forever.
+func TestRestoreToGivesUpOnALinkThatMakesNoProgress(t *testing.T) {
+ src := t.TempDir()
+ dbdoortest.Replicate(t, src, 1)
+ var gets atomic.Int32
+ parent := t.TempDir()
+ if _, err := RestoreTo(context.Background(), NewClient(flaky(newTestDoor(t, src), 0, &gets)), parent); err == nil {
+ t.Fatal("RestoreTo over a dead link succeeded")
+ }
+ if n := gets.Load(); n > 40 {
+ t.Fatalf("made %d attempts at a link that never delivered a byte", n)
+ }
+ assertEmptyDir(t, parent)
+}