fix(store): 新 pod 首次连库在网络策略放行前重试

cmd/felis/api.go

@@ -100,7 +100,7 @@ func cmdAPI(args []string, stdout, stderr io.Writer) int {
 
 	// Before anything serves: an api on a schema it was not built for answers with
 	// errors, or writes rows the other version cannot read.
-	drv, err := openStore(ctx, cfg.Database.URL, false)
+	drv, err := openPodStore(ctx, cfg.Database.URL, "api", stderr)
 	if err != nil {
 		fmt.Fprintf(stderr, "felis api: open database: %v\n", err)
 		return 1

cmd/felis/backup.go

@@ -15,7 +15,6 @@ import (
 	"felis.lolicon.best/internal/config"
 	"felis.lolicon.best/internal/naming"
 	"felis.lolicon.best/internal/reaper"
-	"felis.lolicon.best/internal/store"
 	ctrl "sigs.k8s.io/controller-runtime"
 )
 
@@ -100,7 +99,7 @@ func cmdBackup(args []string, stdout, stderr io.Writer) int {
 			len(a.Skipped), strings.Join(a.Skipped[:min(len(a.Skipped), 10)], ", "))
 	}
 
-	drv, err := store.Open(ctx, cfg.Database.URL)
+	drv, err := openPodStore(ctx, cfg.Database.URL, "backup", stderr)
 	if err != nil {
 		fmt.Fprintf(stderr, "felis backup: open database: %v\n", err)
 		return 1

cmd/felis/migrate.go

@@ -6,6 +6,7 @@ import (
 	"flag"
 	"fmt"
 	"io"
+	"time"
 
 	"felis.lolicon.best/internal/config"
 	"felis.lolicon.best/internal/dbbackup"
@@ -138,6 +139,33 @@ func openStore(ctx context.Context, url string, allowPending bool) (*store.Postg
 	if err != nil {
 		return nil, err
 	}
+	return checkSchema(ctx, drv, allowPending)
+}
+
+// podDBWindow and podDBInterval bound how long a pod that has just started retries its
+// first database dial while the network policy has yet to admit it
+// (store.OpenRetrying). A minute is far past the sync lag and far inside every Job's
+// deadline. Vars so a test can shrink them.
+var (
+	podDBWindow   = time.Minute
+	podDBInterval = time.Second
+)
+
+// openPodStore is openStore for felis-api and the reaper and backup Jobs, whose first
+// dial comes milliseconds after their pod starts.
+func openPodStore(ctx context.Context, url, prog string, stderr io.Writer) (*store.PostgresDriver, error) {
+	drv, err := store.OpenRetrying(ctx, url, podDBWindow, podDBInterval, func(err error) {
+		fmt.Fprintf(stderr, "felis %s: %v; retrying (a pod that has just started waits for the network policy to admit it)\n", prog, err)
+	})
+	if err != nil {
+		return nil, err
+	}
+	return checkSchema(ctx, drv, false)
+}
+
+// checkSchema closes drv and fails when its schema is not the one this build was
+// written against (see openStore).
+func checkSchema(ctx context.Context, drv *store.PostgresDriver, allowPending bool) (*store.PostgresDriver, error) {
 	s, err := store.ReadSchema(ctx, drv)
 	if err == nil {
 		err = s.Err()

cmd/felis/migrate_test.go

@@ -0,0 +1,48 @@
+package main
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"net"
+	"strings"
+	"syscall"
+	"testing"
+	"time"
+)
+
+// openPodStore retries a refused first dial for podDBWindow, saying so on stderr
+// under the calling command's name each time.
+func TestOpenPodStoreRetriesARefusedDial(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	addr := ln.Addr().String()
+	ln.Close()
+
+	window, interval := podDBWindow, podDBInterval
+	podDBWindow, podDBInterval = 200*time.Millisecond, 20*time.Millisecond
+	t.Cleanup(func() { podDBWindow, podDBInterval = window, interval })
+
+	var stderr bytes.Buffer
+	start := time.Now()
+	_, err = openPodStore(context.Background(), fmt.Sprintf("postgres://felis@%s/felis?sslmode=disable", addr), "reaper", &stderr)
+	if !errors.Is(err, syscall.ECONNREFUSED) {
+		t.Fatalf("err = %v, want a refused dial", err)
+	}
+	if elapsed := time.Since(start); elapsed < podDBWindow {
+		t.Fatalf("gave up after %s, inside the %s window", elapsed, podDBWindow)
+	}
+	lines := strings.Split(strings.TrimSuffix(stderr.String(), "\n"), "\n")
+	if len(lines) < 2 || len(lines) > 11 {
+		t.Fatalf("%d retry lines in a 200ms window at 20ms:\n%s", len(lines), stderr.String())
+	}
+	for _, l := range lines {
+		if !strings.HasPrefix(l, "felis reaper: failed to connect to `user=felis database=felis`: ") ||
+			!strings.HasSuffix(l, "; retrying (a pod that has just started waits for the network policy to admit it)") {
+			t.Fatalf("retry line %q", l)
+		}
+	}
+}

cmd/felis/reaper.go

@@ -78,7 +78,7 @@ func cmdReaper(args []string, stdout, stderr io.Writer) int {
 		return 1
 	}
 
-	drv, err := openStore(ctx, cfg.Database.URL, false)
+	drv, err := openPodStore(ctx, cfg.Database.URL, "reaper", stderr)
 	if err != nil {
 		fmt.Fprintf(stderr, "felis reaper: open database: %v\n", err)
 		return 1

internal/store/open_test.go

@@ -0,0 +1,194 @@
+package store
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net"
+	"sync/atomic"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/jackc/pgx/v5/pgconn"
+	"github.com/jackc/pgx/v5/pgproto3"
+)
+
+// servePG answers every connection on ln like a PostgreSQL server would: the i-th
+// connection (from 1) gets a FATAL error with the SQLSTATE answer(i) returns, or,
+// for "", a trust login and an empty reply to each simple query, which is all a
+// ping sends. It counts the connections in accepts.
+func servePG(ln net.Listener, answer func(i int32) string, accepts *atomic.Int32) {
+	for {
+		conn, err := ln.Accept()
+		if err != nil {
+			return
+		}
+		code := answer(accepts.Add(1))
+		go func() {
+			defer conn.Close()
+			be := pgproto3.NewBackend(conn, conn)
+			if _, err := be.ReceiveStartupMessage(); err != nil {
+				return
+			}
+			if code != "" {
+				be.Send(&pgproto3.ErrorResponse{Severity: "FATAL", Code: code, Message: "answered " + code})
+				_ = be.Flush()
+				return
+			}
+			be.Send(&pgproto3.AuthenticationOk{})
+			be.Send(&pgproto3.BackendKeyData{ProcessID: 1, SecretKey: []byte{0, 0, 0, 1}})
+			be.Send(&pgproto3.ReadyForQuery{TxStatus: 'I'})
+			if be.Flush() != nil {
+				return
+			}
+			for {
+				msg, err := be.Receive()
+				if err != nil {
+					return
+				}
+				switch msg.(type) {
+				case *pgproto3.Query:
+					be.Send(&pgproto3.EmptyQueryResponse{})
+					be.Send(&pgproto3.ReadyForQuery{TxStatus: 'I'})
+					if be.Flush() != nil {
+						return
+					}
+				case *pgproto3.Terminate:
+					return
+				}
+			}
+		}()
+	}
+}
+
+// freeAddr returns a loopback address nothing listens on.
+func freeAddr(t *testing.T) string {
+	t.Helper()
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	addr := ln.Addr().String()
+	ln.Close()
+	return addr
+}
+
+func testDSN(addr string) string {
+	return fmt.Sprintf("postgres://felis@%s/felis?sslmode=disable", addr)
+}
+
+// A pod's first dials are refused until the network policy admits it; the open
+// waits that out and connects once the server is reachable.
+func TestOpenRetryingWaitsOutARefusedDial(t *testing.T) {
+	addr := freeAddr(t)
+	var accepts atomic.Int32
+	go func() {
+		time.Sleep(150 * time.Millisecond)
+		ln, err := net.Listen("tcp", addr)
+		if err != nil {
+			return
+		}
+		t.Cleanup(func() { ln.Close() })
+		servePG(ln, func(int32) string { return "" }, &accepts)
+	}()
+	var retries []error
+	drv, err := OpenRetrying(context.Background(), testDSN(addr), 5*time.Second, 20*time.Millisecond, func(err error) { retries = append(retries, err) })
+	if err != nil {
+		t.Fatalf("open: %v (after %d retries)", err, len(retries))
+	}
+	drv.Close()
+	if len(retries) == 0 {
+		t.Fatal("connected without a retry: the listener was up too early for this test to mean anything")
+	}
+	for _, r := range retries {
+		if !errors.Is(r, syscall.ECONNREFUSED) {
+			t.Fatalf("retried on %v, want only refused dials", r)
+		}
+	}
+	if n := accepts.Load(); n != 1 {
+		t.Fatalf("server saw %d connections, want 1", n)
+	}
+}
+
+// The window bounds the wait, and Open itself tries once.
+func TestOpenRetryingGivesUpAtTheWindow(t *testing.T) {
+	addr := freeAddr(t)
+	retries := 0
+	start := time.Now()
+	_, err := OpenRetrying(context.Background(), testDSN(addr), 200*time.Millisecond, 20*time.Millisecond, func(error) { retries++ })
+	if !errors.Is(err, syscall.ECONNREFUSED) {
+		t.Fatalf("err = %v, want a refused dial", err)
+	}
+	if elapsed := time.Since(start); elapsed < 200*time.Millisecond || elapsed > 3*time.Second {
+		t.Fatalf("gave up after %s, want the 200ms window", elapsed)
+	}
+	if retries < 2 || retries > 11 {
+		t.Fatalf("%d retries in a 200ms window at 20ms, want several and at most one per interval", retries)
+	}
+
+	// Open has no window: one refused dial is its answer (a retry would call the nil
+	// onRetry and panic).
+	if _, err := Open(context.Background(), testDSN(addr)); !errors.Is(err, syscall.ECONNREFUSED) {
+		t.Fatalf("Open err = %v, want a refused dial", err)
+	}
+}
+
+// What the server says is final, except that it is still starting.
+func TestOpenRetryingTakesTheServersAnswer(t *testing.T) {
+	for _, tc := range []struct {
+		name    string
+		answer  func(i int32) string
+		wantErr string // SQLSTATE of the final error; "" for a connection
+		accepts int32
+	}{
+		{"bad password", func(int32) string { return "28P01" }, "28P01", 1},
+		{"no such database", func(int32) string { return "3D000" }, "3D000", 1},
+		{"starting up", func(i int32) string {
+			if i <= 2 {
+				return "57P03"
+			}
+			return ""
+		}, "", 3},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			ln, err := net.Listen("tcp", "127.0.0.1:0")
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer ln.Close()
+			var accepts atomic.Int32
+			go servePG(ln, tc.answer, &accepts)
+			retries := 0
+			drv, err := OpenRetrying(context.Background(), testDSN(ln.Addr().String()), 2*time.Second, 10*time.Millisecond, func(error) { retries++ })
+			if tc.wantErr == "" {
+				if err != nil {
+					t.Fatalf("open: %v", err)
+				}
+				drv.Close()
+			} else {
+				var pgErr *pgconn.PgError
+				if !errors.As(err, &pgErr) || pgErr.Code != tc.wantErr {
+					t.Fatalf("err = %v, want SQLSTATE %s", err, tc.wantErr)
+				}
+			}
+			if n := accepts.Load(); n != tc.accepts || retries != int(tc.accepts)-1 {
+				t.Fatalf("server saw %d connections after %d retries, want %d", n, retries, tc.accepts)
+			}
+		})
+	}
+}
+
+// A cancelled context ends the wait without another retry.
+func TestOpenRetryingStopsWhenCancelled(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	retries := 0
+	start := time.Now()
+	if _, err := OpenRetrying(ctx, testDSN(freeAddr(t)), 5*time.Second, 20*time.Millisecond, func(error) { retries++ }); err == nil {
+		t.Fatal("opened on a cancelled context")
+	}
+	if retries != 0 || time.Since(start) > time.Second {
+		t.Fatalf("%d retries over %s on a cancelled context, want none", retries, time.Since(start))
+	}
+}

internal/store/sqldriver.go

@@ -79,15 +79,50 @@ func newPool(dsn string) (*sql.DB, error) {
 
 // Open dials dsn and returns a PostgresDriver. The caller owns Close.
 func Open(ctx context.Context, dsn string) (*PostgresDriver, error) {
+	return OpenRetrying(ctx, dsn, 0, 0, nil)
+}
+
+// OpenRetrying is Open for a pod that has only just started. Until window has passed
+// it retries, every interval, a ping the server never answered (a refused, reset or
+// timed-out dial), reporting each retry to onRetry. k3s admits a new pod's address to
+// a NetworkPolicy's allow set a moment after the pod starts: a drill measured about
+// 70ms of refused dials, and the reaper, which dials within milliseconds of starting,
+// failed every pod of every run on `connection refused` against a healthy database.
+// An error the server answered with is final, except 57P03 (cannot_connect_now: it is
+// starting up or shutting down).
+func OpenRetrying(ctx context.Context, dsn string, window, interval time.Duration, onRetry func(error)) (*PostgresDriver, error) {
 	db, err := newPool(dsn)
 	if err != nil {
 		return nil, fmt.Errorf("open postgres: %w", err)
 	}
-	if err := db.PingContext(ctx); err != nil {
-		db.Close()
-		return nil, fmt.Errorf("ping postgres: %w", err)
+	deadline := time.Now().Add(window)
+	for {
+		err := db.PingContext(ctx)
+		if err == nil {
+			return &PostgresDriver{db: db}, nil
+		}
+		if !unansweredDial(err) || !time.Now().Before(deadline) || ctx.Err() != nil {
+			db.Close()
+			return nil, fmt.Errorf("ping postgres: %w", err)
+		}
+		onRetry(err)
+		select {
+		case <-ctx.Done():
+			db.Close()
+			return nil, fmt.Errorf("ping postgres: %w", err)
+		case <-time.After(interval):
+		}
 	}
-	return &PostgresDriver{db: db}, nil
+}
+
+// unansweredDial reports whether a ping failed before the server said anything, or
+// with the one answer that means "try again shortly".
+func unansweredDial(err error) bool {
+	var pgErr *pgconn.PgError
+	if errors.As(err, &pgErr) {
+		return pgErr.Code == "57P03" // cannot_connect_now
+	}
+	return true
 }
 
 // DB exposes the underlying pool for the access layer.