amadan

carlos / platform Public

db door: bound active streams, cut idle ones, restore one file at a time

The per-minute token budget bounds how many requests start, not how many
are still running, and io.Copy blocked on a stalled reader with no write
deadline, so slow downloads accumulated.

- Active file streams are bounded per principal (4) and across the
  console (32); a full door answers 429 with Retry-After: 5.
- Each stream must move a byte every 30s in both directions: a watchdog
  renewed per chunk cancels a stalled bucket read, and a write deadline
  renewed per chunk through http.ResponseController cuts a client that
  stopped reading. The deadline is cleared after, since the server sets
  no WriteTimeout and a kept-alive connection would carry it over.

The client had to change with it. litestream's compactor holds every
planned file open until the restore ends, so streaming straight through
needed one door stream per planned file and would deadlock against a
per-person bound. RestoreTo now downloads each file whole into its
private 0700 directory before litestream reads it, and removes it when
closed: one stream at a time, pinned by a test.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
claude 51badfd74c45210f230ca46ea317392aa2d0deeb parent ab02ff4
5 files changed, +346 −3
  • internal/console/apidb.go +94 −2
  • internal/console/apidb_test.go +151 −0
  • internal/console/server.go +6 −0
  • internal/dbdoor/restore.go +42 −1
  • internal/dbdoor/restore_test.go +53 −0
