rework context handling throughout router and modules

internal/module/udp-multicast.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"log/slog"
 	"net"
+	"sync"
 	"time"
 
 	"github.com/google/jsonschema-go/jsonschema"
@@ -21,6 +22,7 @@ type UDPMulticast struct {
 	Addr   *net.UDPAddr
 	logger *slog.Logger
 	cancel context.CancelFunc
+	connMu sync.Mutex
 }
 
 func init() {
@@ -86,19 +88,21 @@ func (um *UDPMulticast) Start(ctx context.Context, router common.RouteIO) error
 	}
 	defer client.Close()
 
+	um.connMu.Lock()
 	um.conn = client
+	um.connMu.Unlock()
 
 	buffer := make([]byte, 2048)
 	for {
 		select {
 		case <-um.ctx.Done():
-			// TODO(jwetzell): cleanup?
-			um.logger.Debug("done")
 			return nil
 		default:
+			um.connMu.Lock()
 			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() {
@@ -136,5 +140,14 @@ func (um *UDPMulticast) Output(ctx context.Context, payload any) error {
 }
 
 func (um *UDPMulticast) Stop() {
-	um.cancel()
+	if um.cancel != nil {
+		um.cancel()
+	}
+	um.connMu.Lock()
+	defer um.connMu.Unlock()
+	if um.conn != nil {
+		um.conn.Close()
+		um.conn = nil
+	}
+	um.logger.Debug("done")
 }