migration: add interceptor signal to cleanly shutdown.

Author: ziggie1984

cmd_migrate_db.go

@@ -7,6 +7,7 @@ import (
 	_ "net/http/pprof" // Register pprof handlers
 	"os"
 	"path/filepath"
+	"strings"
 	"syscall"
 	"time"
 
@@ -20,6 +21,7 @@ import (
 	"github.com/lightningnetwork/lnd/kvdb/sqlbase"
 	"github.com/lightningnetwork/lnd/kvdb/sqlite"
 	"github.com/lightningnetwork/lnd/lncfg"
+	"github.com/lightningnetwork/lnd/signal"
 	"github.com/lightningnetwork/lnd/watchtower/wtdb"
 	"go.etcd.io/bbolt"
 )
@@ -196,6 +198,16 @@ func (x *migrateDBCommand) Execute(_ []string) error {
 		}()
 	}
 
+	// Hook interceptor for os signals. We need to except the case where
+	// we call the function multiple times.
+	shutdownInterceptor, err := signal.Intercept()
+	if err != nil && !strings.Contains(err.Error(), "intercept "+
+		"already started") {
+
+		_, _ = fmt.Fprintln(os.Stderr, err)
+		os.Exit(1)
+	}
+
 	for _, prefix := range allDBPrefixes {
 		logger.Infof("Attempting to migrate DB with prefix %s", prefix)
 
@@ -351,7 +363,7 @@ func (x *migrateDBCommand) Execute(_ []string) error {
 			MetaDB:    metaDB,
 		}
 
-		migrator, err := migratekvdb.New(cfg)
+		migrator, err := migratekvdb.New(cfg, shutdownInterceptor)
 		if err != nil {
 			return err
 		}

migratekvdb/migration.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/btcsuite/btclog/v2"
 	"github.com/lightningnetwork/lnd/kvdb"
+	"github.com/lightningnetwork/lnd/signal"
 	"go.etcd.io/bbolt"
 )
 
@@ -114,10 +115,13 @@ type Migrator struct {
 	lastVerifiedKeyCount   int64
 	migratedKeysSinceStart int64
 	verifiedKeysSinceStart int64
+
+	// Interceptor is the interceptor to use for intercepting signals.
+	interceptor signal.Interceptor
 }
 
 // New creates a new Migrator with the given configuration.
-func New(cfg Config) (*Migrator, error) {
+func New(cfg Config, interceptor signal.Interceptor) (*Migrator, error) {
 	// Validate and set defaults for the config.
 	if err := validateConfig(&cfg); err != nil {
 		return nil, fmt.Errorf("invalid config: %w", err)
@@ -128,6 +132,7 @@ func New(cfg Config) (*Migrator, error) {
 		startTimeMigration: time.Now(),
 		lastLogTime:        time.Now(),
 		logInterval:        5 * time.Second,
+		interceptor:        interceptor,
 	}, nil
 }
 
@@ -636,6 +641,8 @@ func (m *Migrator) migrateBucket(ctx context.Context,
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
+		case <-m.interceptor.ShutdownChannel():
+			return fmt.Errorf("shutdown signal received")
 		default:
 		}
 
@@ -980,6 +987,8 @@ func (m *Migrator) walkAndCompare(ctx context.Context, srcBucket,
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
+		case <-m.interceptor.ShutdownChannel():
+			return fmt.Errorf("shutdown signal received")
 		default:
 		}
 

migratekvdb/migration_test.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/btcsuite/btclog/v2"
 	"github.com/lightningnetwork/lnd/kvdb"
+	"github.com/lightningnetwork/lnd/signal"
 	"github.com/stretchr/testify/require"
 	"go.etcd.io/bbolt"
 )
@@ -71,7 +72,10 @@ func TestMigration(t *testing.T) {
 		MetaDB:    metaDb,
 	}
 
-	migrator, err := New(cfg)
+	interceptor, err := signal.Intercept()
+	require.NoError(t, err)
+
+	migrator, err := New(cfg, interceptor)
 	require.NoError(t, err)
 
 	err = migrator.Migrate(context.Background(), sourceDB, targetDB)