diff --git a/internal/console/apidb.go b/internal/console/apidb.go
index 5013973c7..5de0cd0d7 100644
--- a/internal/console/apidb.go
+++ b/internal/console/apidb.go
@@ -294,7 +294,22 @@ func (s *Server) handleAPIDBGet(w http.ResponseWriter, r *http.Request) {
jsonError(w, http.StatusRequestedRangeNotSatisfiable, "unsupported range")
return
}
- rc, err := g.rd.OpenLTXFile(r.Context(), level, ltx.TXID(minTXID), ltx.TXID(maxTXID), offset, size)
+ if !s.dbStreams.acquire(g.actor) {
+ w.Header().Set("Retry-After", dbStreamRetryAfter)
+ jsonError(w, http.StatusTooManyRequests, "too many database downloads at once. Wait a few seconds and try again.")
+ return
+ }
+ defer s.dbStreams.release(g.actor)
+ // The idle watchdog: any stretch of dbStreamIdle with no byte moved
+ // cancels the bucket read. The write side has its own deadline below,
+ // renewed per chunk, because a cancelled context does not unblock a
+ // write to a client that has stopped reading.
+ idle := s.dbStreamIdle
+ ctx, cancel := context.WithCancel(r.Context())
+ defer cancel()
+ watchdog := time.AfterFunc(idle, cancel)
+ defer watchdog.Stop()
+ rc, err := g.rd.OpenLTXFile(ctx, level, ltx.TXID(minTXID), ltx.TXID(maxTXID), offset, size)
if errors.Is(err, os.ErrNotExist) {
http.NotFound(w, r)
return
@@ -321,8 +336,13 @@ func (s *Server) handleAPIDBGet(w http.ResponseWriter, r *http.Request) {
if rangeHdr != "" {
status = http.StatusPartialContent
}
+ rw := http.NewResponseController(w)
+ // The server sets no WriteTimeout, so a deadline left here would carry
+ // over to the next request on a kept-alive connection. Clear it.
+ defer rw.SetWriteDeadline(time.Time{})
+ _ = rw.SetWriteDeadline(time.Now().Add(idle))
w.WriteHeader(status)
- if _, err := io.Copy(w, rc); err != nil {
+ if err := dbStream(w, rw, rc, watchdog, idle); err != nil {
// The status is already on the wire. The client's resumable reader
// sees a short body against the listed size and reopens at the
// offset it reached.
@@ -331,6 +351,31 @@ func (s *Server) handleAPIDBGet(w http.ResponseWriter, r *http.Request) {
}
}
+// dbStream copies rc to w a chunk at a time, renewing the write deadline
+// before each chunk and the read watchdog after it, so a stream ends when
+// either side stops moving for idle.
+func dbStream(w io.Writer, rw *http.ResponseController, rc io.Reader, watchdog *time.Timer, idle time.Duration) error {
+ buf := make([]byte, 32<<10)
+ for {
+ n, rerr := rc.Read(buf)
+ if n > 0 {
+ if err := rw.SetWriteDeadline(time.Now().Add(idle)); err != nil && !errors.Is(err, http.ErrNotSupported) {
+ return err
+ }
+ if _, err := w.Write(buf[:n]); err != nil {
+ return err
+ }
+ watchdog.Reset(idle)
+ }
+ if rerr == io.EOF {
+ return nil
+ }
+ if rerr != nil {
+ return rerr
+ }
+ }
+}
+
// dbReplicaCacheMax bounds the replica-client cache. Each client owns an
// HTTP transport, so one per request would cost a TLS handshake per file and
// leave an idle connection behind each; one per prefix lets a restore's
@@ -363,3 +408,50 @@ func (s *Server) cachedReadReplica(ctx context.Context, prefix string) (replicaR
c.clients[prefix] = rc
return rc, nil
}
+
+// The db door's stream bounds. The per-minute token budget bounds how many
+// requests start, not how many are still running, and a slow or stalled
+// download holds a goroutine, an S3 connection and a client connection for
+// as long as it lasts. So active file streams are bounded per principal and
+// across the console, and refused with 429 and Retry-After when full; and
+// each stream must move a byte at least every dbStreamIdle, in both
+// directions, or it is cut. The CLI's restore downloads one file at a time
+// (dbdoor.RestoreTo), so a person's restore needs one slot.
+const (
+ dbMaxStreams = 32
+ dbMaxStreamsPerPrincipal = 4
+ dbStreamIdleDefault = 30 * time.Second
+ dbStreamRetryAfter = "5"
+)
+
+type dbStreamLimits struct {
+ mu sync.Mutex
+ maxTotal, maxPer int
+ total int
+ per map[string]int
+}
+
+func newDBStreamLimits(maxTotal, maxPer int) *dbStreamLimits {
+ return &dbStreamLimits{maxTotal: maxTotal, maxPer: maxPer, per: map[string]int{}}
+}
+
+// acquire takes a slot for key, or reports that the console or key is full.
+func (l *dbStreamLimits) acquire(key string) bool {
+ l.mu.Lock()
+ defer l.mu.Unlock()
+ if l.total >= l.maxTotal || l.per[key] >= l.maxPer {
+ return false
+ }
+ l.total++
+ l.per[key]++
+ return true
+}
+
+func (l *dbStreamLimits) release(key string) {
+ l.mu.Lock()
+ defer l.mu.Unlock()
+ l.total--
+ if l.per[key]--; l.per[key] <= 0 {
+ delete(l.per, key)
+ }
+}
diff --git a/internal/console/apidb_test.go b/internal/console/apidb_test.go
index 4aec2ddd3..334f83b63 100644
--- a/internal/console/apidb_test.go
+++ b/internal/console/apidb_test.go
@@ -552,6 +552,10 @@ func TestAPIDBHasItsOwnBudget(t *testing.T) {
func TestAPIDBRestoresThroughTheDoor(t *testing.T) {
fx := dbFixture(t)
fx.replicate(t, dbTestHost, 3)
+ // The tightest bound there is, and a restore still completes under it.
+ // (Small files finish server-side before the client reads them, so the
+ // one-stream-at-a-time property itself is pinned in internal/dbdoor.)
+ fx.srv.dbStreams = newDBStreamLimits(1, 1)
do := func(ctx context.Context, tail string, h http.Header) (*http.Response, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, fx.url(dbdoor.Base("bab", "titogo", dbTestHost)+tail), nil)
if err != nil {
@@ -611,3 +615,150 @@ func TestAPIDBRangedGetIsAudited(t *testing.T) {
}
t.Fatalf("a ranged read wrote no db-read row naming its start: %+v", evs)
}
+
+// stallReplica serves listings from a file replica but opens files whose
+// reads never finish: each blocks until the request's context ends, the way
+// a stalled S3 read does, or produces zeros forever when endless is set.
+type stallReplica struct {
+ replicaReader
+ endless bool
+ opened chan struct{}
+}
+
+func (s *stallReplica) OpenLTXFile(ctx context.Context, level int, min, max ltx.TXID, offset, size int64) (io.ReadCloser, error) {
+ select {
+ case s.opened <- struct{}{}:
+ default:
+ }
+ if s.endless {
+ return io.NopCloser(zeros{}), nil
+ }
+ return io.NopCloser(ctxReader{ctx}), nil
+}
+
+type ctxReader struct{ ctx context.Context }
+
+func (r ctxReader) Read([]byte) (int, error) { <-r.ctx.Done(); return 0, r.ctx.Err() }
+
+type zeros struct{}
+
+func (zeros) Read(p []byte) (int, error) { clear(p); return len(p), nil }
+
+func (fx *dbFx) stall(endless bool) *stallReplica {
+ sr := &stallReplica{endless: endless, opened: make(chan struct{}, 16)}
+ fx.srv.dbReplica = func(_ context.Context, prefix string) (replicaReader, error) {
+ sr.replicaReader = lsfile.NewReplicaClient(filepath.Join(fx.root, prefix))
+ return sr, nil
+ }
+ return sr
+}
+
+func (fx *dbFx) memberToken(t *testing.T, addr, code string) string {
+ t.Helper()
+ if err := fx.st.AddMember(OpsAccountID, addr, RoleMember, time.Now()); err != nil {
+ t.Fatal(err)
+ }
+ return cliToken(t, fx.consoleFixture, addr, code)
+}
+
+// startGet begins a GET and returns once the door has opened the file, so
+// the stream holds its slot; cancel ends it.
+func (fx *dbFx) startGet(t *testing.T, sr *stallReplica, tok string) (stop func()) {
+ t.Helper()
+ ctx, cancel := context.WithCancel(context.Background())
+ req, _ := http.NewRequestWithContext(ctx, http.MethodGet, fx.url(dbdoor.GetPath("bab", "titogo", dbTestHost, 0, 1, 1)), nil)
+ req.Header.Set("Authorization", "Bearer "+tok)
+ done := make(chan struct{})
+ go func() {
+ defer close(done)
+ if res, err := http.DefaultClient.Do(req); err == nil {
+ io.Copy(io.Discard, res.Body)
+ res.Body.Close()
+ }
+ }()
+ select {
+ case <-sr.opened:
+ case <-time.After(5 * time.Second):
+ t.Fatal("the stream never opened")
+ }
+ return func() { cancel(); <-done }
+}
+
+func (fx *dbFx) activeStreams() int {
+ fx.srv.dbStreams.mu.Lock()
+ defer fx.srv.dbStreams.mu.Unlock()
+ return fx.srv.dbStreams.total
+}
+
+func (fx *dbFx) waitIdle(t *testing.T, within time.Duration) {
+ t.Helper()
+ deadline := time.Now().Add(within)
+ for fx.activeStreams() != 0 {
+ if time.Now().After(deadline) {
+ t.Fatalf("%d streams still hold slots after %v", fx.activeStreams(), within)
+ }
+ time.Sleep(10 * time.Millisecond)
+ }
+}
+
+// Slow downloads must not accumulate: active file streams are bounded per
+// principal and across the console, and a refusal says when to come back.
+func TestAPIDBBoundsConcurrentStreams(t *testing.T) {
+ fx := dbFixture(t)
+ fx.replicate(t, dbTestHost, 1)
+ fx.srv.dbStreams = newDBStreamLimits(2, 1)
+ sr := fx.stall(false)
+ bob := fx.memberToken(t, "bob@keymail.dev", "code-db-bob")
+ carol := fx.memberToken(t, "carol@keymail.dev", "code-db-carol")
+ u := fx.url(dbdoor.GetPath("bab", "titogo", dbTestHost, 0, 1, 1))
+
+ stopAlice := fx.startGet(t, sr, fx.member)
+ status, h, body := rawReq(t, http.MethodGet, u, fx.member, "")
+ if status != http.StatusTooManyRequests || h.Get("Retry-After") == "" {
+ t.Fatalf("second stream for one person = %d %q %s, want 429 with Retry-After", status, h.Get("Retry-After"), body)
+ }
+ stopBob := fx.startGet(t, sr, bob)
+ status, h, body = rawReq(t, http.MethodGet, u, carol, "")
+ if status != http.StatusTooManyRequests || h.Get("Retry-After") == "" {
+ t.Fatalf("stream past the console's bound = %d %q %s, want 429 with Retry-After", status, h.Get("Retry-After"), body)
+ }
+ stopAlice()
+ stopBob()
+ fx.waitIdle(t, 5*time.Second)
+ stopAgain := fx.startGet(t, sr, fx.member)
+ stopAgain()
+}
+
+// A stalled bucket read ends after the idle deadline and frees its slot.
+func TestAPIDBStalledReadIsCut(t *testing.T) {
+ fx := dbFixture(t)
+ fx.replicate(t, dbTestHost, 1)
+ fx.srv.dbStreamIdle = 200 * time.Millisecond
+ sr := fx.stall(false)
+ start := time.Now()
+ status, _, _ := rawReq(t, http.MethodGet, fx.url(dbdoor.GetPath("bab", "titogo", dbTestHost, 0, 1, 1)), fx.member, "")
+ <-sr.opened
+ if time.Since(start) > 5*time.Second {
+ t.Fatalf("a stalled read held the request for %v", time.Since(start))
+ }
+ _ = status
+ fx.waitIdle(t, 5*time.Second)
+}
+
+// A client that stops reading is cut by the renewable write deadline, so a
+// trickling or abandoned download cannot hold a slot forever.
+func TestAPIDBStalledClientIsCut(t *testing.T) {
+ fx := dbFixture(t)
+ fx.replicate(t, dbTestHost, 1)
+ fx.srv.dbStreamIdle = 300 * time.Millisecond
+ sr := fx.stall(true)
+ req, _ := http.NewRequest(http.MethodGet, fx.url(dbdoor.GetPath("bab", "titogo", dbTestHost, 0, 1, 1)), nil)
+ req.Header.Set("Authorization", "Bearer "+fx.member)
+ res, err := http.DefaultClient.Do(req)
+ if err != nil {
+ t.Fatal(err)
+ }
+ defer res.Body.Close() // never read: the client has stalled
+ <-sr.opened
+ fx.waitIdle(t, 15*time.Second)
+}
diff --git a/internal/console/server.go b/internal/console/server.go
index f3e7e9f9e..c8ff5c1af 100644
--- a/internal/console/server.go
+++ b/internal/console/server.go
@@ -729,6 +729,10 @@ type Server struct {
// lease.ReadReplicaClient values; tests point it at a file replica.
dbReplica func(ctx context.Context, prefix string) (replicaReader, error)
dbClients dbReplicaCache
+ // dbStreams bounds the db door's active file streams, and dbStreamIdle
+ // is how long one may go without moving a byte (apidb.go).
+ dbStreams *dbStreamLimits
+ dbStreamIdle time.Duration
// cliRoles is pattern -> allowed roles, recorded by cliRoute as it
// registers each /api/cli/ door (authcli.go). Written once, during
// registerCLIRoutes, before the server serves anything, and only read
@@ -1154,6 +1158,8 @@ func NewServer(cfg Config, st *Store) (*Server, error) {
// approval page, and the tokens an owner can see and kill (authcli.go).
// The db door's production replica reader (apidb.go); tests repoint it.
s.dbReplica = s.cachedReadReplica
+ s.dbStreams = newDBStreamLimits(dbMaxStreams, dbMaxStreamsPerPrincipal)
+ s.dbStreamIdle = dbStreamIdleDefault
s.registerCLIRoutes()
// The email API's ten doors (apimail.go, tenant-email task 17) —
// registered separately from registerCLIRoutes because this task's file
diff --git a/internal/dbdoor/restore.go b/internal/dbdoor/restore.go
index f96f79b57..824d9859b 100644
--- a/internal/dbdoor/restore.go
+++ b/internal/dbdoor/restore.go
@@ -74,7 +74,7 @@ func RestoreTo(ctx context.Context, c *Client, parent string) (*Restored, error)
out := filepath.Join(dir, restoredName)
var rec *recorder
for attempt := 0; ; attempt++ {
- rec = &recorder{Client: c}
+ rec = &recorder{Client: c, dir: dir}
err = restoreOnce(ctx, rec, out)
if err == nil {
break
@@ -121,6 +121,7 @@ func removeCopy(out string) {
// newest file opened is the copy's "as of" (R1).
type recorder struct {
*Client
+ dir string // the private directory each file is downloaded into
mu sync.Mutex
vanished bool
maxTXID ltx.TXID
@@ -137,6 +138,9 @@ func (r *recorder) OpenLTXFile(ctx context.Context, level int, minTXID, maxTXID
}
return nil, err
}
+ if rc, err = r.download(rc); err != nil {
+ return nil, err
+ }
if offset != 0 {
return rc, nil
}
@@ -150,6 +154,43 @@ func (r *recorder) OpenLTXFile(ctx context.Context, level int, minTXID, maxTXID
return &headerTee{ReadCloser: rc, buf: r.header}, nil
}
+// 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()
+ f, err := os.CreateTemp(r.dir, "ltx-*")
+ if err != nil {
+ return nil, err
+ }
+ _, err = io.Copy(f, rc)
+ if err == nil {
+ _, err = f.Seek(0, io.SeekStart)
+ }
+ if err != nil {
+ f.Close()
+ os.Remove(f.Name())
+ return nil, err
+ }
+ return &downloaded{f}, nil
+}
+
+// downloaded is a file read once and removed on Close.
+type downloaded struct{ *os.File }
+
+func (d *downloaded) Close() error {
+ err := d.File.Close()
+ if rerr := os.Remove(d.File.Name()); err == nil {
+ err = rerr
+ }
+ return err
+}
+
func (r *recorder) sawVanish() bool {
r.mu.Lock()
defer r.mu.Unlock()
diff --git a/internal/dbdoor/restore_test.go b/internal/dbdoor/restore_test.go
index 3d4dbeba5..bf20f749a 100644
--- a/internal/dbdoor/restore_test.go
+++ b/internal/dbdoor/restore_test.go
@@ -3,9 +3,13 @@ package dbdoor
import (
"context"
"errors"
+ "io"
"net/http"
"os"
"path/filepath"
+ "strings"
+ "sync"
+ "sync/atomic"
"testing"
"time"
@@ -153,3 +157,52 @@ func assertEmptyDir(t *testing.T, dir string) {
t.Fatalf("%s holds %d entries after a failed restore, want none (no copy of production left behind)", dir, len(ents))
}
}
+
+// countingBody tracks how many door responses are open at once.
+type countingBody struct {
+ io.ReadCloser
+ open *atomic.Int32
+ once sync.Once
+}
+
+func (b *countingBody) Close() error {
+ b.once.Do(func() { b.open.Add(-1) })
+ return b.ReadCloser.Close()
+}
+
+// The door bounds a person's concurrent file streams, and litestream's
+// compactor holds every planned file open until the restore ends. Streaming
+// straight through would need one slot per file, so RestoreTo downloads each
+// file whole into its private directory before handing it on: one stream at
+// a time, whatever the plan's length.
+func TestRestoreToHoldsOneStreamAtATime(t *testing.T) {
+ src := t.TempDir()
+ dbdoortest.Replicate(t, src, 4)
+ door := newTestDoor(t, src)
+ var open, peak atomic.Int32
+ do := 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 // listings are read and closed at once
+ }
+ n := open.Add(1)
+ for p := peak.Load(); n > p && !peak.CompareAndSwap(p, n); p = peak.Load() {
+ }
+ res.Body = &countingBody{ReadCloser: res.Body, open: &open}
+ return res, nil
+ }
+ got, err := RestoreTo(context.Background(), NewClient(do), t.TempDir())
+ if err != nil {
+ t.Fatal(err)
+ }
+ defer os.RemoveAll(got.Dir)
+ if p := peak.Load(); p != 1 {
+ t.Fatalf("restore held %d file streams open at once, want 1", p)
+ }
+ if rows := countRows(t, got.Path); rows != 4*dbdoortest.Rows {
+ t.Fatalf("restored %d rows, want %d", rows, 4*dbdoortest.Rows)
+ }
+ if ents, _ := os.ReadDir(got.Dir); len(ents) > 3 {
+ t.Fatalf("downloaded files left in the private directory: %v", ents)
+ }
+}