v0.1.0 (#17)

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):
 		}
 	}
 }