Fix: Use proper context in acquireReadLock during DB close (#701)

Co-authored-by: Claude <[email protected]>

Author: corylanou

db.go

@@ -326,7 +326,7 @@ func (db *DB) Close(ctx context.Context) (err error) {
 
 // init initializes the connection to the database.
 // Skipped if already initialized or if the database file does not exist.
-func (db *DB) init() (err error) {
+func (db *DB) init(ctx context.Context) (err error) {
 	// Exit if already initialized.
 	if db.db != nil {
 		return nil
@@ -382,7 +382,7 @@ func (db *DB) init() (err error) {
 	}
 
 	// Disable autocheckpoint for litestream's connection.
-	if _, err := db.db.ExecContext(db.ctx, `PRAGMA wal_autocheckpoint = 0;`); err != nil {
+	if _, err := db.db.ExecContext(ctx, `PRAGMA wal_autocheckpoint = 0;`); err != nil {
 		return fmt.Errorf("disable autocheckpoint: %w", err)
 	}
 
@@ -400,7 +400,7 @@ func (db *DB) init() (err error) {
 
 	// Start a long-running read transaction to prevent other transactions
 	// from checkpointing.
-	if err := db.acquireReadLock(); err != nil {
+	if err := db.acquireReadLock(ctx); err != nil {
 		return fmt.Errorf("acquire read lock: %w", err)
 	}
 
@@ -479,7 +479,7 @@ func (db *DB) verifyHeadersMatch() error {
 */
 
 // acquireReadLock begins a read transaction on the database to prevent checkpointing.
-func (db *DB) acquireReadLock() error {
+func (db *DB) acquireReadLock(ctx context.Context) error {
 	if db.rtx != nil {
 		return nil
 	}
@@ -491,7 +491,7 @@ func (db *DB) acquireReadLock() error {
 	}
 
 	// Execute read query to obtain read lock.
-	if _, err := tx.ExecContext(db.ctx, `SELECT COUNT(1) FROM _litestream_seq;`); err != nil {
+	if _, err := tx.ExecContext(ctx, `SELECT COUNT(1) FROM _litestream_seq;`); err != nil {
 		_ = tx.Rollback()
 		return err
 	}
@@ -520,7 +520,7 @@ func (db *DB) Sync(ctx context.Context) (err error) {
 	defer db.mu.Unlock()
 
 	// Initialize database, if necessary. Exit if no DB exists.
-	if err := db.init(); err != nil {
+	if err := db.init(ctx); err != nil {
 		return err
 	} else if db.db == nil {
 		db.Logger.Debug("sync: no database found")
@@ -1017,7 +1017,7 @@ func (db *DB) checkpoint(ctx context.Context, mode string) error {
 
 	// Execute checkpoint and immediately issue a write to the WAL to ensure
 	// a new page is written.
-	if err := db.execCheckpoint(mode); err != nil {
+	if err := db.execCheckpoint(ctx, mode); err != nil {
 		return err
 	} else if _, err = db.db.Exec(`INSERT INTO _litestream_seq (id, seq) VALUES (1, 1) ON CONFLICT (id) DO UPDATE SET seq = seq + 1`); err != nil {
 		return err
@@ -1057,7 +1057,7 @@ func (db *DB) checkpoint(ctx context.Context, mode string) error {
 	return nil
 }
 
-func (db *DB) execCheckpoint(mode string) (err error) {
+func (db *DB) execCheckpoint(ctx context.Context, mode string) (err error) {
 	// Ignore if there is no underlying database.
 	if db.db == nil {
 		return nil
@@ -1079,7 +1079,7 @@ func (db *DB) execCheckpoint(mode string) (err error) {
 	if err := db.releaseReadLock(); err != nil {
 		return fmt.Errorf("release read lock: %w", err)
 	}
-	defer func() { _ = db.acquireReadLock() }()
+	defer func() { _ = db.acquireReadLock(ctx) }()
 
 	// A non-forced checkpoint is issued as "PASSIVE". This will only checkpoint
 	// if there are not pending transactions. A forced checkpoint ("RESTART")
@@ -1096,8 +1096,8 @@ func (db *DB) execCheckpoint(mode string) (err error) {
 	db.Logger.Debug("checkpoint", "mode", mode, "result", fmt.Sprintf("%d,%d,%d", row[0], row[1], row[2]))
 
 	// Reacquire the read lock immediately after the checkpoint.
-	if err := db.acquireReadLock(); err != nil {
-		return fmt.Errorf("release read lock: %w", err)
+	if err := db.acquireReadLock(ctx); err != nil {
+		return fmt.Errorf("reacquire read lock: %w", err)
 	}
 
 	return nil
@@ -1412,7 +1412,7 @@ func (db *DB) CRC64(ctx context.Context) (uint64, ltx.Pos, error) {
 	db.mu.Lock()
 	defer db.mu.Unlock()
 
-	if err := db.init(); err != nil {
+	if err := db.init(ctx); err != nil {
 		return 0, ltx.Pos{}, err
 	} else if db.db == nil {
 		return 0, ltx.Pos{}, os.ErrNotExist