v0.1.0 (#17)

envd/internal/port/forward.go

@@ -1,4 +1,5 @@
 // SPDX-License-Identifier: Apache-2.0
+// Modifications by M/S Omukk
 
 // portf (port forward) periodaically scans opened TCP ports on the 127.0.0.1 (or localhost)
 // and launches `socat` process for every such port in the background.
@@ -80,8 +81,16 @@ func (f *Forwarder) StartForwarding(ctx context.Context) {
 	}
 
 	for {
-		// procs is an array of currently opened ports.
-		if procs, ok := <-f.scannerSubscriber.Messages; ok {
+		select {
+		case <-ctx.Done():
+			f.stopAllForwarding()
+			return
+		case procs, ok := <-f.scannerSubscriber.Messages:
+			if !ok {
+				f.stopAllForwarding()
+				return
+			}
+
 			// Now we are going to refresh all ports that are being forwarded in the `ports` map. Maybe add new ones
 			// and maybe remove some.
 
@@ -133,11 +142,22 @@ func (f *Forwarder) StartForwarding(ctx context.Context) {
 	}
 }
 
-func (f *Forwarder) startPortForwarding(ctx context.Context, p *PortToForward) {
+func (f *Forwarder) stopAllForwarding() {
+	for _, p := range f.ports {
+		f.stopPortForwarding(p)
+	}
+	f.ports = make(map[string]*PortToForward)
+}
+
+func (f *Forwarder) startPortForwarding(_ context.Context, p *PortToForward) {
 	// https://unix.stackexchange.com/questions/311492/redirect-application-listening-on-localhost-to-listening-on-external-interface
 	// socat -d -d TCP4-LISTEN:4000,bind=169.254.0.21,fork TCP4:localhost:4000
 	// reuseaddr is used to fix the "Address already in use" error when restarting socat quickly.
-	cmd := exec.CommandContext(ctx,
+	//
+	// We use exec.Command (not CommandContext) because stopAllForwarding kills
+	// socat via SIGKILL to the process group. CommandContext would also call
+	// cmd.Wait() on context cancellation, racing with the wait goroutine below.
+	cmd := exec.Command(
 		"socat", "-d", "-d", "-d",
 		fmt.Sprintf("TCP4-LISTEN:%v,bind=%s,reuseaddr,fork", p.port, f.sourceIP.To4()),
 		fmt.Sprintf("TCP%d:localhost:%v", p.family, p.port),

envd/internal/port/scan.go

@@ -1,8 +1,10 @@
 // SPDX-License-Identifier: Apache-2.0
+// Modifications by M/S Omukk
 
 package port
 
 import (
+	"context"
 	"sync"
 	"time"
 
@@ -10,8 +12,7 @@ import (
 )
 
 type Scanner struct {
-	scanExit chan struct{}
-	period   time.Duration
+	period time.Duration
 
 	// Plain mutex-protected map instead of concurrent-map. The concurrent-map
 	// library's Items() spawns goroutines and uses a WaitGroup internally,
@@ -20,15 +21,10 @@ type Scanner struct {
 	subs map[string]*ScannerSubscriber
 }
 
-func (s *Scanner) Destroy() {
-	close(s.scanExit)
-}
-
 func NewScanner(period time.Duration) *Scanner {
 	return &Scanner{
-		period:   period,
-		subs:     make(map[string]*ScannerSubscriber),
-		scanExit: make(chan struct{}),
+		period: period,
+		subs:   make(map[string]*ScannerSubscriber),
 	}
 }
 
@@ -51,7 +47,8 @@ func (s *Scanner) Unsubscribe(sub *ScannerSubscriber) {
 }
 
 // ScanAndBroadcast starts scanning open TCP ports and broadcasts every open port to all subscribers.
-func (s *Scanner) ScanAndBroadcast() {
+// It exits when ctx is cancelled.
+func (s *Scanner) ScanAndBroadcast(ctx context.Context) {
 	for {
 		// Read directly from /proc/net/tcp and /proc/net/tcp6 instead of
 		// using gopsutil's net.Connections(), which walks /proc/{pid}/fd
@@ -60,15 +57,14 @@ func (s *Scanner) ScanAndBroadcast() {
 
 		s.mu.RLock()
 		for _, sub := range s.subs {
-			sub.Signal(conns)
+			sub.Signal(ctx, conns)
 		}
 		s.mu.RUnlock()
 
 		select {
-		case <-s.scanExit:
+		case <-ctx.Done():
 			return
-		default:
-			time.Sleep(s.period)
+		case <-time.After(s.period):
 		}
 	}
 }

envd/internal/port/scanSubscriber.go

@@ -1,8 +1,11 @@
 // SPDX-License-Identifier: Apache-2.0
+// Modifications by M/S Omukk
 
 package port
 
 import (
+	"context"
+
 	"github.com/rs/zerolog"
 )
 
@@ -33,19 +36,26 @@ func (ss *ScannerSubscriber) Destroy() {
 	close(ss.Messages)
 }
 
-func (ss *ScannerSubscriber) Signal(conns []ConnStat) {
-	// Filter isn't specified. Accept everything.
+// Signal sends the (filtered) connection list to the subscriber. It respects
+// ctx cancellation so the scanner goroutine is never stuck waiting for a
+// consumer that has already exited.
+func (ss *ScannerSubscriber) Signal(ctx context.Context, conns []ConnStat) {
+	var payload []ConnStat
+
 	if ss.filter == nil {
-		ss.Messages <- conns
+		payload = conns
 	} else {
 		filtered := []ConnStat{}
 		for i := range conns {
-			// We need to access the list directly otherwise there will be implicit memory aliasing
-			// If the filter matched a connection, we will send it to a channel.
 			if ss.filter.Match(&conns[i]) {
 				filtered = append(filtered, conns[i])
 			}
 		}
-		ss.Messages <- filtered
+		payload = filtered
+	}
+
+	select {
+	case ss.Messages <- payload:
+	case <-ctx.Done():
 	}
 }

envd/internal/port/subsystem.go

@@ -0,0 +1,106 @@
+// SPDX-License-Identifier: Apache-2.0
+// Modifications by M/S Omukk
+
+package port
+
+import (
+	"context"
+	"runtime"
+	"runtime/debug"
+	"sync"
+	"time"
+
+	"github.com/rs/zerolog"
+
+	"git.omukk.dev/wrenn/sandbox/envd/internal/services/cgroups"
+)
+
+// PortSubsystem owns the port scanner and forwarder lifecycle.
+// It supports stop/restart across Firecracker snapshot/restore cycles.
+type PortSubsystem struct {
+	logger        *zerolog.Logger
+	cgroupManager cgroups.Manager
+	period        time.Duration
+
+	mu      sync.Mutex
+	cancel  context.CancelFunc
+	wg      *sync.WaitGroup // per-cycle WaitGroup; nil when not running
+	running bool
+}
+
+// NewPortSubsystem creates a new PortSubsystem. Call Start() to begin scanning.
+func NewPortSubsystem(logger *zerolog.Logger, cgroupManager cgroups.Manager, period time.Duration) *PortSubsystem {
+	return &PortSubsystem{
+		logger:        logger,
+		cgroupManager: cgroupManager,
+		period:        period,
+	}
+}
+
+// Start creates a fresh scanner and forwarder, launching their goroutines.
+// Safe to call multiple times; does nothing if already running.
+func (p *PortSubsystem) Start(parentCtx context.Context) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	if p.running {
+		return
+	}
+
+	ctx, cancel := context.WithCancel(parentCtx)
+	p.cancel = cancel
+	p.running = true
+
+	// Allocate a fresh WaitGroup for this lifecycle so a concurrent Stop
+	// on the previous cycle's WaitGroup cannot interfere.
+	wg := &sync.WaitGroup{}
+	p.wg = wg
+
+	scanner := NewScanner(p.period)
+	forwarder := NewForwarder(p.logger, scanner, p.cgroupManager)
+
+	wg.Add(2)
+
+	go func() {
+		defer wg.Done()
+		forwarder.StartForwarding(ctx)
+	}()
+
+	go func() {
+		defer wg.Done()
+		scanner.ScanAndBroadcast(ctx)
+	}()
+}
+
+// Stop quiesces the scanner and forwarder goroutines and forces a GC cycle
+// to put the Go runtime's page allocator in a consistent state before snapshot.
+// Blocks until both goroutines have exited. Safe to call if already stopped.
+func (p *PortSubsystem) Stop() {
+	p.mu.Lock()
+	if !p.running {
+		p.mu.Unlock()
+		return
+	}
+	cancelFn := p.cancel
+	wg := p.wg
+	p.cancel = nil
+	p.wg = nil
+	p.running = false
+	p.mu.Unlock()
+
+	cancelFn()
+	wg.Wait()
+
+	// Force two GC cycles to ensure all spans are swept and the page
+	// allocator summary tree is fully consistent before the VM is frozen.
+	runtime.GC()
+	runtime.GC()
+	debug.FreeOSMemory()
+}
+
+// Restart stops the subsystem (if running) and starts it again with a fresh
+// scanner and forwarder. Used after snapshot restore via PostInit.
+func (p *PortSubsystem) Restart(parentCtx context.Context) {
+	p.Stop()
+	p.Start(parentCtx)
+}