rkrtechnology/litestream
2
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Allow giving context to database close (#520)

Browse Source
This commit is contained in:
Toni Spets
2023-11-14 19:28:58 +02:00
committed by GitHub
parent 977d4a5ee4
commit ae4c9918d9
3 changed files with 10 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
cmd/litestream/replicate.go
View File

@@ -168,7 +168,7 @@ func (c *ReplicateCommand) Run() (err error) {
// Close closes all open databases.
func (c *ReplicateCommand) Close() (err error) {
for _, db := range c.DBs {
if e := db.Close(); e != nil {
if e := db.Close(context.Background()); e != nil {
db.Logger.Error("error closing db", "error", e)
if err == nil {
err = e

7
db.go
View File

@@ -317,14 +317,11 @@ func (db *DB) Open() (err error) {
}
// Close flushes outstanding WAL writes to replicas, releases the read lock,
// and closes the database.
func (db *DB) Close() (err error) {
// and closes the database. Takes a context for final sync.
func (db *DB) Close(ctx context.Context) (err error) {
db.cancel()
db.wg.Wait()
// Start a new context for shutdown since we canceled the DB context.
ctx := context.Background()
// Perform a final db sync, if initialized.
if db.db != nil {
if e := db.Sync(ctx); e != nil && err == nil {

14
db_test.go
View File

@@ -264,7 +264,7 @@ func TestDB_Sync(t *testing.T) {
// Checkpoint & fully close which should close WAL file.
if err := db.Checkpoint(context.Background(), litestream.CheckpointModeTruncate); err != nil {
t.Fatal(err)
} else if err := db.Close(); err != nil {
} else if err := db.Close(context.Background()); err != nil {
t.Fatal(err)
} else if err := sqldb.Close(); err != nil {
t.Fatal(err)
@@ -314,7 +314,7 @@ func TestDB_Sync(t *testing.T) {
}
// Fully close which should close WAL file.
if err := db.Close(); err != nil {
if err := db.Close(context.Background()); err != nil {
t.Fatal(err)
} else if err := sqldb.Close(); err != nil {
t.Fatal(err)
@@ -367,7 +367,7 @@ func TestDB_Sync(t *testing.T) {
pos0, err := db.Pos()
if err != nil {
t.Fatal(err)
} else if err := db.Close(); err != nil {
} else if err := db.Close(context.Background()); err != nil {
t.Fatal(err)
}
@@ -413,7 +413,7 @@ func TestDB_Sync(t *testing.T) {
}
// Close & truncate shadow WAL to simulate a partial header write.
if err := db.Close(); err != nil {
if err := db.Close(context.Background()); err != nil {
t.Fatal(err)
} else if err := os.Truncate(db.ShadowWALPath(pos0.Generation, pos0.Index), litestream.WALHeaderSize-1); err != nil {
t.Fatal(err)
@@ -458,7 +458,7 @@ func TestDB_Sync(t *testing.T) {
}
// Close & truncate shadow WAL to simulate a partial frame write.
if err := db.Close(); err != nil {
if err := db.Close(context.Background()); err != nil {
t.Fatal(err)
} else if err := os.Truncate(db.ShadowWALPath(pos0.Generation, pos0.Index), fi.Size()-1); err != nil {
t.Fatal(err)
@@ -504,7 +504,7 @@ func TestDB_Sync(t *testing.T) {
}
// Close & delete shadow WAL to simulate dir created but not WAL.
if err := db.Close(); err != nil {
if err := db.Close(context.Background()); err != nil {
t.Fatal(err)
} else if err := os.Remove(db.ShadowWALPath(pos0.Generation, pos0.Index)); err != nil {
t.Fatal(err)
@@ -626,7 +626,7 @@ func MustOpenDBAt(tb testing.TB, path string) *litestream.DB {
// MustCloseDB closes db and removes its parent directory.
func MustCloseDB(tb testing.TB, db *litestream.DB) {
tb.Helper()
if err := db.Close(); err != nil && !strings.Contains(err.Error(), `database is closed`) {
if err := db.Close(context.Background()); err != nil && !strings.Contains(err.Error(), `database is closed`) {
tb.Fatal(err)
} else if err := os.RemoveAll(filepath.Dir(db.Path())); err != nil {
tb.Fatal(err)