chore: add ability to limit the number of events that are output (#15)

Author: mikhailswift

dns.go

@@ -16,7 +16,7 @@ import (
 	"github.com/google/gopacket/layers"
 )
 
-func processUDPPackets(ctx context.Context, reader *ringbuf.Reader) {
+func processUDPPackets(ctx context.Context, reader *ringbuf.Reader, outputCh chan<- []statEntry) {
 	seenDNSPackets := map[statEntry]struct{}{}
 	resetSeenPacketsTick := time.NewTicker(time.Minute)
 	defer resetSeenPacketsTick.Stop()
@@ -32,11 +32,9 @@ func processUDPPackets(ctx context.Context, reader *ringbuf.Reader) {
 			seenDNSPackets = map[statEntry]struct{}{}
 		default:
 			record, err := reader.Read()
-			if err != nil {
-				if errors.Is(err, ringbuf.ErrClosed) {
-					return
-				}
-
+			if errors.Is(err, ringbuf.ErrClosed) {
+				return
+			} else if err != nil {
 				log.Printf("Error reading UDP Packet: %v", err)
 				continue
 			}
@@ -97,7 +95,7 @@ func processUDPPackets(ctx context.Context, reader *ringbuf.Reader) {
 			}
 
 			entry.Timestamp = time.Now().UTC()
-			fmt.Print(outputJSON([]statEntry{entry}))
+			outputCh <- []statEntry{entry}
 		}
 	}
 }

flags.go

@@ -32,6 +32,7 @@ import (
 var (
 	uniqueOutput, version, help, externalOnly *bool
 	kubeconfig                                *string
+	maxEvents                                 int
 )
 
 func parseFlags() {
@@ -41,6 +42,7 @@ func parseFlags() {
 	uniqueOutput = fs.Bool('u', "unique", "if true, only show the first instance of each connection")
 	kubeconfig = fs.StringLong("kubeconfig", "", "path to kubeconfig file (if not set, dynamic kubeconfig discovery is used)")
 	externalOnly = fs.BoolDefault('e', "external", true, "if true, only show traffic to external destinations")
+	maxEventsPtr := fs.Int('m', "max-events", 10000, "max number of events to output, 0 for unlimited")
 	version = fs.BoolLong("version", "display program version")
 
 	var err error
@@ -63,4 +65,6 @@ func parseFlags() {
 
 		os.Exit(0)
 	}
+
+	maxEvents = *maxEventsPtr
 }

go.mod

@@ -2,7 +2,7 @@ module github.com/dkorunic/pktstat-bpf
 
 go 1.24.0
 
-toolchain go1.24.1
+toolchain go1.25.4
 
 require (
 	github.com/cilium/ebpf v0.18.0

main.go

@@ -142,31 +142,51 @@ func main() {
 	signalCh := make(chan os.Signal, 1)
 	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM, syscall.SIGINT)
 
-	// Create a ticker to process the map every second
-	ticker := time.NewTicker(1 * time.Second)
-	defer ticker.Stop()
-
 	// Set up signal handler
 	go func() {
 		s := <-signalCh
 		_, _ = fmt.Fprintf(os.Stderr, "Received %v signal, exiting...\n", s)
 		cancel()
 	}()
 
-	// Create DNS lookup map and mutex for sharing between goroutines
-	dnsLookupMap := make(map[uint32]string)
-	dnsLookupMapMutex := &sync.RWMutex{}
+	wg := &sync.WaitGroup{}
+	outputCh := make(chan []statEntry)
+	go outputEvents(outputCh)
+	defer close(outputCh)
 
 	udpPktReader, err := ringbuf.NewReader(objs.UdpPkts)
 	if err != nil {
 		log.Printf("Failed to create ringbuf reader for UDP packets: %v", err)
 	} else {
 		log.Printf("Created UDP packet ringbuf reader successfully")
-		go processUDPPackets(ctx, udpPktReader)
-		defer udpPktReader.Close()
+		wg.Go(func() {
+			go func() {
+				<-ctx.Done()
+				udpPktReader.Close()
+			}()
+
+			processUDPPackets(ctx, udpPktReader, outputCh)
+		})
 	}
 
 	// Run the main loop
+	wg.Go(func() {
+		processPktCountMap(ctx, &objs, outputCh)
+	})
+
+	wg.Wait()
+}
+
+// processPktCountMap starts processing the PktCount map for connection events
+func processPktCountMap(ctx context.Context, objs *counterObjects, outputCh chan<- []statEntry) {
+	// Create a ticker to process the map every second
+	ticker := time.NewTicker(1 * time.Second)
+	defer ticker.Stop()
+
+	// Create DNS lookup map and mutex for sharing between goroutines
+	dnsLookupMap := make(map[uint32]string)
+	dnsLookupMapMutex := &sync.RWMutex{}
+
 	seenEntries := make(map[string]bool)
 	for {
 		select {
@@ -226,23 +246,41 @@ func main() {
 				continue
 			}
 
-			// Format output as JSON Lines
-			output := outputJSON(newEntries)
-
-			// Add newline if needed
-			if output != "" && !strings.HasSuffix(output, "\n") {
-				output += "\n"
-			}
-
-			// Write output to stdout
-			fmt.Print(output)
+			outputCh <- newEntries
 
 		case <-ctx.Done():
 			return
 		}
 	}
 }
 
+func outputEvents(ch <-chan []statEntry) {
+	numEvents := 0
+
+	for entries := range ch {
+		if maxEvents > 0 && numEvents >= maxEvents {
+			continue
+		}
+		// Format output as JSON Lines
+		output := outputJSON(entries)
+
+		// Add newline if needed
+		if output != "" && !strings.HasSuffix(output, "\n") {
+			output += "\n"
+		}
+
+		// Write output to stdout
+		fmt.Print(output)
+
+		if maxEvents > 0 {
+			numEvents += len(entries)
+			if numEvents >= maxEvents {
+				log.Printf("Max number of events has been met, no longer outputting any events")
+			}
+		}
+	}
+}
+
 // startKProbes attaches a series of eBPF programs to kernel functions using KProbes.
 //
 // This function iterates over a slice of kprobeHook structs, each containing a kernel function