mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-20 00:18:55 +00:00
fix(backup): serialize pg_dump/psql with the migrations advisory lock
This commit is contained in:
TT
@@ -210,7 +210,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
dataManagementService := service.NewDataManagementService()
|
||||
dataManagementHandler := admin.NewDataManagementHandler(dataManagementService)
|
||||
backupObjectStoreFactory := repository.NewS3BackupStoreFactory()
|
||||
dbDumper := repository.NewPgDumper(configConfig)
|
||||
dbDumper := repository.NewPgDumper(configConfig, db)
|
||||
backupService := service.ProvideBackupService(settingRepository, configConfig, secretEncryptor, backupObjectStoreFactory, dbDumper, leaderLockCache, db)
|
||||
imageStorageFactory := repository.ProvideImageStorageFactory()
|
||||
imageStorageSettingService := service.ProvideImageStorageSettingService(settingRepository, secretEncryptor, backupService, imageStorageFactory, configConfig)
|
||||
|
||||
@@ -2,9 +2,14 @@ package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"database/sql/driver"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os/exec"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
@@ -12,16 +17,39 @@ import (
|
||||
|
||||
// PgDumper implements service.DBDumper using pg_dump/psql
|
||||
type PgDumper struct {
|
||||
cfg *config.DatabaseConfig
|
||||
cfg *config.DatabaseConfig
|
||||
db *sql.DB
|
||||
commandContext func(context.Context, string, ...string) *exec.Cmd
|
||||
}
|
||||
|
||||
// NewPgDumper creates a new PgDumper
|
||||
func NewPgDumper(cfg *config.Config) service.DBDumper {
|
||||
return &PgDumper{cfg: &cfg.Database}
|
||||
func NewPgDumper(cfg *config.Config, db *sql.DB) service.DBDumper {
|
||||
return &PgDumper{
|
||||
cfg: &cfg.Database,
|
||||
db: db,
|
||||
commandContext: exec.CommandContext,
|
||||
}
|
||||
}
|
||||
|
||||
// Dump executes pg_dump and returns a streaming reader of the output
|
||||
func (d *PgDumper) Dump(ctx context.Context) (io.ReadCloser, error) {
|
||||
if d.db == nil {
|
||||
return nil, errors.New("acquire backup migration lock: nil sql db")
|
||||
}
|
||||
lockConn, err := d.db.Conn(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("acquire backup migration lock connection: %w", err)
|
||||
}
|
||||
if err := pgAdvisoryLock(ctx, lockConn); err != nil {
|
||||
// Lock acquisition can fail after PostgreSQL accepted the request. Discard
|
||||
// the session because ownership is ambiguous and advisory locks are session-scoped.
|
||||
discardSQLConnection(lockConn)
|
||||
return nil, fmt.Errorf("acquire backup migration lock: %w", err)
|
||||
}
|
||||
releaseLock := func() error {
|
||||
return releaseBackupMigrationLock(lockConn)
|
||||
}
|
||||
|
||||
args := []string{
|
||||
"-h", d.cfg.Host,
|
||||
"-p", fmt.Sprintf("%d", d.cfg.Port),
|
||||
@@ -33,7 +61,11 @@ func (d *PgDumper) Dump(ctx context.Context) (io.ReadCloser, error) {
|
||||
"--if-exists",
|
||||
}
|
||||
|
||||
cmd := exec.CommandContext(ctx, "pg_dump", args...)
|
||||
commandContext := d.commandContext
|
||||
if commandContext == nil {
|
||||
commandContext = exec.CommandContext
|
||||
}
|
||||
cmd := commandContext(ctx, "pg_dump", args...)
|
||||
if d.cfg.Password != "" {
|
||||
cmd.Env = append(cmd.Environ(), "PGPASSWORD="+d.cfg.Password)
|
||||
}
|
||||
@@ -43,15 +75,34 @@ func (d *PgDumper) Dump(ctx context.Context) (io.ReadCloser, error) {
|
||||
|
||||
stdout, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create stdout pipe: %w", err)
|
||||
return nil, errors.Join(fmt.Errorf("create stdout pipe: %w", err), releaseLock())
|
||||
}
|
||||
|
||||
if err := cmd.Start(); err != nil {
|
||||
return nil, fmt.Errorf("start pg_dump: %w", err)
|
||||
_ = stdout.Close()
|
||||
return nil, errors.Join(fmt.Errorf("start pg_dump: %w", err), releaseLock())
|
||||
}
|
||||
|
||||
// 返回一个 ReadCloser：读 stdout，关闭时等待进程退出
|
||||
return &cmdReadCloser{ReadCloser: stdout, cmd: cmd}, nil
|
||||
return &cmdReadCloser{ReadCloser: stdout, cmd: cmd, release: releaseLock}, nil
|
||||
}
|
||||
|
||||
func releaseBackupMigrationLock(conn *sql.Conn) error {
|
||||
unlockCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
if err := pgAdvisoryUnlock(unlockCtx, conn); err != nil {
|
||||
// A failed unlock must not return a possibly lock-owning session to the pool.
|
||||
discardSQLConnection(conn)
|
||||
return fmt.Errorf("release backup migration lock: %w", err)
|
||||
}
|
||||
if err := conn.Close(); err != nil {
|
||||
return fmt.Errorf("close backup migration lock connection: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func discardSQLConnection(conn *sql.Conn) {
|
||||
_ = conn.Raw(func(any) error { return driver.ErrBadConn })
|
||||
_ = conn.Close()
|
||||
}
|
||||
|
||||
// Restore executes psql to restore from a streaming reader
|
||||
@@ -84,15 +135,27 @@ func (d *PgDumper) Restore(ctx context.Context, data io.Reader) error {
|
||||
// cmdReadCloser wraps a command stdout pipe and waits for the process on Close
|
||||
type cmdReadCloser struct {
|
||||
io.ReadCloser
|
||||
cmd *exec.Cmd
|
||||
cmd *exec.Cmd
|
||||
release func() error
|
||||
closeOnce sync.Once
|
||||
closeErr error
|
||||
}
|
||||
|
||||
func (c *cmdReadCloser) Close() error {
|
||||
// Close the pipe first
|
||||
_ = c.ReadCloser.Close()
|
||||
// Wait for the process to exit
|
||||
if err := c.cmd.Wait(); err != nil {
|
||||
return fmt.Errorf("pg_dump exited with error: %w", err)
|
||||
}
|
||||
return nil
|
||||
c.closeOnce.Do(func() {
|
||||
var closeErrs []error
|
||||
if err := c.ReadCloser.Close(); err != nil {
|
||||
closeErrs = append(closeErrs, fmt.Errorf("close pg_dump stdout: %w", err))
|
||||
}
|
||||
if err := c.cmd.Wait(); err != nil {
|
||||
closeErrs = append(closeErrs, fmt.Errorf("pg_dump exited with error: %w", err))
|
||||
}
|
||||
if c.release != nil {
|
||||
if err := c.release(); err != nil {
|
||||
closeErrs = append(closeErrs, err)
|
||||
}
|
||||
}
|
||||
c.closeErr = errors.Join(closeErrs...)
|
||||
})
|
||||
return c.closeErr
|
||||
}
|
||||
|
||||
@@ -0,0 +1,157 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"os/exec"
|
||||
"regexp"
|
||||
"testing"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func newTestPgDumper(t *testing.T, commandContext func(context.Context, string, ...string) *exec.Cmd) (*PgDumper, sqlmock.Sqlmock) {
|
||||
t.Helper()
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
return &PgDumper{
|
||||
cfg: &config.DatabaseConfig{
|
||||
Host: "db.example.test",
|
||||
Port: 5432,
|
||||
User: "sub2api",
|
||||
Password: "secret",
|
||||
DBName: "sub2api",
|
||||
SSLMode: "require",
|
||||
},
|
||||
db: db,
|
||||
commandContext: commandContext,
|
||||
}, mock
|
||||
}
|
||||
|
||||
func expectBackupMigrationLock(mock sqlmock.Sqlmock) {
|
||||
mock.ExpectQuery(regexp.QuoteMeta("SELECT pg_try_advisory_lock($1)")).
|
||||
WithArgs(migrationsAdvisoryLockID).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"pg_try_advisory_lock"}).AddRow(true))
|
||||
}
|
||||
|
||||
func expectBackupMigrationUnlock(mock sqlmock.Sqlmock) {
|
||||
mock.ExpectExec(regexp.QuoteMeta("SELECT pg_advisory_unlock($1)")).
|
||||
WithArgs(migrationsAdvisoryLockID).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
}
|
||||
|
||||
func TestPgDumperHoldsMigrationLockThroughReaderClose(t *testing.T) {
|
||||
var mock sqlmock.Sqlmock
|
||||
commandCreated := false
|
||||
dumper, createdMock := newTestPgDumper(t, func(ctx context.Context, name string, args ...string) *exec.Cmd {
|
||||
commandCreated = true
|
||||
require.Equal(t, "pg_dump", name)
|
||||
require.Contains(t, args, "--clean")
|
||||
require.NoError(t, mock.ExpectationsWereMet(), "migration lock must be acquired before pg_dump is created")
|
||||
return exec.CommandContext(ctx, "sh", "-c", "printf backup-data")
|
||||
})
|
||||
mock = createdMock
|
||||
expectBackupMigrationLock(mock)
|
||||
|
||||
reader, err := dumper.Dump(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.True(t, commandCreated)
|
||||
data, err := io.ReadAll(reader)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "backup-data", string(data))
|
||||
|
||||
expectBackupMigrationUnlock(mock)
|
||||
require.Error(t, mock.ExpectationsWereMet(), "migration lock was released before the reader closed")
|
||||
require.NoError(t, reader.Close())
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
require.NoError(t, reader.Close(), "reader close must be idempotent")
|
||||
}
|
||||
|
||||
func TestPgDumperReleasesMigrationLockWhenStdoutPipeSetupFails(t *testing.T) {
|
||||
dumper, mock := newTestPgDumper(t, func(ctx context.Context, _ string, _ ...string) *exec.Cmd {
|
||||
cmd := exec.CommandContext(ctx, "sh", "-c", "true")
|
||||
cmd.Stdout = io.Discard
|
||||
return cmd
|
||||
})
|
||||
expectBackupMigrationLock(mock)
|
||||
expectBackupMigrationUnlock(mock)
|
||||
|
||||
reader, err := dumper.Dump(context.Background())
|
||||
require.Nil(t, reader)
|
||||
require.ErrorContains(t, err, "create stdout pipe")
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestPgDumperReleasesMigrationLockWhenProcessStartFails(t *testing.T) {
|
||||
dumper, mock := newTestPgDumper(t, func(ctx context.Context, _ string, _ ...string) *exec.Cmd {
|
||||
return exec.CommandContext(ctx, "/path/that/does/not/exist/pg_dump")
|
||||
})
|
||||
expectBackupMigrationLock(mock)
|
||||
expectBackupMigrationUnlock(mock)
|
||||
|
||||
reader, err := dumper.Dump(context.Background())
|
||||
require.Nil(t, reader)
|
||||
require.ErrorContains(t, err, "start pg_dump")
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestPgDumperReleasesMigrationLockWhenProcessFails(t *testing.T) {
|
||||
dumper, mock := newTestPgDumper(t, func(ctx context.Context, _ string, _ ...string) *exec.Cmd {
|
||||
return exec.CommandContext(ctx, "sh", "-c", "printf partial-backup; exit 7")
|
||||
})
|
||||
expectBackupMigrationLock(mock)
|
||||
|
||||
reader, err := dumper.Dump(context.Background())
|
||||
require.NoError(t, err)
|
||||
data, err := io.ReadAll(reader)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "partial-backup", string(data))
|
||||
expectBackupMigrationUnlock(mock)
|
||||
require.ErrorContains(t, reader.Close(), "pg_dump exited with error")
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestPgDumperReportsUnlockFailureAndDiscardsConnection(t *testing.T) {
|
||||
dumper, mock := newTestPgDumper(t, func(ctx context.Context, _ string, _ ...string) *exec.Cmd {
|
||||
return exec.CommandContext(ctx, "sh", "-c", "printf backup-data")
|
||||
})
|
||||
expectBackupMigrationLock(mock)
|
||||
|
||||
reader, err := dumper.Dump(context.Background())
|
||||
require.NoError(t, err)
|
||||
_, err = io.ReadAll(reader)
|
||||
require.NoError(t, err)
|
||||
mock.ExpectExec(regexp.QuoteMeta("SELECT pg_advisory_unlock($1)")).
|
||||
WithArgs(migrationsAdvisoryLockID).
|
||||
WillReturnError(errors.New("unlock unavailable"))
|
||||
require.ErrorContains(t, reader.Close(), "release backup migration lock")
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestPgDumperDoesNotStartProcessWhenMigrationLockFails(t *testing.T) {
|
||||
commandCreated := false
|
||||
dumper, mock := newTestPgDumper(t, func(ctx context.Context, _ string, _ ...string) *exec.Cmd {
|
||||
commandCreated = true
|
||||
return exec.CommandContext(ctx, "sh", "-c", "true")
|
||||
})
|
||||
mock.ExpectQuery(regexp.QuoteMeta("SELECT pg_try_advisory_lock($1)")).
|
||||
WithArgs(migrationsAdvisoryLockID).
|
||||
WillReturnError(errors.New("database unavailable"))
|
||||
|
||||
reader, err := dumper.Dump(context.Background())
|
||||
require.Nil(t, reader)
|
||||
require.ErrorContains(t, err, "acquire backup migration lock")
|
||||
require.False(t, commandCreated)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestPgDumperRejectsNilDatabase(t *testing.T) {
|
||||
dumper := &PgDumper{cfg: &config.DatabaseConfig{}}
|
||||
reader, err := dumper.Dump(context.Background())
|
||||
require.Nil(t, reader)
|
||||
require.ErrorContains(t, err, "nil sql db")
|
||||
}
|
||||
Reference in New Issue
Block a user