feat: clean shutdown

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

Author: wolf31o2

connmanager/connection_manager.go

@@ -15,8 +15,11 @@
 package connmanager
 
 import (
+	"context"
+	"errors"
 	"io"
 	"log/slog"
+	"net"
 	"sync"
 
 	"github.com/blinklabs-io/dingo/event"
@@ -36,8 +39,8 @@ const (
 
 type connectionInfo struct {
 	conn      *ouroboros.Connection
-	isInbound bool
 	peerAddr  string
+	isInbound bool
 }
 
 type peerConnectionState struct {
@@ -47,19 +50,22 @@ type peerConnectionState struct {
 
 type ConnectionManager struct {
 	connections      map[ouroboros.ConnectionId]*connectionInfo
+	metrics          *connectionManagerMetrics
+	listeners        []net.Listener
 	config           ConnectionManagerConfig
 	connectionsMutex sync.Mutex
-	metrics          *connectionManagerMetrics
+	listenersMutex   sync.Mutex
+	closing          bool
 }
 
 type ConnectionManagerConfig struct {
+	PromRegistry       prometheus.Registerer
 	Logger             *slog.Logger
 	EventBus           *event.EventBus
 	ConnClosedFunc     ConnectionManagerConnClosedFunc
 	Listeners          []ListenerConfig
 	OutboundConnOpts   []ouroboros.ConnectionOptionFunc
 	OutboundSourcePort uint
-	PromRegistry       prometheus.Registerer
 }
 
 type connectionManagerMetrics struct {
@@ -211,6 +217,57 @@ func (c *ConnectionManager) Start() error {
 	return nil
 }
 
+func (c *ConnectionManager) Stop(ctx context.Context) error {
+	var err error
+
+	c.config.Logger.Info("stopping connection manager")
+
+	// Mark closing to suppress accept-loop noise
+	c.listenersMutex.Lock()
+	c.closing = true
+	c.listenersMutex.Unlock()
+
+	// Stop accepting new connections
+	c.stopListeners()
+
+	// Close all existing connections gracefully
+	c.connectionsMutex.Lock()
+	conns := make([]*ouroboros.Connection, 0, len(c.connections))
+	for _, info := range c.connections {
+		conns = append(conns, info.conn)
+	}
+	c.connectionsMutex.Unlock()
+
+	// Close connections
+	for _, conn := range conns {
+		if conn != nil {
+			if closeErr := conn.Close(); closeErr != nil {
+				err = errors.Join(err, closeErr)
+			}
+		}
+	}
+
+	c.config.Logger.Info("connection manager stopped")
+	return err
+}
+
+func (c *ConnectionManager) stopListeners() {
+	c.listenersMutex.Lock()
+	defer c.listenersMutex.Unlock()
+
+	for _, listener := range c.listeners {
+		if listener != nil {
+			if err := listener.Close(); err != nil {
+				c.config.Logger.Warn(
+					"error closing listener",
+					"error", err,
+				)
+			}
+		}
+	}
+	c.listeners = nil
+}
+
 func (c *ConnectionManager) AddConnection(
 	conn *ouroboros.Connection,
 	isInbound bool,

connmanager/listener.go

@@ -16,9 +16,11 @@ package connmanager
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"net"
 	"runtime"
+	"strings"
 
 	"github.com/blinklabs-io/dingo/event"
 	ouroboros "github.com/blinklabs-io/gouroboros"
@@ -81,6 +83,11 @@ func (c *ConnectionManager) startListener(l ListenerConfig) error {
 			)
 		}
 	}
+	// Track listener for shutdown
+	c.listenersMutex.Lock()
+	c.listeners = append(c.listeners, l.Listener)
+	c.listenersMutex.Unlock()
+
 	// Build connection options
 	defaultConnOpts := []ouroboros.ConnectionOptionFunc{
 		ouroboros.WithLogger(c.config.Logger),
@@ -96,6 +103,29 @@ func (c *ConnectionManager) startListener(l ListenerConfig) error {
 			// Accept connection
 			conn, err := l.Listener.Accept()
 			if err != nil {
+				// If we're closing, exit quietly
+				c.listenersMutex.Lock()
+				isClosing := c.closing
+				c.listenersMutex.Unlock()
+				if isClosing {
+					c.config.Logger.Debug("listener: shutting down, stopping accept loop")
+					return
+				}
+				// During shutdown, closing the listener will cause Accept to return
+				// a net.ErrClosed. Treat this as a normal termination and exit the loop
+				if errors.Is(err, net.ErrClosed) || strings.Contains(err.Error(), "use of closed network connection") {
+					c.config.Logger.Debug("listener: closed, stopping accept loop")
+					return
+				}
+				// Some platforms may return temporary errors; handle and continue
+				var ne net.Error
+				if errors.As(err, &ne) && ne.Temporary() {
+					c.config.Logger.Warn(
+						fmt.Sprintf("listener: temporary accept error: %s", err),
+					)
+					continue
+				}
+				// Otherwise, log at error level and continue
 				c.config.Logger.Error(
 					fmt.Sprintf("listener: accept failed: %s", err),
 				)

internal/integration/benchmark_test.go

@@ -99,12 +99,12 @@ func loadBlockData(numBlocks int) ([][]byte, error) {
 
 // getTestBackends returns a slice of test backends for benchmarking
 func getTestBackends(b *testing.B, diskDataDir string) []struct {
-	name   string
 	config *database.Config
+	name   string
 } {
 	backends := []struct {
-		name   string
 		config *database.Config
+		name   string
 	}{
 		{
 			name: "memory",
@@ -133,8 +133,8 @@ func getTestBackends(b *testing.B, diskDataDir string) []struct {
 		// Use path prefix for isolation instead of unique bucket names
 		testPrefix := strings.ReplaceAll(b.Name(), "/", "-")
 		backends = append(backends, struct {
-			name   string
 			config *database.Config
+			name   string
 		}{
 			name: "GCS",
 			config: &database.Config{
@@ -153,8 +153,8 @@ func getTestBackends(b *testing.B, diskDataDir string) []struct {
 		// Use path prefix for isolation instead of unique bucket names
 		testPrefix := strings.ReplaceAll(b.Name(), "/", "-")
 		backends = append(backends, struct {
-			name   string
 			config *database.Config
+			name   string
 		}{
 			name: "S3",
 			config: &database.Config{

internal/node/node.go

@@ -142,16 +142,17 @@ func Run(cfg *config.Config, logger *slog.Logger) error {
 		"component",
 		"node",
 	)
+	metricsServer := &http.Server{
+		Addr: fmt.Sprintf(
+			"%s:%d",
+			cfg.BindAddr,
+			cfg.MetricsPort,
+		),
+		ReadHeaderTimeout: 60 * time.Second,
+	}
 	go func() {
-		debugger := &http.Server{
-			Addr: fmt.Sprintf(
-				"%s:%d",
-				cfg.BindAddr,
-				cfg.MetricsPort,
-			),
-			ReadHeaderTimeout: 60 * time.Second,
-		}
-		if err := debugger.ListenAndServe(); err != nil {
+		if err := metricsServer.ListenAndServe(); err != nil &&
+			err != http.ErrServerClosed {
 			logger.Error(
 				fmt.Sprintf("failed to start metrics listener: %s", err),
 				"component", "node",
@@ -166,21 +167,52 @@ func Run(cfg *config.Config, logger *slog.Logger) error {
 		syscall.SIGTERM,
 	)
 	defer signalCtxStop()
+
+	// Run node in goroutine
+	errChan := make(chan error, 1)
 	go func() {
-		<-signalCtx.Done()
-		logger.Info("signal received, shutting down")
-		if err := d.Stop(); err != nil { //nolint:contextcheck
-			logger.Error(
-				"failure(s) while shutting down",
-				"error",
-				err,
-			)
+		if err := d.Run(); err != nil {
+			errChan <- err
 		}
-		os.Exit(0)
 	}()
-	// Run node
-	if err := d.Run(); err != nil {
+
+	// Wait for signal or error
+	select {
+	case <-signalCtx.Done():
+		logger.Info("signal received, initiating graceful shutdown")
+
+		// Shutdown metrics server
+		shutdownCtx, cancel := context.WithTimeout(
+			context.Background(),
+			5*time.Second,
+		)
+		if err := metricsServer.Shutdown(shutdownCtx); err != nil {
+			logger.Error("metrics server shutdown error", "error", err)
+		}
+		cancel()
+
+		// Shutdown node
+		if err := d.Stop(); err != nil {
+			logger.Error("shutdown errors occurred", "error", err)
+			return err
+		}
+		logger.Info("shutdown complete")
+		return nil
+
+	case err := <-errChan:
+		logger.Error("node error", "error", err)
+		signalCtxStop()
+
+		// Cleanup on error
+		shutdownCtx, cancel := context.WithTimeout(
+			context.Background(),
+			5*time.Second,
+		)
+		if shutdownErr := metricsServer.Shutdown(shutdownCtx); shutdownErr != nil {
+			logger.Error("metrics server shutdown error", "error", shutdownErr)
+		}
+		cancel()
+
 		return err
 	}
-	return nil
 }

mempool/mempool.go

@@ -15,6 +15,7 @@
 package mempool
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"log/slog"
@@ -145,6 +146,28 @@ func (m *Mempool) RemoveConsumer(connId ouroboros.ConnectionId) {
 	m.consumersMutex.Unlock()
 }
 
+func (m *Mempool) Stop(ctx context.Context) error {
+	m.logger.Debug("stopping mempool")
+
+	// Stop all consumers
+	m.consumersMutex.Lock()
+	for _, consumer := range m.consumers {
+		if consumer != nil {
+			consumer.ClearCache()
+		}
+	}
+	m.consumers = make(map[ouroboros.ConnectionId]*MempoolConsumer)
+	m.consumersMutex.Unlock()
+
+	// Clear transactions
+	m.Lock()
+	m.transactions = nil
+	m.Unlock()
+
+	m.logger.Debug("mempool stopped")
+	return nil
+}
+
 func (m *Mempool) Consumer(connId ouroboros.ConnectionId) *MempoolConsumer {
 	m.consumersMutex.Lock()
 	defer m.consumersMutex.Unlock()