engine: more performance improvements

engine/udp.go

@@ -2,6 +2,7 @@ package engine
 
 import (
 	"bytes"
+	"container/list"
 	"errors"
 	"net"
 	"sync"
@@ -12,7 +13,6 @@ import (
 	"git.difuse.io/Difuse/Mellaris/ruleset"
 
 	"github.com/bwmarrin/snowflake"
-	lru "github.com/hashicorp/golang-lru/v2"
 )
 
 // udpVerdict is a subset of io.Verdict for UDP streams.
@@ -116,15 +116,18 @@ func (f *udpStreamFactory) currentRuleset() (ruleset.Ruleset, uint64) {
 
 type udpStreamManager struct {
 	factory      *udpStreamFactory
-	streams      *lru.Cache[uint32, *udpStreamValue]
+	streams      map[uint32]*list.Element
+	order        *list.List
+	maxStreams   int
 	tupleIndex   map[udpTupleKey]uint32
 	streamTuples map[uint32]udpTupleKey
 	stats        *statsCounters
 }
 
 type udpStreamValue struct {
-	Stream *udpStream
-	Tuple  udpTupleKey
+	StreamID uint32
+	Stream   *udpStream
+	Tuple    udpTupleKey
 }
 
 func (v *udpStreamValue) Match(k udpTupleKey) (ok, rev bool) {
@@ -143,27 +146,23 @@ type udpTupleKey struct {
 }
 
 func newUDPStreamManager(factory *udpStreamFactory, maxStreams int, stats *statsCounters) (*udpStreamManager, error) {
+	if maxStreams <= 0 {
+		maxStreams = 1
+	}
 	m := &udpStreamManager{
 		factory:      factory,
+		streams:      make(map[uint32]*list.Element, maxStreams),
+		order:        list.New(),
+		maxStreams:   maxStreams,
 		tupleIndex:   make(map[udpTupleKey]uint32, maxStreams),
 		streamTuples: make(map[uint32]udpTupleKey, maxStreams),
 		stats:        stats,
 	}
-	ss, err := lru.NewWithEvict[uint32, *udpStreamValue](maxStreams, func(k uint32, v *udpStreamValue) {
-		if v != nil && v.Stream != nil {
-			v.Stream.Close()
-		}
-		m.removeTupleMappingLocked(k)
-	})
-	if err != nil {
-		return nil, err
-	}
-	m.streams = ss
 	return m, nil
 }
 
 func (m *udpStreamManager) MatchWithContext(streamID uint32, tuple udpTupleKey, rev bool, payload []byte, uc *udpContext) {
-	value, ok := m.streams.Get(streamID)
+	value, ok := m.get(streamID)
 	if !ok {
 		if m.stats != nil {
 			m.stats.UDPTupleLookups.Add(1)
@@ -176,7 +175,7 @@ func (m *udpStreamManager) MatchWithContext(streamID uint32, tuple udpTupleKey,
 				m.stats.UDPTupleHits.Add(1)
 			}
 			var hasValue bool
-			matchedValue, hasValue = m.streams.Get(matchedKey)
+			matchedValue, hasValue = m.get(matchedKey)
 			if !hasValue || matchedValue == nil {
 				delete(m.tupleIndex, tuple)
 				delete(m.streamTuples, matchedKey)
@@ -188,16 +187,18 @@ func (m *udpStreamManager) MatchWithContext(streamID uint32, tuple udpTupleKey,
 			value = matchedValue
 			rev = matchedRev
 			if matchedKey != streamID {
-				m.streams.Remove(matchedKey)
-				m.streams.Add(streamID, matchedValue)
+				m.remove(matchedKey, false)
+				matchedValue.StreamID = streamID
+				m.add(streamID, matchedValue)
 				m.bindTupleLocked(streamID, tuple)
 			}
 		} else {
 			value = &udpStreamValue{
-				Stream: m.factory.New(tuple, payload, uc),
-				Tuple:  tuple,
+				StreamID: streamID,
+				Stream:   m.factory.New(tuple, payload, uc),
+				Tuple:    tuple,
 			}
-			m.streams.Add(streamID, value)
+			m.add(streamID, value)
 			m.bindTupleLocked(streamID, tuple)
 		}
 	} else {
@@ -205,10 +206,11 @@ func (m *udpStreamManager) MatchWithContext(streamID uint32, tuple udpTupleKey,
 		if !ok {
 			value.Stream.Close()
 			value = &udpStreamValue{
-				Stream: m.factory.New(tuple, payload, uc),
-				Tuple:  tuple,
+				StreamID: streamID,
+				Stream:   m.factory.New(tuple, payload, uc),
+				Tuple:    tuple,
 			}
-			m.streams.Add(streamID, value)
+			m.add(streamID, value)
 			m.bindTupleLocked(streamID, tuple)
 		}
 	}
@@ -217,6 +219,55 @@ func (m *udpStreamManager) MatchWithContext(streamID uint32, tuple udpTupleKey,
 	}
 }
 
+func (m *udpStreamManager) get(streamID uint32) (*udpStreamValue, bool) {
+	ele, ok := m.streams[streamID]
+	if !ok || ele == nil {
+		return nil, false
+	}
+	m.order.MoveToFront(ele)
+	value, ok := ele.Value.(*udpStreamValue)
+	return value, ok && value != nil
+}
+
+func (m *udpStreamManager) add(streamID uint32, value *udpStreamValue) {
+	if value == nil {
+		return
+	}
+	if existing, ok := m.streams[streamID]; ok {
+		existing.Value = value
+		m.order.MoveToFront(existing)
+		return
+	}
+	value.StreamID = streamID
+	m.streams[streamID] = m.order.PushFront(value)
+	for len(m.streams) > m.maxStreams {
+		back := m.order.Back()
+		if back == nil {
+			return
+		}
+		evicted, _ := back.Value.(*udpStreamValue)
+		if evicted == nil {
+			m.order.Remove(back)
+			continue
+		}
+		m.remove(evicted.StreamID, true)
+	}
+}
+
+func (m *udpStreamManager) remove(streamID uint32, closeStream bool) {
+	ele, ok := m.streams[streamID]
+	if !ok || ele == nil {
+		return
+	}
+	value, _ := ele.Value.(*udpStreamValue)
+	delete(m.streams, streamID)
+	m.order.Remove(ele)
+	m.removeTupleMappingLocked(streamID)
+	if closeStream && value != nil && value.Stream != nil {
+		value.Stream.Close()
+	}
+}
+
 func (m *udpStreamManager) bindTupleLocked(streamID uint32, key udpTupleKey) {
 	m.removeTupleMappingLocked(streamID)
 	m.tupleIndex[key] = streamID