From ae4c9918d9949f6fe9062dfb2f22f32da0851f0d Mon Sep 17 00:00:00 2001
From: Toni Spets <toni.spets@beeper.com>
Date: Tue, 14 Nov 2023 19:28:58 +0200
Subject: [PATCH] Allow giving context to database close (#520)

---
 cmd/litestream/replicate.go |  2 +-
 db.go                       |  7 ++-----
 db_test.go                  | 14 +++++++-------
 3 files changed, 10 insertions(+), 13 deletions(-)

diff --git a/cmd/litestream/replicate.go b/cmd/litestream/replicate.go
index 898e79d..988407c 100644
--- a/cmd/litestream/replicate.go
+++ b/cmd/litestream/replicate.go
@@ -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
diff --git a/db.go b/db.go
index ba914e1..ac6564f 100644
--- a/db.go
+++ b/db.go
@@ -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 {
diff --git a/db_test.go b/db_test.go
index ef82987..c882858 100644
--- a/db_test.go
+++ b/db_test.go
@@ -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)