feat(node): per-node life-cycle management

Signed-off-by: Chris Gianelloni <[email protected]>

Author: wolf31o2

internal/node/load.go

@@ -15,6 +15,7 @@
 package node
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"log/slog"
@@ -95,7 +96,7 @@ func Load(cfg *config.Config, logger *slog.Logger, immutableDir string) error {
 	if err != nil {
 		return fmt.Errorf("failed to load state: %w", err)
 	}
-	if err := ls.Start(); err != nil {
+	if err := ls.Start(context.Background()); err != nil {
 		return fmt.Errorf("failed to load state: %w", err)
 	}
 	// Open immutable DB

ledger/state.go

@@ -15,6 +15,7 @@
 package ledger
 
 import (
+	"context"
 	"encoding/hex"
 	"errors"
 	"fmt"
@@ -300,6 +301,7 @@ type LedgerState struct {
 	currentTip                       ochainsync.Tip
 	currentEpoch                     models.Epoch
 	dbWorkerPool                     *DatabaseWorkerPool
+	ctx                              context.Context
 	sync.RWMutex
 	chainsyncMutex             sync.Mutex
 	chainsyncBlockfetchMutex   sync.Mutex
@@ -334,7 +336,8 @@ func NewLedgerState(cfg LedgerStateConfig) (*LedgerState, error) {
 	return ls, nil
 }
 
-func (ls *LedgerState) Start() error {
+func (ls *LedgerState) Start(ctx context.Context) error {
+	ls.ctx = ctx
 	// Init metrics
 	ls.metrics.init(ls.config.PromRegistry)
 

node.go

@@ -46,6 +46,8 @@ type Node struct {
 	ouroboros      *ouroborosPkg.Ouroboros
 	shutdownFuncs  []func(context.Context) error
 	config         Config
+	ctx            context.Context
+	cancel         context.CancelFunc
 	done           chan struct{}
 	shutdownOnce   sync.Once
 }
@@ -73,6 +75,7 @@ func (n *Node) Run() error {
 			return err
 		}
 	}
+	n.ctx, n.cancel = context.WithCancel(context.Background())
 	// Load database
 	dbNeedsRecovery := false
 	dbConfig := &database.Config{
@@ -149,9 +152,12 @@ func (n *Node) Run() error {
 		}
 	}
 	// Start ledger
-	if err := n.ledgerState.Start(); err != nil {
+	if err := n.ledgerState.Start(n.ctx); err != nil {
 		return fmt.Errorf("failed to start ledger: %w", err)
 	}
+	n.shutdownFuncs = append(n.shutdownFuncs, func(ctx context.Context) error {
+		return n.ledgerState.Close()
+	})
 	// Initialize mempool
 	n.mempool = mempool.NewMempool(mempool.MempoolConfig{
 		MempoolCapacity: n.config.mempoolCapacity,
@@ -192,6 +198,9 @@ func (n *Node) Run() error {
 	if err := n.connManager.Start(); err != nil {
 		return err
 	}
+	n.shutdownFuncs = append(n.shutdownFuncs, func(ctx context.Context) error {
+		return n.connManager.Stop(ctx)
+	})
 	// Configure peer governor
 	n.peerGov = peergov.NewPeerGovernor(
 		peergov.PeerGovernorConfig{
@@ -214,6 +223,10 @@ func (n *Node) Run() error {
 	if err := n.peerGov.Start(); err != nil {
 		return err
 	}
+	n.shutdownFuncs = append(n.shutdownFuncs, func(ctx context.Context) error {
+		n.peerGov.Stop()
+		return nil
+	})
 	// Configure UTxO RPC
 	n.utxorpc = utxorpc.NewUtxorpc(
 		utxorpc.UtxorpcConfig{
@@ -227,9 +240,16 @@ func (n *Node) Run() error {
 	if err := n.utxorpc.Start(); err != nil {
 		return err
 	}
+	n.shutdownFuncs = append(n.shutdownFuncs, func(ctx context.Context) error {
+		return n.utxorpc.Stop(ctx)
+	})
 
 	// Wait for shutdown signal
-	<-n.done
+	if n.ctx != nil {
+		<-n.ctx.Done()
+	} else {
+		<-n.done
+	}
 	return nil
 }
 
@@ -249,6 +269,9 @@ func (n *Node) shutdown() error {
 	}
 	ctx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
 	defer cancel()
+	if n.cancel != nil {
+		n.cancel()
+	}
 
 	var err error