Repository: benbjohnson/litestream
Pull request #287 (2022-01-28 22:07:41)

Use sqlite3_file_control(SQLITE_FCNTL_PERSIST_WAL) to persist WAL
Previously, Litestream would avoid closing the SQLite3 connection in order to ensure that the WAL file was not cleaned up by the database if it was the last connection. This commit changes the behavior by introducing a file control call to perform the same action. This allows us to close the database file normally in all cases.


Patch:
diff --git a/db.go b/db.go
index ed5a97f..49fbf21 100644
--- a/db.go
+++ b/db.go
@@ -405,20 +405,9 @@ func (db *DB) Open() (err error) {
 	return nil
 }
 
-// Close releases the read lock & closes the database. This method should only
-// be called by tests as it causes the underlying database to be checkpointed.
+// Close flushes outstanding WAL writes to replicas, releases the read lock,
+// and closes the database.
 func (db *DB) Close() (err error) {
-	return db.close(false)
-}
-
-// SoftClose closes everything but the underlying db connection. This method
-// is available because the binary needs to avoid closing the database on exit
-// to prevent autocheckpointing.
-func (db *DB) SoftClose() (err error) {
-	return db.close(true)
-}
-
-func (db *DB) close(soft bool) (err error) {
 	db.cancel()
 	db.wg.Wait()
 
@@ -439,7 +428,7 @@ func (db *DB) close(soft bool) (err error) {
 				err = e
 			}
 		}
-		r.Stop(!soft)
+		r.Stop(true)
 	}
 
 	// Release the read lock to allow other applications to handle checkpointing.
@@ -449,9 +438,7 @@ func (db *DB) close(soft bool) (err error) {
 		}
 	}
 
-	// Only perform full close if this is not a soft close.
-	// This closes the underlying database connection which can clean up the WAL.
-	if !soft && db.db != nil {
+	if db.db != nil {
 		if e := db.db.Close(); e != nil && err == nil {
 			err = e
 		}
@@ -507,8 +494,9 @@ func (db *DB) init() (err error) {
 	dsn := db.path
 	dsn += fmt.Sprintf("?_busy_timeout=%d", BusyTimeout.Milliseconds())
 
-	// Connect to SQLite database.
-	if db.db, err = sql.Open("sqlite3", dsn); err != nil {
+	// Connect to SQLite database. Use the driver registered with a hook to
+	// prevent WAL files from being removed.
+	if db.db, err = sql.Open("litestream-sqlite3", dsn); err != nil {
 		return err
 	}
 
@@ -1536,7 +1524,7 @@ func ApplyWAL(ctx context.Context, dbPath, walPath string) error {
 	}
 
 	// Open SQLite database and force a truncating checkpoint.
-	d, err := sql.Open("sqlite3", dbPath)
+	d, err := sql.Open("litestream-sqlite3", dbPath)
 	if err != nil {
 		return err
 	}
diff --git a/go.mod b/go.mod
index 286a1d0..169ea19 100644
--- a/go.mod
+++ b/go.mod
@@ -8,7 +8,7 @@ require (
 	github.com/Azure/azure-storage-blob-go v0.14.0
 	github.com/aws/aws-sdk-go v1.42.40
 	github.com/mattn/go-shellwords v1.0.12
-	github.com/mattn/go-sqlite3 v1.14.10
+	github.com/mattn/go-sqlite3 v1.14.11
 	github.com/pierrec/lz4/v4 v4.1.12
 	github.com/pkg/sftp v1.13.4
 	github.com/prometheus/client_golang v1.12.0
diff --git a/go.sum b/go.sum
index ef3f492..7c57594 100644
--- a/go.sum
+++ b/go.sum
@@ -236,6 +236,8 @@ github.com/mattn/go-shellwords v1.0.12 h1:M2zGm7EW6UQJvDeQxo4T51eKPurbeFbe8WtebG
 github.com/mattn/go-shellwords v1.0.12/go.mod h1:EZzvwXDESEeg03EKmM+RmDnNOPKG4lLtQsUlTZDWQ8Y=
 github.com/mattn/go-sqlite3 v1.14.10 h1:MLn+5bFRlWMGoSRmJour3CL1w/qL96mvipqpwQW/Sfk=
 github.com/mattn/go-sqlite3 v1.14.10/go.mod h1:NyWgC/yNuGj7Q9rpYnZvas74GogHl5/Z4A/KQRfk6bU=
+github.com/mattn/go-sqlite3 v1.14.11 h1:gt+cp9c0XGqe9S/wAHTL3n/7MqY+siPWgWJgqdsFrzQ=
+github.com/mattn/go-sqlite3 v1.14.11/go.mod h1:NyWgC/yNuGj7Q9rpYnZvas74GogHl5/Z4A/KQRfk6bU=
 github.com/matttproud/golang_protobuf_extensions v1.0.1 h1:4hp9jkHxhMHkqkrB3Ix0jegS5sx/RkqARlsWZ6pIwiU=
 github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=
 github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=


Test patch:
diff --git a/cmd/litestream/replicate.go b/cmd/litestream/replicate.go
index 7f9a9aa..fa84907 100644
--- a/cmd/litestream/replicate.go
+++ b/cmd/litestream/replicate.go
@@ -181,7 +181,7 @@ func (c *ReplicateCommand) Run(ctx context.Context) (err error) {
 // Close closes all open databases.
 func (c *ReplicateCommand) Close() (err error) {
 	for _, db := range c.DBs {
-		if e := db.SoftClose(); e != nil {
+		if e := db.Close(); e != nil {
 			log.Printf("error closing db: path=%s err=%s", db.Path(), e)
 			if err == nil {
 				err = e
diff --git a/db_test.go b/db_test.go
index fe74225..a9dbb58 100644
--- a/db_test.go
+++ b/db_test.go
@@ -254,8 +254,8 @@ func TestDB_Sync(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		// Verify WAL does not exist.
-		if _, err := os.Stat(db.WALPath()); !os.IsNotExist(err) {
+		// Remove WAL file.
+		if err := os.Remove(db.WALPath()); err != nil {
 			t.Fatal(err)
 		}
 
diff --git a/litestream.go b/litestream.go
index 98f94a8..5381e05 100644
--- a/litestream.go
+++ b/litestream.go
@@ -13,6 +13,8 @@ import (
 	"strconv"
 	"strings"
 	"time"
+
+	"github.com/mattn/go-sqlite3"
 )
 
 // Naming constants.
@@ -51,6 +53,17 @@ var (
 	LogFlags = 0
 )
 
+func init() {
+	sql.Register("litestream-sqlite3", &sqlite3.SQLiteDriver{
+		ConnectHook: func(conn *sqlite3.SQLiteConn) error {
+			if err := conn.SetFileControlInt("main", sqlite3.SQLITE_FCNTL_PERSIST_WAL, 1); err != nil {
+				return fmt.Errorf("cannot set file control: %w", err)
+			}
+			return nil
+		},
+	})
+}
+
 // SnapshotIterator represents an iterator over a collection of snapshot metadata.
 type SnapshotIterator interface {
 	io.Closer