waitfordb: initial tool — env-configured wait-for-DB init container

Small Go tool + distroless container used as a K8s initContainer to block an app until its database (Postgres/MySQL) is reachable. Env-var configured (WAITFORDB_* + libpq PG* fallback), configurable timeout/interval, redacted logs, exit codes. Woodpecker CI publishes docker-internal/waitfordb on tag.
2026-08-22 13:31:44 +10:00
parent 98971eb1b7
commit 2b37986218
19 changed files with 1467 additions and 1 deletions
@@ -0,0 +1,113 @@
+// Package wait implements the retry/timeout loop that polls a database until a
+// liveness attempt succeeds. The clock and the attempt are injected so the loop
+// is fully testable without a real database or real time.
+package wait
+
+import (
+	"context"
+	"time"
+)
+
+// Clock abstracts time so tests can advance it instantly.
+type Clock interface {
+	Now() time.Time
+	// Sleep blocks for d or until ctx is done, returning ctx.Err() if it was
+	// cancelled first.
+	Sleep(ctx context.Context, d time.Duration) error
+}
+
+// AttemptFunc performs one connect+ping. The ctx carries the per-attempt
+// connect timeout.
+type AttemptFunc func(ctx context.Context) error
+
+// OnFailure is invoked after each failed attempt that will be retried.
+type OnFailure func(attempt int, err error, elapsed, timeout, retryIn time.Duration)
+
+// Params configures the loop.
+type Params struct {
+	Timeout        time.Duration // 0 = wait forever
+	Interval       time.Duration
+	ConnectTimeout time.Duration
+}
+
+// Result reports how a run ended.
+type Result struct {
+	OK        bool
+	TimedOut  bool
+	Cancelled bool
+	Attempts  int
+	Elapsed   time.Duration
+	LastErr   error
+}
+
+// Run polls attempt until it succeeds, the timeout is exhausted, or ctx is
+// cancelled. It always makes at least one attempt.
+func Run(ctx context.Context, p Params, attempt AttemptFunc, onFail OnFailure, clk Clock) Result {
+	start := clk.Now()
+	var deadline time.Time
+	if p.Timeout > 0 {
+		deadline = start.Add(p.Timeout)
+	}
+
+	res := Result{}
+	for {
+		res.Attempts++
+
+		actx, cancel := context.WithTimeout(ctx, p.ConnectTimeout)
+		err := attempt(actx)
+		cancel()
+
+		now := clk.Now()
+		res.Elapsed = now.Sub(start)
+
+		if err == nil {
+			res.OK = true
+			return res
+		}
+		res.LastErr = err
+
+		// A cancelled parent context (SIGTERM/SIGINT) wins over a retry.
+		if ctx.Err() != nil {
+			res.Cancelled = true
+			return res
+		}
+
+		// No time budget left for another attempt.
+		if p.Timeout > 0 && !now.Before(deadline) {
+			res.TimedOut = true
+			return res
+		}
+
+		sleep := p.Interval
+		if p.Timeout > 0 {
+			if remaining := deadline.Sub(now); remaining < sleep {
+				sleep = remaining
+			}
+		}
+		onFail(res.Attempts, err, res.Elapsed, p.Timeout, sleep)
+
+		if serr := clk.Sleep(ctx, sleep); serr != nil {
+			res.Cancelled = true
+			return res
+		}
+	}
+}
+
+// RealClock is the production Clock backed by the wall clock.
+type RealClock struct{}
+
+func (RealClock) Now() time.Time { return time.Now() }
+
+func (RealClock) Sleep(ctx context.Context, d time.Duration) error {
+	if d <= 0 {
+		return ctx.Err()
+	}
+	t := time.NewTimer(d)
+	defer t.Stop()
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-t.C:
+		return nil
+	}
+}
@@ -0,0 +1,125 @@
+package wait
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+)
+
+// fakeClock advances instantly on Sleep so the loop runs with no real delay.
+type fakeClock struct {
+	t         time.Time
+	cancelAt  time.Duration // if >0, cancel the run once elapsed reaches this
+	cancel    context.CancelFunc
+	start     time.Time
+	sleepCall int
+}
+
+func newFakeClock() *fakeClock {
+	start := time.Unix(0, 0)
+	return &fakeClock{t: start, start: start}
+}
+
+func (c *fakeClock) Now() time.Time { return c.t }
+
+func (c *fakeClock) Sleep(ctx context.Context, d time.Duration) error {
+	c.sleepCall++
+	c.t = c.t.Add(d)
+	if c.cancelAt > 0 && c.t.Sub(c.start) >= c.cancelAt && c.cancel != nil {
+		c.cancel()
+	}
+	return ctx.Err()
+}
+
+var errDown = errors.New("connection refused")
+
+// failNThenOK returns an AttemptFunc that fails the first n calls then succeeds.
+func failNThenOK(n int, calls *int) AttemptFunc {
+	return func(ctx context.Context) error {
+		*calls++
+		if *calls <= n {
+			return errDown
+		}
+		return nil
+	}
+}
+
+func noFail(int, error, time.Duration, time.Duration, time.Duration) {}
+
+func TestSucceedsFirstAttempt(t *testing.T) {
+	calls := 0
+	res := Run(context.Background(),
+		Params{Timeout: time.Minute, Interval: 2 * time.Second, ConnectTimeout: time.Second},
+		failNThenOK(0, &calls), noFail, newFakeClock())
+	if !res.OK || res.Attempts != 1 {
+		t.Fatalf("want OK after 1 attempt, got %+v", res)
+	}
+}
+
+func TestWaitsThenSucceeds(t *testing.T) {
+	calls := 0
+	clk := newFakeClock()
+	failures := 0
+	res := Run(context.Background(),
+		Params{Timeout: time.Minute, Interval: 2 * time.Second, ConnectTimeout: time.Second},
+		failNThenOK(3, &calls),
+		func(int, error, time.Duration, time.Duration, time.Duration) { failures++ },
+		clk)
+	if !res.OK {
+		t.Fatalf("want OK, got %+v", res)
+	}
+	if res.Attempts != 4 {
+		t.Errorf("attempts = %d, want 4", res.Attempts)
+	}
+	if failures != 3 {
+		t.Errorf("onFail called %d times, want 3", failures)
+	}
+	// 3 sleeps of 2s each.
+	if got := res.Elapsed; got != 6*time.Second {
+		t.Errorf("elapsed = %v, want 6s", got)
+	}
+}
+
+func TestTimesOut(t *testing.T) {
+	calls := 0
+	alwaysFail := func(ctx context.Context) error { calls++; return errDown }
+	res := Run(context.Background(),
+		Params{Timeout: 10 * time.Second, Interval: 3 * time.Second, ConnectTimeout: time.Second},
+		alwaysFail, noFail, newFakeClock())
+	if res.OK || !res.TimedOut {
+		t.Fatalf("want timeout, got %+v", res)
+	}
+	if !errors.Is(res.LastErr, errDown) {
+		t.Errorf("LastErr = %v, want errDown", res.LastErr)
+	}
+	// Deadline 10s, interval 3s: attempts at 0,3,6,9, then next check at ~12s > deadline.
+	if res.Attempts < 3 {
+		t.Errorf("attempts = %d, want several before timeout", res.Attempts)
+	}
+}
+
+func TestWaitForeverEventuallySucceeds(t *testing.T) {
+	calls := 0
+	res := Run(context.Background(),
+		Params{Timeout: 0, Interval: time.Second, ConnectTimeout: time.Second},
+		failNThenOK(100, &calls), noFail, newFakeClock())
+	if !res.OK || res.Attempts != 101 {
+		t.Fatalf("want OK after 101 attempts with no timeout, got %+v", res)
+	}
+}
+
+func TestCancelledDuringSleep(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	clk := newFakeClock()
+	clk.cancelAt = 4 * time.Second
+	clk.cancel = cancel
+	calls := 0
+	alwaysFail := func(ctx context.Context) error { calls++; return errDown }
+	res := Run(ctx,
+		Params{Timeout: time.Hour, Interval: 2 * time.Second, ConnectTimeout: time.Second},
+		alwaysFail, noFail, clk)
+	if !res.Cancelled {
+		t.Fatalf("want Cancelled, got %+v", res)
+	}
+}