refactor: engine/tcp/worker perf improvements

engine/engine.go

@@ -2,16 +2,11 @@ package engine
 
 import (
 	"context"
-	"encoding/binary"
-	"runtime"
 	"sync"
 	"sync/atomic"
 
 	"git.difuse.io/Difuse/Mellaris/io"
 	"git.difuse.io/Difuse/Mellaris/ruleset"
-
-	"github.com/google/gopacket"
-	"github.com/google/gopacket/layers"
 )
 
 var _ Engine = (*engine)(nil)
@@ -27,12 +22,15 @@ type engine struct {
 	workers     []*worker
 	verdicts    sync.Map     // streamID(uint32) → verdictEntry
 	verdictsGen atomic.Int64 // incremented on ruleset update
+
+	overflowCh   chan *workerPacket
+	overflowOnce sync.Once
 }
 
 func NewEngine(config Config) (Engine, error) {
 	workerCount := config.Workers
 	if workerCount <= 0 {
-		workerCount = runtime.NumCPU()
+		workerCount = 1
 	}
 	macResolver := newSourceMACResolver()
 	var err error
@@ -53,9 +51,10 @@ func NewEngine(config Config) (Engine, error) {
 		}
 	}
 	e := &engine{
-		logger:  config.Logger,
-		io:      config.IO,
-		workers: workers,
+		logger:     config.Logger,
+		io:         config.IO,
+		workers:    workers,
+		overflowCh: make(chan *workerPacket, 1024),
 	}
 	return e, nil
 }
@@ -75,6 +74,10 @@ func (e *engine) Run(ctx context.Context) error {
 	ioCtx, ioCancel := context.WithCancel(ctx)
 	defer ioCancel()
 
+	e.overflowOnce.Do(func() {
+		go e.drainOverflow(ioCtx)
+	})
+
 	for _, w := range e.workers {
 		go w.Run(ioCtx)
 	}
@@ -111,55 +114,55 @@ func (e *engine) dispatch(p io.Packet) bool {
 	}
 
 	data := p.Data()
-	layerType, srcMAC, dstMAC, ok := classifyPacket(data)
-	if !ok {
+	if !validPacket(data) {
 		_ = e.io.SetVerdict(p, io.VerdictAcceptStream, nil)
 		return true
 	}
 	gen := e.verdictsGen.Load()
 	index := streamID % uint32(len(e.workers))
-	e.workers[index].Feed(&workerPacket{
+	wp := &workerPacket{
 		StreamID:  streamID,
 		Data:      data,
-		LayerType: layerType,
-		SrcMAC:    srcMAC,
-		DstMAC:    dstMAC,
 		SetVerdict: func(v io.Verdict, b []byte) error {
 			if v == io.VerdictAcceptStream || v == io.VerdictDropStream {
 				e.verdicts.Store(streamID, verdictEntry{Verdict: v, Gen: gen})
 			}
 			return e.io.SetVerdict(p, v, b)
 		},
-	})
+	}
+	if !e.workers[index].Feed(wp) {
+		select {
+		case e.overflowCh <- wp:
+		default:
+		}
+	}
 	return true
 }
 
-// classifyPacket detects packet framing and returns a gopacket decode layer
-// plus best-effort source/destination MAC addresses when available.
-func classifyPacket(data []byte) (gopacket.LayerType, []byte, []byte, bool) {
+func validPacket(data []byte) bool {
 	if len(data) == 0 {
-		return 0, nil, nil, false
+		return false
 	}
-
-	// Fast path for IP packets (NFQUEUE payloads are typically IP-only).
 	ipVersion := data[0] >> 4
-	if ipVersion == 4 {
-		return layers.LayerTypeIPv4, nil, nil, true
+	if ipVersion == 4 || ipVersion == 6 {
+		return true
 	}
-	if ipVersion == 6 {
-		return layers.LayerTypeIPv6, nil, nil, true
-	}
-
-	// Ethernet frame path (for custom PacketIO implementations).
 	if len(data) >= 14 {
-		etherType := binary.BigEndian.Uint16(data[12:14])
-		if etherType == uint16(layers.EthernetTypeIPv4) || etherType == uint16(layers.EthernetTypeIPv6) {
-			return layers.LayerTypeEthernet,
-				append([]byte(nil), data[6:12]...),
-				append([]byte(nil), data[:6]...),
-				true
+		etherType := uint16(data[12])<<8 | uint16(data[13])
+		if etherType == 0x0800 || etherType == 0x86DD {
+			return true
+		}
+	}
+	return false
+}
+
+func (e *engine) drainOverflow(ctx context.Context) {
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case wp := <-e.overflowCh:
+			_ = wp.SetVerdict(io.VerdictAccept, nil)
 		}
 	}
-
-	return 0, nil, nil, false
 }