Improve closure of sqlite backend (#63828)

The lifetime of `runPeriodicOperations` was only loosely tied to
the lite.Backend via lite.Backend.ctx. When closing the backend
the context was cancelled which would eventually terminate
`runPeriodicOperations`. However, `lite.Backend.Close` could return
while `runPeriodicOperations` was still processing as a result
of the periodic ticker firing. To prevent this race the backend
now utilizies a sync.WaitGroup to launch `runPeriodicOperations`
and waits for the group to be completed _prior_ to returning from
`(lite.Backend) Close`.

`runPeriodicOperations` also attempted to close the sql.DB when it
detected that the context had been cancelled. While the double close
is fine due to the idempotence of (lite.Backed) closeDatabase, it
was confusing and unnecesarry. Ownership of closing the database now
lies exclusively on explicitly calling `(lite.Backend) Close`.

The closed flag has also been removed. It existed solely to prevent
logging or rolling back transactions that had failed after the
underlying database had been closed. Both operations are always attempted
now. All logs are silenced in the event that the database is unreachable
during process termination.

On failure to create the database schema the underlying sqlite database
is now closed before returning an error from lite.NewWithConfig.

lib/backend/lite/lite.go

@@ -32,7 +32,7 @@ import (
 	"path/filepath"
 	"runtime/debug"
 	"strconv"
-	"sync/atomic"
+	"sync"
 	"time"
 
 	"github.com/gravitational/trace"
@@ -230,21 +230,22 @@ func NewWithConfig(ctx context.Context, cfg Config) (*Backend, error) {
 		return nil, trace.Wrap(err, "error opening URI: %v", connectionURI)
 	}
 
+	// serialize access to sqlite, as we're using immediate transactions anyway,
+	// and in-memory go locks are faster than sqlite locks
+	db.SetMaxOpenConns(1)
+
 	if setPermissions {
 		// Ensure the database has restrictive access permissions.
-		err = db.PingContext(ctx)
-		if err != nil {
+		if err := db.PingContext(ctx); err != nil {
+			_ = db.Close()
 			return nil, trace.Wrap(err)
 		}
-		err = os.Chmod(path, dbMode)
-		if err != nil {
+		if err := os.Chmod(path, dbMode); err != nil {
+			_ = db.Close()
 			return nil, trace.ConvertSystemError(err)
 		}
 	}
 
-	// serialize access to sqlite, as we're using immediate transactions anyway,
-	// and in-memory go locks are faster than sqlite locks
-	db.SetMaxOpenConns(1)
 	buf := backend.NewCircularBuffer(
 		backend.BufferCapacity(cfg.BufferSize),
 	)
@@ -260,12 +261,13 @@ func NewWithConfig(ctx context.Context, cfg Config) (*Backend, error) {
 	}
 	l.logger.DebugContext(ctx, "Connected to database", "database", connectionURI, "poll_stream_period", cfg.PollStreamPeriod)
 	if err := l.createSchema(); err != nil {
-		return nil, trace.Wrap(err, "error creating schema: %v", connectionURI)
+		cancel()
+		return nil, trace.NewAggregate(trace.Wrap(err, "error creating schema: %v", connectionURI), db.Close())
 	}
 	if err := l.showPragmas(); err != nil {
 		l.logger.WarnContext(ctx, "Failed to show pragma settings", "error", err)
 	}
-	go l.runPeriodicOperations()
+	l.wg.Go(l.runPeriodicOperations)
 	return l, nil
 }
 
@@ -282,8 +284,7 @@ type Backend struct {
 	ctx    context.Context
 	cancel context.CancelFunc
 
-	// closedFlag is set to indicate that the database is closed
-	closedFlag int32
+	wg sync.WaitGroup
 }
 
 func (l *Backend) GetName() string {
@@ -992,7 +993,9 @@ func (l *Backend) NewWatcher(ctx context.Context, watch backend.Watch) (backend.
 // Close closes all associated resources
 func (l *Backend) Close() error {
 	l.cancel()
-	return l.closeDatabase()
+	l.buf.Close()
+	l.wg.Wait()
+	return trace.Wrap(l.db.Close())
 }
 
 // CloseWatchers closes all the watchers
@@ -1001,20 +1004,6 @@ func (l *Backend) CloseWatchers() {
 	l.buf.Clear()
 }
 
-func (l *Backend) isClosed() bool {
-	return atomic.LoadInt32(&l.closedFlag) == 1
-}
-
-func (l *Backend) setClosed() {
-	atomic.StoreInt32(&l.closedFlag, 1)
-}
-
-func (l *Backend) closeDatabase() error {
-	l.setClosed()
-	l.buf.Close()
-	return l.db.Close()
-}
-
 func (l *Backend) inTransaction(ctx context.Context, f func(tx *sql.Tx) error) (err error) {
 	start := time.Now()
 	defer func() {
@@ -1023,50 +1012,44 @@ func (l *Backend) inTransaction(ctx context.Context, f func(tx *sql.Tx) error) (
 			l.logger.WarnContext(ctx, "SLOW TRANSACTION", "duration", diff, "stack", string(debug.Stack()))
 		}
 	}()
+
 	tx, err := l.db.BeginTx(ctx, nil)
 	if err != nil {
 		return trace.Wrap(convertError(err))
 	}
-	commit := func() error {
-		return tx.Commit()
-	}
-	rollback := func() error {
-		return tx.Rollback()
-	}
 	defer func() {
 		if r := recover(); r != nil {
 			l.logger.ErrorContext(ctx, "Unexpected panic in inTransaction, trying to rollback.", "error", r)
 			err = trace.BadParameter("panic: %v", r)
-			if e2 := rollback(); e2 != nil {
+			if e2 := tx.Rollback(); e2 != nil {
 				l.logger.ErrorContext(ctx, "Failed to rollback", "error", e2)
 			}
 			return
 		}
-		if err != nil && !trace.IsNotFound(err) {
-			if isConstraintError(trace.Unwrap(err)) {
-				err = trace.AlreadyExists("%s", err)
-			}
-			// transaction aborted by interrupt, no action needed
-			if isInterrupt(trace.Unwrap(err)) {
-				return
-			}
+
+		if err != nil {
 			if isLockedError(trace.Unwrap(err)) {
 				err = trace.ConnectionProblem(err, "database is locked")
 			}
 			if isReadonlyError(trace.Unwrap(err)) {
 				err = trace.ConnectionProblem(err, "database is in readonly mode")
 			}
-			if !l.isClosed() {
-				if !trace.IsCompareFailed(err) && !trace.IsAlreadyExists(err) && !trace.IsConnectionProblem(err) {
-					l.logger.WarnContext(ctx, "Unexpected error in inTransaction, rolling back.", "error", err)
-				}
-				if e2 := rollback(); e2 != nil {
-					l.logger.ErrorContext(ctx, "Failed to rollback too", "error", e2)
-				}
+
+			if isConstraintError(trace.Unwrap(err)) {
+				err = trace.AlreadyExists("%s", err)
+			}
+
+			if e2 := tx.Rollback(); e2 != nil && !trace.IsConnectionProblem(err) && l.ctx.Err() == nil {
+				l.logger.ErrorContext(ctx, "Failed to rollback too", "error", e2)
+			}
+
+			if !trace.IsNotFound(err) && !trace.IsCompareFailed(err) && !trace.IsAlreadyExists(err) && !trace.IsConnectionProblem(err) && l.ctx.Err() == nil {
+				l.logger.WarnContext(ctx, "Unexpected error in inTransaction, rolling back.", "error", err)
 			}
 			return
 		}
-		if err2 := commit(); err2 != nil {
+
+		if err2 := tx.Commit(); err2 != nil {
 			err = trace.Wrap(err2)
 		}
 	}()
@@ -1109,14 +1092,6 @@ func isLockedError(err error) bool {
 	return e.Code() == sqlite3.SQLITE_BUSY
 }
 
-func isInterrupt(err error) bool {
-	var e *sqlite.Error
-	if ok := errors.As(err, &e); !ok {
-		return false
-	}
-	return e.Code() == sqlite3.SQLITE_INTERRUPT
-}
-
 func isReadonlyError(err error) bool {
 	var e *sqlite.Error
 	if ok := errors.As(err, &e); !ok {

lib/backend/lite/periodic.go

@@ -38,29 +38,23 @@ func (l *Backend) runPeriodicOperations() {
 	for {
 		select {
 		case <-l.ctx.Done():
-			if err := l.closeDatabase(); err != nil {
-				l.logger.WarnContext(l.ctx, "Error closing database", "error", err)
-			}
 			return
 		case <-t.C:
-			err := l.removeExpiredKeys()
-			if err != nil {
+			if err := l.removeExpiredKeys(); err != nil {
 				// connection problem means that database is closed
-				// or is closing, downgrade the log to debug
-				// to avoid polluting logs in production
-				if trace.IsConnectionProblem(err) {
-					l.logger.DebugContext(l.ctx, "Failed to run remove expired keys", "error", err)
-				} else {
+				// or is closing, omit the log to avoid polluting logs in production
+				if !trace.IsConnectionProblem(err) && l.ctx.Err() == nil {
 					l.logger.DebugContext(l.ctx, "Failed to run remove expired keys", "error", err)
 				}
 			}
 			if !l.EventsOff {
-				err = l.removeOldEvents()
-				if err != nil {
+				err := l.removeOldEvents()
+				if err != nil && !trace.IsConnectionProblem(err) && l.ctx.Err() == nil {
 					l.logger.WarnContext(l.ctx, "Failed to run remove old events", "error", err)
 				}
+
 				rowid, err = l.pollEvents(rowid)
-				if err != nil {
+				if err != nil && !trace.IsConnectionProblem(err) && l.ctx.Err() == nil {
 					l.logger.WarnContext(l.ctx, "Failed to run poll events", "error", err)
 				}
 			}