add --event-chan-size and --delay-before-handle-packet-events for imp…

…rove e2e test
mozillazg · May 5, 2024 · d94507e · d94507e
1 parent 502390d
commit d94507e
Show file tree

Hide file tree

Showing 7 changed files with 23 additions and 10 deletions.
diff --git a/bpf/event.go b/bpf/event.go
@@ -11,12 +11,12 @@ import (
 	"unsafe"
 )
 
-func (b *BPF) PullPacketEvents(ctx context.Context) (<-chan BpfPacketEventT, error) {
+func (b *BPF) PullPacketEvents(ctx context.Context, chanSize int) (<-chan BpfPacketEventT, error) {
 	reader, err := perf.NewReader(b.objs.PacketEvents, 1500*1000)
 	if err != nil {
 		return nil, xerrors.Errorf(": %w", err)
 	}
-	ch := make(chan BpfPacketEventT, 10)
+	ch := make(chan BpfPacketEventT, chanSize)
 	go func() {
 		defer close(ch)
 		defer reader.Close()
@@ -63,12 +63,12 @@ func parsePacketEvent(rawSample []byte) (*BpfPacketEventT, error) {
 	return &event, nil
 }
 
-func (b *BPF) PullExecEvents(ctx context.Context) (<-chan BpfExecEventT, error) {
+func (b *BPF) PullExecEvents(ctx context.Context, chanSize int) (<-chan BpfExecEventT, error) {
 	reader, err := perf.NewReader(b.objs.ExecEvents, 1024*256)
 	if err != nil {
 		return nil, xerrors.Errorf(": %w", err)
 	}
-	ch := make(chan BpfExecEventT, 10)
+	ch := make(chan BpfExecEventT, chanSize)
 	go func() {
 		defer close(ch)
 		defer reader.Close()

diff --git a/cmd/capture.go b/cmd/capture.go
@@ -2,9 +2,11 @@ package cmd
 
 import (
 	"context"
+	"log"
+	"time"
+
 	"github.com/mozillazg/ptcpdump/internal/consumer"
 	"github.com/mozillazg/ptcpdump/internal/metadata"
-	"log"
 )
 
 func capture(ctx context.Context, opts Options) error {
@@ -29,11 +31,11 @@ func capture(ctx context.Context, opts Options) error {
 	}
 	defer bf.Close()
 
-	packetEvensCh, err := bf.PullPacketEvents(ctx)
+	packetEvensCh, err := bf.PullPacketEvents(ctx, int(opts.eventChanSize))
 	if err != nil {
 		return err
 	}
-	execEvensCh, err := bf.PullExecEvents(ctx)
+	execEvensCh, err := bf.PullExecEvents(ctx, int(opts.eventChanSize))
 	if err != nil {
 		return err
 	}
@@ -44,6 +46,9 @@ func capture(ctx context.Context, opts Options) error {
 	log.Println("capturing...")
 
 	packetConsumer := consumer.NewPacketEventConsumer(writers)
+	if opts.delayBeforeHandlePacketEvents > 0 {
+		time.Sleep(opts.delayBeforeHandlePacketEvents)
+	}
 	packetConsumer.Start(ctx, packetEvensCh, opts.maxPacketCount)
 
 	return nil

diff --git a/cmd/options.go b/cmd/options.go
@@ -1,5 +1,7 @@
 package cmd
 
+import "time"
+
 type Options struct {
 	ifaces         []string
 	pid            uint
@@ -13,6 +15,9 @@ type Options struct {
 	print          bool
 	maxPacketCount uint
 	direction      string
+
+	eventChanSize                 uint
+	delayBeforeHandlePacketEvents time.Duration
 }
 
 func (o Options) WritePath() string {

diff --git a/cmd/root.go b/cmd/root.go
@@ -57,6 +57,9 @@ func init() {
 		"Exit after receiving count packets")
 	rootCmd.Flags().StringVarP(&opts.direction, "direction", "Q",
 		"inout", "Choose send/receive direction for which packets should be captured. Possible values are 'in', 'out' and 'inout'")
+	rootCmd.Flags().UintVar(&opts.eventChanSize, "event-chan-size", 10, "Size of event chan")
+	rootCmd.Flags().DurationVar(&opts.delayBeforeHandlePacketEvents, "delay-before-handle-packet-events", 0,
+		"Delay some durations before handle packet events")
 }
 
 func Execute() error {

diff --git a/testdata/test_base.sh b/testdata/test_base.sh
@@ -10,7 +10,7 @@ RNAME="${FILE_PREFIX}_base.read.txt"
 
 
 function test_ptcpdump() {
-  timeout 30s ${CMD} -c 1 -i any --print -w "${FNAME}" \
+  timeout 30s ${CMD} -c 1 -i any --print -w "${FNAME}" --delay-before-handle-packet-events=1s --event-chan-size=1024 \
     'dst host 1.1.1.1 and tcp[tcpflags] = tcp-syn' | tee "${LNAME}" &
   sleep 10
   curl -m 10 1.1.1.1 &>/dev/null || true

diff --git a/testdata/test_pid_filter.sh b/testdata/test_pid_filter.sh
@@ -10,7 +10,7 @@ RNAME="${FILE_PREFIX}_filter_by_pid.read.txt"
 
 
 function test_ptcpdump() {
-  timeout 30s ${CMD} -c 6 --pid $$ -f -i any --print -w "${FNAME}" | tee "${LNAME}" &
+  timeout 30s ${CMD} -c 6 --pid $$ -f -i any --print -w "${FNAME}"  --delay-before-handle-packet-events=1s --event-chan-size=1024 | tee "${LNAME}" &
   sleep 10
   curl -m 10 1.1.1.1 &>/dev/null || true
   wait

diff --git a/testdata/test_pname_filter.sh b/testdata/test_pname_filter.sh
@@ -10,7 +10,7 @@ RNAME="${FILE_PREFIX}_filter_by_pname.read.txt"
 
 
 function test_ptcpdump() {
-  timeout 30s ${CMD} -c 6 --pname curl -f -i any --print -w "${FNAME}" | tee "${LNAME}" &
+  timeout 30s ${CMD} -c 6 --pname curl -f -i any --print -w "${FNAME}"  --delay-before-handle-packet-events=1s --event-chan-size=1024 | tee "${LNAME}" &
   sleep 10
   curl -m 10 1.1.1.1 &>/dev/null || true
   wait