rework module shutdown process to cancel module context in defer

internal/module/db-sqlite.go

@@ -76,20 +76,19 @@ func (dbs *DbSqlite) Start(ctx context.Context, inputHandler common.InputHandler
 	dbs.db = db
 	dbs.dbMu.Unlock()
 	<-dbs.ctx.Done()
+	dbs.logger.Debug("done")
 	return nil
 }
 
 func (dbs *DbSqlite) Stop() {
 	if dbs.cancel != nil {
-		dbs.cancel()
+		defer dbs.cancel()
 	}
 	dbs.dbMu.Lock()
 	defer dbs.dbMu.Unlock()
 	if dbs.db != nil {
 		dbs.db.Close()
-		dbs.db = nil
 	}
-	dbs.logger.Debug("done")
 }
 
 func (dbs *DbSqlite) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) {

internal/module/http-server.go

@@ -167,12 +167,13 @@ func (hs *HTTPServer) Start(ctx context.Context, inputHandler common.InputHandle
 	err := httpServer.ListenAndServe()
 	// TODO(jwetzell): handle server closed error differently
 	if err != nil {
-		if err.Error() != "http: Server closed" {
+		if !errors.Is(err, http.ErrServerClosed) {
 			return err
 		}
 	}
 
 	<-hs.ctx.Done()
+	hs.logger.Debug("done")
 	return nil
 }
 
@@ -200,7 +201,7 @@ func (hs *HTTPServer) Output(ctx context.Context, payload any) error {
 
 func (hs *HTTPServer) Stop() {
 	if hs.cancel != nil {
-		hs.cancel()
+		defer hs.cancel()
 	}
 	hs.serverMu.Lock()
 	defer hs.serverMu.Unlock()
@@ -209,7 +210,5 @@ func (hs *HTTPServer) Stop() {
 		hs.server.Shutdown(shutdownCtx)
 		shutdownCancel()
 		<-shutdownCtx.Done()
-		hs.server = nil
 	}
-	hs.logger.Debug("done")
 }

internal/module/midi-input.go

@@ -83,17 +83,16 @@ func (mi *MIDIInput) Start(ctx context.Context, inputHandler common.InputHandler
 	mi.stop = stop
 
 	<-mi.ctx.Done()
+	mi.logger.Debug("done")
 	return nil
 }
 
 func (mi *MIDIInput) Stop() {
 	if mi.cancel != nil {
-		mi.cancel()
+		defer mi.cancel()
 	}
 	if mi.stop != nil {
 		mi.stop()
-		mi.stop = nil
 	}
 	midi.CloseDriver()
-	mi.logger.Debug("done")
 }

internal/module/midi-output.go

@@ -86,6 +86,7 @@ func (mo *MIDIOutput) Start(ctx context.Context, inputHandler common.InputHandle
 	mo.sendFuncMu.Unlock()
 
 	<-mo.ctx.Done()
+	mo.logger.Debug("done")
 	return nil
 }
 
@@ -107,8 +108,7 @@ func (mo *MIDIOutput) Output(ctx context.Context, payload any) error {
 
 func (mo *MIDIOutput) Stop() {
 	if mo.cancel != nil {
-		mo.cancel()
+		defer mo.cancel()
 	}
 	midi.CloseDriver()
-	mo.logger.Debug("done")
 }

internal/module/mqtt-client.go

@@ -152,6 +152,7 @@ func (mc *MQTTClient) Start(ctx context.Context, inputHandler common.InputHandle
 	mc.clientMu.Unlock()
 
 	<-mc.ctx.Done()
+	mc.logger.Debug("done")
 	return nil
 }
 
@@ -189,7 +190,5 @@ func (mc *MQTTClient) Stop() {
 	defer mc.clientMu.Unlock()
 	if mc.client != nil {
 		mc.client.Disconnect(250)
-		mc.client = nil
 	}
-	mc.logger.Debug("done")
 }

internal/module/nats-client.go

@@ -102,6 +102,7 @@ func (nc *NATSClient) Start(ctx context.Context, inputHandler common.InputHandle
 	nc.subMu.Unlock()
 
 	<-nc.ctx.Done()
+	nc.logger.Debug("done")
 	return nil
 }
 
@@ -135,13 +136,12 @@ func (nc *NATSClient) Publish(ctx context.Context, topic string, payload any) er
 
 func (nc *NATSClient) Stop() {
 	if nc.cancel != nil {
-		nc.cancel()
+		defer nc.cancel()
 	}
 	nc.subMu.Lock()
 	defer nc.subMu.Unlock()
 	if nc.sub != nil {
 		nc.sub.Unsubscribe()
-		nc.sub = nil
 	}
 
 	nc.clientMu.Lock()
@@ -150,7 +150,5 @@ func (nc *NATSClient) Stop() {
 		nc.client.Drain()
 		// TODO(jwetzell): setup closed callback to get when client is fully closed
 		nc.client.Close()
-		nc.client = nil
 	}
-	nc.logger.Debug("done")
 }

internal/module/nats-server.go

@@ -96,9 +96,10 @@ func (ns *NATSServer) Start(ctx context.Context, inputHandler common.InputHandle
 	ns.cancel = cancel
 
 	natsServer, err := server.NewServer(&server.Options{
-		Host:  ns.Ip,
-		Port:  ns.Port,
-		NoLog: true,
+		Host:   ns.Ip,
+		Port:   ns.Port,
+		NoLog:  true,
+		NoSigs: true,
 	})
 
 	if err != nil {
@@ -107,26 +108,26 @@ func (ns *NATSServer) Start(ctx context.Context, inputHandler common.InputHandle
 
 	ns.serverMu.Lock()
 	ns.server = natsServer
-	defer ns.serverMu.Unlock()
 	natsServer.Start()
 
 	if !natsServer.ReadyForConnections(5 * time.Second) {
 		return errors.New("nats.server failed to start")
 	}
+	ns.serverMu.Unlock()
 
 	<-ns.ctx.Done()
-
+	ns.logger.Debug("done")
 	return nil
 }
 
 func (ns *NATSServer) Stop() {
 	if ns.cancel != nil {
-		ns.cancel()
+		defer ns.cancel()
 	}
 	ns.serverMu.Lock()
 	defer ns.serverMu.Unlock()
 	if ns.server != nil {
 		ns.server.Shutdown()
+		ns.server.WaitForShutdown()
 	}
-	ns.logger.Debug("done")
 }

internal/module/psn-client.go

@@ -63,7 +63,7 @@ func (pc *PSNClient) Start(ctx context.Context, inputHandler common.InputHandler
 	pc.connMu.Unlock()
 
 	buffer := make([]byte, 2048)
-	for {
+	for pc.ctx.Err() == nil {
 		select {
 		case <-pc.ctx.Done():
 			return nil
@@ -99,17 +99,18 @@ func (pc *PSNClient) Start(ctx context.Context, inputHandler common.InputHandler
 			}
 		}
 	}
+	<-pc.ctx.Done()
+	pc.logger.Debug("done")
+	return nil
 }
 
 func (pc *PSNClient) Stop() {
 	if pc.cancel != nil {
-		pc.cancel()
+		defer pc.cancel()
 	}
 	pc.connMu.Lock()
 	defer pc.connMu.Unlock()
 	if pc.conn != nil {
 		pc.conn.Close()
-		pc.conn = nil
 	}
-	pc.logger.Debug("done")
 }

internal/module/redis-client.go

@@ -94,20 +94,19 @@ func (rc *RedisClient) Start(ctx context.Context, inputHandler common.InputHandl
 	rc.clientMu.Unlock()
 
 	<-rc.ctx.Done()
+	rc.logger.Debug("done")
 	return nil
 }
 
 func (rc *RedisClient) Stop() {
 	if rc.cancel != nil {
-		rc.cancel()
+		defer rc.cancel()
 	}
 	rc.clientMu.Lock()
 	defer rc.clientMu.Unlock()
 	if rc.client != nil {
 		rc.client.Close()
-		rc.client = nil
 	}
-	rc.logger.Debug("done")
 }
 
 func (rc *RedisClient) Get(ctx context.Context, key string) (any, error) {

internal/module/serial-client.go

@@ -159,6 +159,8 @@ func (sc *SerialClient) Start(ctx context.Context, inputHandler common.InputHand
 			}
 		}
 	}
+	<-sc.ctx.Done()
+	sc.logger.Debug("done")
 	return nil
 }
 
@@ -176,13 +178,11 @@ func (sc *SerialClient) Output(ctx context.Context, payload any) error {
 
 func (sc *SerialClient) Stop() {
 	if sc.cancel != nil {
-		sc.cancel()
+		defer sc.cancel()
 	}
 	sc.portMu.Lock()
 	defer sc.portMu.Unlock()
 	if sc.port != nil {
 		sc.port.Close()
-		sc.port = nil
 	}
-	sc.logger.Debug("done")
 }

internal/module/sip-call-server.go

@@ -168,6 +168,7 @@ func (scs *SIPCallServer) Start(ctx context.Context, inputHandler common.InputHa
 	}
 
 	<-scs.ctx.Done()
+	scs.logger.Debug("done")
 	return nil
 }
 
@@ -251,12 +252,11 @@ func (scs *SIPCallServer) Output(ctx context.Context, payload any) error {
 
 func (scs *SIPCallServer) Stop() {
 	if scs.cancel != nil {
-		scs.cancel()
+		defer scs.cancel()
 	}
 	scs.uaMu.Lock()
 	defer scs.uaMu.Unlock()
 	if scs.ua != nil {
 		scs.ua.Close()
 	}
-	scs.logger.Debug("done")
 }

internal/module/sip-dtmf-server.go

@@ -190,6 +190,7 @@ func (sds *SIPDTMFServer) Start(ctx context.Context, inputHandler common.InputHa
 	}
 
 	<-sds.ctx.Done()
+	sds.logger.Debug("done")
 	return nil
 }
 
@@ -286,12 +287,11 @@ func (sds *SIPDTMFServer) Output(ctx context.Context, payload any) error {
 
 func (sds *SIPDTMFServer) Stop() {
 	if sds.cancel != nil {
-		sds.cancel()
+		defer sds.cancel()
 	}
 	sds.uaMu.Lock()
 	defer sds.uaMu.Unlock()
 	if sds.ua != nil {
 		sds.ua.Close()
 	}
-	sds.logger.Debug("done")
 }

internal/module/tcp-client.go

@@ -100,11 +100,12 @@ func (tc *TCPClient) Start(ctx context.Context, inputHandler common.InputHandler
 	tc.ctx = moduleContext
 	tc.cancel = cancel
 
+CONNECT_RETRY:
 	for tc.ctx.Err() == nil {
 		err := tc.SetupConn()
 		if err != nil {
 			if tc.ctx.Err() != nil {
-				return nil
+				break CONNECT_RETRY
 			}
 			tc.logger.Error("connection error", "error", err.Error())
 			time.Sleep(time.Second * 2)
@@ -112,39 +113,38 @@ func (tc *TCPClient) Start(ctx context.Context, inputHandler common.InputHandler
 		}
 
 		buffer := make([]byte, 1024)
-		select {
-		case <-tc.ctx.Done():
-			return nil
-		default:
-		READ:
-			for {
-				select {
-				case <-tc.ctx.Done():
-					return nil
-				default:
-					byteCount, err := tc.conn.Read(buffer)
 
-					if err != nil {
-						tc.framer.Clear()
-						break READ
+	READ:
+		for tc.ctx.Err() == nil {
+			tc.conn.SetReadDeadline(time.Now().Add(time.Millisecond * 200))
+			byteCount, err := tc.conn.Read(buffer)
+
+			if err != nil {
+				if opErr, ok := err.(*net.OpError); ok {
+					//NOTE(jwetzell) we hit deadline
+					if opErr.Timeout() {
+						continue
 					}
+				}
+				break READ
+			}
 
-					if tc.framer != nil {
-						if byteCount > 0 {
-							messages := tc.framer.Decode(buffer[0:byteCount])
-							for _, message := range messages {
-								if tc.inputHandler != nil {
-									tc.inputHandler(tc.ctx, tc.Id(), message)
-								} else {
-									tc.logger.Error("input received but no input handler is configured")
-								}
-							}
+			if tc.framer != nil {
+				if byteCount > 0 {
+					messages := tc.framer.Decode(buffer[0:byteCount])
+					for _, message := range messages {
+						if tc.inputHandler != nil {
+							tc.inputHandler(tc.ctx, tc.Id(), message)
+						} else {
+							tc.logger.Error("input received but no input handler is configured")
 						}
 					}
 				}
 			}
 		}
 	}
+	<-tc.ctx.Done()
+	tc.logger.Debug("done")
 	return nil
 }
 
@@ -172,14 +172,11 @@ func (tc *TCPClient) Output(ctx context.Context, payload any) error {
 
 func (tc *TCPClient) Stop() {
 	if tc.cancel != nil {
-		tc.cancel()
+		defer tc.cancel()
 	}
 	tc.connMu.Lock()
 	defer tc.connMu.Unlock()
 	if tc.conn != nil {
 		tc.conn.Close()
-		tc.conn = nil
 	}
-	tc.logger.Debug("done")
-
 }

internal/module/tcp-server.go

@@ -10,7 +10,6 @@ import (
 	"slices"
 	"strings"
 	"sync"
-	"syscall"
 	"time"
 
 	"github.com/google/jsonschema-go/jsonschema"
@@ -58,12 +57,14 @@ func init() {
 				return nil, fmt.Errorf("net.tcp.server framing error: %w", err)
 			}
 
-			framer := framer.GetFramer(framingMethodString)
+			inFramer := framer.GetFramer(framingMethodString)
 
-			if framer == nil {
+			if inFramer == nil {
 				return nil, fmt.Errorf("net.tcp.server unknown framing method: %s", framingMethodString)
 			}
 
+			outFramer := framer.GetFramer(framingMethodString)
+
 			ipString, err := params.GetString("ip")
 			if err != nil {
 				if errors.Is(err, config.ErrParamNotFound) {
@@ -77,24 +78,33 @@ func init() {
 			if err != nil {
 				return nil, err
 			}
-			return &TCPServer{Framer: framer, Addr: addr, config: moduleConfig, logger: CreateLogger(moduleConfig)}, nil
+			return &TCPServer{InFramer: inFramer, OutFramer: outFramer, framerType: framingMethodString, Addr: addr, config: moduleConfig, logger: CreateLogger(moduleConfig)}, nil
 		},
 	})
 }
 
+type tcpConnection struct {
+	conn   *net.TCPConn
+	framer framer.Framer
+}
+
 type TCPServer struct {
-	config        config.ModuleConfig
-	Addr          *net.TCPAddr
-	Framer        framer.Framer
-	ctx           context.Context
-	inputHandler  common.InputHandler
-	wg            sync.WaitGroup
-	connections   []*net.TCPConn
-	connectionsMu sync.RWMutex
-	logger        *slog.Logger
-	cancel        context.CancelFunc
-	listener      *net.TCPListener
-	listenerMu    sync.Mutex
+	config                config.ModuleConfig
+	Addr                  *net.TCPAddr
+	InFramer              framer.Framer
+	OutFramer             framer.Framer
+	framerType            string
+	ctx                   context.Context
+	inputHandler          common.InputHandler
+	wg                    sync.WaitGroup
+	connections           []tcpConnection
+	connectionsMu         sync.RWMutex
+	logger                *slog.Logger
+	cancel                context.CancelFunc
+	listener              *net.TCPListener
+	listenerMu            sync.Mutex
+	connectionShutdownCtx context.Context
+	connectionShutdown    context.CancelFunc
 }
 
 func (ts *TCPServer) Id() string {
@@ -107,70 +117,43 @@ func (ts *TCPServer) Type() string {
 
 func (ts *TCPServer) handleClient(client *net.TCPConn) {
 	ts.connectionsMu.Lock()
-	ts.connections = append(ts.connections, client)
+	ts.connections = append(ts.connections, tcpConnection{conn: client, framer: framer.GetFramer(ts.framerType)})
 	ts.connectionsMu.Unlock()
 	ts.logger.Debug("connection accepted", "remoteAddr", client.RemoteAddr().String())
-	defer client.Close()
+	defer func() {
+		client.Close()
+		ts.connectionsMu.Lock()
+		for i := 0; i < len(ts.connections); i++ {
+			if ts.connections[i].conn == client {
+				ts.connections = slices.Delete(ts.connections, i, i+1)
+				break
+			}
+		}
+		ts.connectionsMu.Unlock()
+		ts.logger.Debug("connection closed", "remoteAddr", client.RemoteAddr().String())
+	}()
 
 	buffer := make([]byte, 1024)
-ClientRead:
-	for ts.ctx.Err() == nil {
-		select {
-		case <-ts.ctx.Done():
-			client.Close()
-			ts.connectionsMu.Lock()
-			for i := 0; i < len(ts.connections); i++ {
-				if ts.connections[i] == client {
-					ts.connections = slices.Delete(ts.connections, i, i+1)
-					break
+	for ts.ctx.Err() == nil && ts.connectionShutdownCtx.Err() == nil {
+		client.SetDeadline(time.Now().Add(time.Millisecond * 200))
+		byteCount, err := client.Read(buffer)
+		if err != nil {
+			if opErr, ok := err.(*net.OpError); ok {
+				//NOTE(jwetzell) we hit deadline
+				if opErr.Timeout() {
+					continue
 				}
 			}
-			ts.connectionsMu.Unlock()
-			return
-		default:
-			client.SetDeadline(time.Now().Add(time.Millisecond * 200))
-			byteCount, err := client.Read(buffer)
-
-			if err != nil {
-				if opErr, ok := err.(*net.OpError); ok {
-					//NOTE(jwetzell) we hit deadline
-					if opErr.Timeout() {
-						continue ClientRead
-					}
-					if errors.Is(opErr, syscall.ECONNRESET) {
-						ts.connectionsMu.Lock()
-						for i := 0; i < len(ts.connections); i++ {
-							if ts.connections[i] == client {
-								ts.connections = slices.Delete(ts.connections, i, i+1)
-								break
-							}
-						}
-						ts.logger.Debug("connection reset", "remoteAddr", client.RemoteAddr().String())
-						ts.connectionsMu.Unlock()
-					}
-				}
-
-				if err.Error() == "EOF" {
-					ts.connectionsMu.Lock()
-					for i := 0; i < len(ts.connections); i++ {
-						if ts.connections[i] == client {
-							ts.connections = slices.Delete(ts.connections, i, i+1)
-							break
-						}
-					}
-					ts.connectionsMu.Unlock()
-				}
-				return
-			}
-			if ts.Framer != nil {
-				if byteCount > 0 {
-					messages := ts.Framer.Decode(buffer[0:byteCount])
-					for _, message := range messages {
-						if ts.inputHandler != nil {
-							ts.inputHandler(ts.ctx, ts.Id(), message)
-						} else {
-							ts.logger.Error("input received but no input handler is configured")
-						}
+			break
+		}
+		if ts.InFramer != nil {
+			if byteCount > 0 {
+				messages := ts.InFramer.Decode(buffer[0:byteCount])
+				for _, message := range messages {
+					if ts.inputHandler != nil {
+						ts.inputHandler(ts.ctx, ts.Id(), message)
+					} else {
+						ts.logger.Error("input received but no input handler is configured")
 					}
 				}
 			}
@@ -185,6 +168,10 @@ func (ts *TCPServer) Start(ctx context.Context, inputHandler common.InputHandler
 	ts.ctx = moduleContext
 	ts.cancel = cancel
 
+	shutdownCtx, shutdownCancel := context.WithCancel(context.Background())
+	ts.connectionShutdownCtx = shutdownCtx
+	ts.connectionShutdown = shutdownCancel
+
 	listener, err := net.ListenTCP("tcp", ts.Addr)
 	if err != nil {
 		return err
@@ -198,6 +185,9 @@ AcceptLoop:
 	for ts.ctx.Err() == nil {
 		conn, err := listener.AcceptTCP()
 		if err != nil {
+			if errors.Is(err, net.ErrClosed) {
+				break AcceptLoop
+			}
 			select {
 			case <-ts.ctx.Done():
 				break AcceptLoop
@@ -211,6 +201,8 @@ AcceptLoop:
 		}
 	}
 	ts.wg.Done()
+	<-ts.ctx.Done()
+	ts.logger.Debug("done")
 	return nil
 }
 
@@ -221,15 +213,21 @@ func (ts *TCPServer) Output(ctx context.Context, payload any) error {
 		return errors.New("net.tcp.server is only able to output bytes")
 	}
 	ts.connectionsMu.Lock()
+	defer ts.connectionsMu.Unlock()
 	var errorString strings.Builder
 
+	if ts.OutFramer == nil {
+		return errors.New("no output framer configured")
+	}
+
+	outputBytes := ts.OutFramer.Encode(payloadBytes)
+
 	for _, connection := range ts.connections {
-		_, err := connection.Write(payloadBytes)
+		_, err := connection.conn.Write(outputBytes)
 		if err != nil {
 			fmt.Fprintf(&errorString, "%s\n", err.Error())
 		}
 	}
-	ts.connectionsMu.Unlock()
 
 	if errorString.String() == "" {
 		return nil
@@ -239,14 +237,18 @@ func (ts *TCPServer) Output(ctx context.Context, payload any) error {
 
 func (ts *TCPServer) Stop() {
 	if ts.cancel != nil {
-		ts.cancel()
+		defer ts.cancel()
 	}
+	if ts.connectionShutdown != nil {
+		ts.connectionShutdown()
+	}
+
 	ts.listenerMu.Lock()
 	defer ts.listenerMu.Unlock()
 	if ts.listener != nil {
 		ts.listener.Close()
-		ts.listener = nil
 	}
+	ts.logger.Debug("waiting for connections to close")
 	ts.wg.Wait()
-	ts.logger.Debug("done")
+	ts.logger.Debug("all connections closed")
 }

internal/module/time-interval.go

@@ -67,25 +67,26 @@ func (i *TimeInterval) Start(ctx context.Context, inputHandler common.InputHandl
 	ticker := time.NewTicker(time.Millisecond * time.Duration(i.Duration))
 	i.ticker = ticker
 
-	for {
+	for i.ctx.Err() == nil {
 		select {
-		case <-i.ctx.Done():
-			return nil
 		case <-ticker.C:
 			if i.inputHandler != nil {
 				i.inputHandler(i.ctx, i.Id(), time.Now())
 			}
+		default:
+			continue
 		}
 	}
+	<-i.ctx.Done()
+	i.logger.Debug("done")
+	return nil
 }
 
 func (i *TimeInterval) Stop() {
 	if i.cancel != nil {
-		i.cancel()
+		defer i.cancel()
 	}
 	if i.ticker != nil {
 		i.ticker.Stop()
-		i.ticker = nil
 	}
-	i.logger.Debug("done")
 }

internal/module/time-timer.go

@@ -66,7 +66,7 @@ func (t *TimeTimer) Start(ctx context.Context, inputHandler common.InputHandler)
 	t.cancel = cancel
 
 	t.timer = time.NewTimer(time.Millisecond * time.Duration(t.Duration))
-	for {
+	for t.ctx.Err() == nil {
 		select {
 		case <-t.ctx.Done():
 			return nil
@@ -76,15 +76,16 @@ func (t *TimeTimer) Start(ctx context.Context, inputHandler common.InputHandler)
 			}
 		}
 	}
+	<-t.ctx.Done()
+	t.logger.Debug("done")
+	return nil
 }
 
 func (t *TimeTimer) Stop() {
 	if t.cancel != nil {
-		t.cancel()
+		defer t.cancel()
 	}
 	if t.timer != nil {
 		t.timer.Stop()
-		t.timer = nil
 	}
-	t.logger.Debug("done")
 }

internal/module/udp-client.go

@@ -96,6 +96,7 @@ func (uc *UDPClient) Start(ctx context.Context, inputHandler common.InputHandler
 	}
 
 	<-uc.ctx.Done()
+	uc.logger.Debug("done")
 	return nil
 }
 
@@ -120,14 +121,11 @@ func (uc *UDPClient) Output(ctx context.Context, payload any) error {
 
 func (uc *UDPClient) Stop() {
 	if uc.cancel != nil {
-		uc.cancel()
+		defer uc.cancel()
 	}
 	uc.connMu.Lock()
 	defer uc.connMu.Unlock()
 	if uc.conn != nil {
 		uc.conn.Close()
-		uc.conn = nil
 	}
-
-	uc.logger.Debug("done")
 }

internal/module/udp-multicast.go

@@ -93,35 +93,34 @@ func (um *UDPMulticast) Start(ctx context.Context, inputHandler common.InputHand
 	um.connMu.Unlock()
 
 	buffer := make([]byte, 2048)
-	for {
-		select {
-		case <-um.ctx.Done():
-			return nil
-		default:
-			um.connMu.Lock()
-			um.conn.SetDeadline(time.Now().Add(time.Millisecond * 200))
+	for um.ctx.Err() == nil {
+		um.conn.SetDeadline(time.Now().Add(time.Millisecond * 200))
 
-			numBytes, _, err := um.conn.ReadFromUDP(buffer)
-			um.connMu.Unlock()
-			if err != nil {
-				//NOTE(jwetzell) we hit deadline
-				if opErr, ok := err.(*net.OpError); ok && opErr.Timeout() {
-					continue
-				}
-				return err
+		numBytes, _, err := um.conn.ReadFromUDP(buffer)
+		if err != nil {
+			if errors.Is(err, net.ErrClosed) {
+				break
 			}
+			//NOTE(jwetzell) we hit deadline
+			if opErr, ok := err.(*net.OpError); ok && opErr.Timeout() {
+				continue
+			}
+			return err
+		}
 
-			if numBytes > 0 {
-				message := buffer[:numBytes]
+		if numBytes > 0 {
+			message := buffer[:numBytes]
 
-				if um.inputHandler != nil {
-					um.inputHandler(um.ctx, um.Id(), message)
-				} else {
-					um.logger.Error("input received but no input handler is configured")
-				}
+			if um.inputHandler != nil {
+				um.inputHandler(um.ctx, um.Id(), message)
+			} else {
+				um.logger.Error("input received but no input handler is configured")
 			}
 		}
 	}
+	<-um.ctx.Done()
+	um.logger.Debug("done")
+	return nil
 }
 
 func (um *UDPMulticast) Output(ctx context.Context, payload any) error {
@@ -141,13 +140,11 @@ func (um *UDPMulticast) Output(ctx context.Context, payload any) error {
 
 func (um *UDPMulticast) Stop() {
 	if um.cancel != nil {
-		um.cancel()
+		defer um.cancel()
 	}
 	um.connMu.Lock()
 	defer um.connMu.Unlock()
 	if um.conn != nil {
 		um.conn.Close()
-		um.conn = nil
 	}
-	um.logger.Debug("done")
 }

internal/module/udp-server.go

@@ -111,32 +111,29 @@ func (us *UDPServer) Start(ctx context.Context, inputHandler common.InputHandler
 	}
 	us.listenerMu.Lock()
 	us.listener = listener
+	us.listenerMu.Unlock()
 
 	buffer := make([]byte, us.BufferSize)
 	for us.ctx.Err() == nil {
-		select {
-		case <-us.ctx.Done():
-			return nil
-		default:
-			listener.SetDeadline(time.Now().Add(time.Millisecond * 200))
+		listener.SetDeadline(time.Now().Add(time.Millisecond * 200))
 
-			numBytes, _, err := listener.ReadFromUDP(buffer)
-			if err != nil {
-				//NOTE(jwetzell) we hit deadline
-				if opErr, ok := err.(*net.OpError); ok && opErr.Timeout() {
-					continue
-				}
-				return err
-			}
-			message := buffer[:numBytes]
-			if us.inputHandler != nil {
-				us.inputHandler(us.ctx, us.Id(), message)
-			} else {
-				us.logger.Error("input received but no input handler is configured")
+		numBytes, _, err := listener.ReadFromUDP(buffer)
+		if err != nil {
+			//NOTE(jwetzell) we hit deadline
+			if opErr, ok := err.(*net.OpError); ok && opErr.Timeout() {
+				continue
 			}
+			break
+		}
+		message := buffer[:numBytes]
+		if us.inputHandler != nil {
+			us.inputHandler(us.ctx, us.Id(), message)
+		} else {
+			us.logger.Error("input received but no input handler is configured")
 		}
 	}
-	us.listenerMu.Unlock()
+	<-us.ctx.Done()
+	us.logger.Debug("done")
 	return nil
 }
 
@@ -146,13 +143,11 @@ func (us *UDPServer) Output(ctx context.Context, payload any) error {
 
 func (us *UDPServer) Stop() {
 	if us.cancel != nil {
-		us.cancel()
+		defer us.cancel()
 	}
 	us.listenerMu.Lock()
 	defer us.listenerMu.Unlock()
 	if us.listener != nil {
 		us.listener.Close()
-		us.listener = nil
 	}
-	us.logger.Debug("done")
 }

internal/module/websocket-client.go

@@ -99,6 +99,7 @@ func (wc *WebSocketClient) Start(ctx context.Context, inputHandler common.InputH
 		time.Sleep(2 * time.Second)
 	}
 	<-wc.ctx.Done()
+	wc.logger.Debug("done")
 	return nil
 }
 
@@ -116,10 +117,6 @@ func (wc *WebSocketClient) readLoop() {
 				if opErr.Timeout() {
 					continue
 				}
-				// NOTE(jwetzell) connection was closed
-				if errors.Is(opErr, net.ErrClosed) {
-					continue
-				}
 			}
 			wc.logger.Error("websocket read error", "error", err)
 			return
@@ -182,13 +179,11 @@ func (wc *WebSocketClient) Output(ctx context.Context, payload any) error {
 
 func (wc *WebSocketClient) Stop() {
 	if wc.cancel != nil {
-		wc.cancel()
+		defer wc.cancel()
 	}
 	wc.connMu.Lock()
 	defer wc.connMu.Unlock()
 	if wc.conn != nil {
 		wc.conn.Close()
-		wc.conn = nil
 	}
-	wc.logger.Debug("done")
 }

router_test.go

@@ -63,7 +63,7 @@ func (mcm *MockCounterModule) Type() string {
 
 func (mcm *MockCounterModule) Stop() {
 	if mcm.cancel != nil {
-		mcm.cancel()
+		defer mcm.cancel()
 	}
 }