diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index 5ec4892e16..8f4a0482bd 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -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) diff --git a/backend/internal/repository/backup_pg_dumper.go b/backend/internal/repository/backup_pg_dumper.go index e9a92ef29d..01510bf831 100644 --- a/backend/internal/repository/backup_pg_dumper.go +++ b/backend/internal/repository/backup_pg_dumper.go @@ -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 } diff --git a/backend/internal/repository/backup_pg_dumper_test.go b/backend/internal/repository/backup_pg_dumper_test.go new file mode 100644 index 0000000000..586182b70f --- /dev/null +++ b/backend/internal/repository/backup_pg_dumper_test.go @@ -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") +}