[sqlite] Add shutdown hook

cmd/omniwitness/monolith.go

@@ -138,13 +138,22 @@ func main() {
 	if len(*dbFile) > 0 {
 		// Start up local database.
 		klog.Infof("Connecting to local DB at %q", *dbFile)
-		p = psql.New(psql.Opts{
+		ps, shutdown, err := psql.New(ctx, psql.Opts{
 			Path:         *dbFile,
 			MaxOpenConns: *dbMaxConns,
 		})
-		if err := p.Init(ctx); err != nil {
+		if err != nil {
+			klog.Exitf("Failed to construct SQL persistence: %v", err)
+		}
+		if err := ps.Init(ctx); err != nil {
 			klog.Exitf("Failed to init SQL persistence: %v", err)
 		}
+		p = ps
+		defer func() {
+			if err := shutdown(); err != nil {
+				klog.Errorf("Persistence shutdown failed: %v", err)
+			}
+		}()
 	} else {
 		klog.Warning("No persistence configured for witness. Reboots will lose guarantees of witness correctness. Use --db_file for production deployments.")
 		p = inmemory.New()

persistence/sqlite/sql.go

@@ -20,7 +20,6 @@ import (
 	"database/sql"
 	"fmt"
 	"iter"
-	"sync"
 
 	"github.com/transparency-dev/formats/log"
 	"github.com/transparency-dev/formats/note"
@@ -41,43 +40,40 @@ type Opts struct {
 }
 
 // New returns a persistence object that is backed by the provided database.
-func New(opts Opts) *Persistence {
-	return &Persistence{
-		opts: opts,
-		dbInit: &sync.Mutex{},
+//
+// The returned shutdown func should be called once the persistence instance is no longer required.
+func New(ctx context.Context, opts Opts) (*Persistence, func() error, error) {
+	// Open database with some flags:
+	// - use WAL mode as this allows for read concurrency while writes are happening.
+	// - set a busy_timeout so that sqlite will queue write transactions rather than immediately return ErrBusy
+	// - set synchronous to FULL to ensure durability of commitments
+	db, err := sql.Open("sqlite", fmt.Sprintf("%s?_pragma=journal_mode(WAL)&_pragma=busy_timeout(1000)&_pragma=synchronous(FULL)", opts.Path))
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to connect to DB: %v", err)
 	}
+	if opts.MaxOpenConns != 0 {
+		db.SetMaxOpenConns(opts.MaxOpenConns)
+	}
+	r := &Persistence{
+		db: db,
+	}
+	if err := r.createTablesIfNotExist(ctx); err != nil {
+		return nil, nil, fmt.Errorf("failed to create tables: %v", err)
+	}
+	return r, db.Close, nil
 }
 
 // Persistence is an implementation of witness.Persistence which knows how to use an sqlite
 // database to safely store witness state.
 type Persistence struct {
-	opts Opts
-
-	dbInit *sync.Mutex
 	db *sql.DB
 }
 
 func (p *Persistence) Init(ctx context.Context) error {
-	p.dbInit.Lock()
-	defer p.dbInit.Unlock()
-
-	if p.db != nil {
-		return nil
-	}
-
-	// Open database with some flags:
-	// - use WAL mode as this allows for read concurrency while writes are happening.
-	// - set a busy_timeout so that sqlite will queue write transactions rather than immediately return ErrBusy
-	// - set synchronous to FULL to ensure durability of commitments
-	db, err := sql.Open("sqlite", fmt.Sprintf("%s?_pragma=journal_mode(WAL)&_pragma=busy_timeout(1000)&_pragma=synchronous(FULL)", p.opts.Path))
-	if err != nil {
-		return fmt.Errorf("failed to connect to DB: %v", err)
-	}
-	if p.opts.MaxOpenConns != 0 {
-		db.SetMaxOpenConns(p.opts.MaxOpenConns)
-	}
-	p.db = db
+	return nil
+}
 
+func (p *Persistence) createTablesIfNotExist(ctx context.Context) error {
 	for _, ddl := range []string{
 		"CREATE TABLE IF NOT EXISTS chkpts (logID BLOB PRIMARY KEY, chkpt BLOB)",
 		"CREATE TABLE IF NOT EXISTS logs (logID BLOB PRIMARY KEY, origin STRING NOT NULL, vkey STRING NOT NULL, contact STRING, qpd FLOAT64, disabled BOOL)",

persistence/sqlite/sql_test.go

@@ -29,23 +29,25 @@ import (
 )
 
 func TestUpdate(t *testing.T) {
+	t.Helper()
 	ptest.TestUpdate(t, func() (*Persistence, func() error) {
-		p := New(Opts{Path: ":memory:", MaxOpenConns: 1})
-		return p, func() error { 
-			if p.db != nil {
-				return p.db.Close()
-			}
-			return nil
+		p, shutdown, err := New(t.Context(), Opts{Path: ":memory:", MaxOpenConns: 1})
+		if err != nil {
+			t.Fatalf("Failed to create new persistence instance: %v", err)
 		}
+		return p, shutdown
 	})
 }
 
 func TestLogConfig(t *testing.T) {
-	p := New(Opts{Path: ":memory:", MaxOpenConns: 1})
+	p, shutdown, err := New(t.Context(), Opts{Path: ":memory:", MaxOpenConns: 1})
+	if err != nil {
+		t.Fatalf("New: %v", err)
+	}
 	if err := p.Init(t.Context()); err != nil {
 		t.Fatalf("Init(): %v", err)
 	}
-	defer func() { _ = p.db.Close() }()
+	defer func() { _ = shutdown() }()
 
 	vkey := "sum.golang.org+033de0ae+Ac4zctda0e5eza+HJyk9SxEdh+s3Ux18htTTAD8OuAn8"
 	logs := []omniwitness.Log{
@@ -109,11 +111,14 @@ func TestLogConfig(t *testing.T) {
 }
 
 func TestDisabledLogs(t *testing.T) {
-	p := New(Opts{Path: ":memory:", MaxOpenConns: 1})
+	p, shutdown, err := New(t.Context(), Opts{Path: ":memory:", MaxOpenConns: 1})
+	if err != nil {
+		t.Fatalf("New: %v", err)
+	}
 	if err := p.Init(t.Context()); err != nil {
 		t.Fatalf("Init(): %v", err)
 	}
-	defer func() { _ = p.db.Close() }()
+	defer func() { _ = shutdown() }()
 
 	vkey := "sum.golang.org+033de0ae+Ac4zctda0e5eza+HJyk9SxEdh+s3Ux18htTTAD8OuAn8"
 
@@ -145,11 +150,14 @@ func TestDisabledLogs(t *testing.T) {
 }
 
 func TestDeadlock(t *testing.T) {
-	p := New(Opts{Path: ":memory:", MaxOpenConns: 1})
+	p, shutdown, err := New(t.Context(), Opts{Path: ":memory:", MaxOpenConns: 1})
+	if err != nil {
+		t.Fatalf("New: %v", err)
+	}
 	if err := p.Init(t.Context()); err != nil {
 		t.Fatalf("Init(): %v", err)
 	}
-	defer func() { _ = p.db.Close() }()
+	defer func() { _ = shutdown() }()
 
 	mPK := "monkeys+db4d9f7e+AULaJMvTtDLHPUcUrjdDad9vDlh/PTfC2VV60JUtCfWT"
 	wSK := "PRIVATE+KEY+witness+f13a86db+AaLa/dfyBhyo/m0Z7WCi98ENVZWtrP8pxgRNrx7tIWiA"