Reduce allocation and unnecessary layers (#822)

The focus here was to **reduce allocation in TCP parser** but speed may have hopeful improved too!
pool no longer use map's key of **string** it uses **uint64**
**Benchmarks** was revamped to be more clear
if you want to compare these results copy the benchmark in tcp/bench_test.go@reduce-allocation to tcp/bench_test.go@master:

**before(master)**:
```
BenchmarkPacketParseAndSort-4         	 1000000	      1006 ns/op	      64 B/op	       2 allocs/op
BenchmarkMessageParserWithoutHint-4   	     625	   1772309 ns/op	      1000 packets/op	  419096 B/op	   10045 allocs/op
BenchmarkMessageParserWithHint-4      	      74	  14969926 ns/op	      1000 chunks/op	      1002 packets/op	  450992 B/op	   10126 allocs/op
```

**After(this branch)**:
```
BenchmarkPacketParseAndSort-4         	 1267662	       941 ns/op	      64 B/op	       2 allocs/op
BenchmarkMessageParserWithoutHint-4   	    2256	    523474 ns/op	      1000 packets/op	  243530 B/op	    1037 allocs/op
BenchmarkMessageParserWithHint-4      	      80	  13990955 ns/op	      1000 chunks/op	      1002 packets/op	  268609 B/op	    1099 allocs/op

```

byteutils/byteutils.go

@@ -2,7 +2,6 @@
 package byteutils
 
 import (
-	"reflect"
 	"unsafe"
 )
 
@@ -48,9 +47,6 @@ func Replace(a []byte, from, to int, new []byte) []byte {
 }
 
 // SliceToString preferred for large body payload (zero allocation and faster)
-func SliceToString(buf *[]byte, s *string) {
-	bHeader := (*reflect.SliceHeader)(unsafe.Pointer(buf))
-	sHeader := (*reflect.StringHeader)(unsafe.Pointer(s))
-	sHeader.Data = bHeader.Data
-	sHeader.Len = bHeader.Len
+func SliceToString(buf []byte) string {
+	return *(*string)(unsafe.Pointer(&buf))
 }

byteutils/byteutils_test.go

@@ -32,9 +32,10 @@ func TestReplace(t *testing.T) {
 }
 
 func BenchmarkStringtoSlice(b *testing.B) {
-	b.StopTimer()
-	buf := make([]byte, b.N)
-	b.StartTimer()
-	s := new(string)
-	SliceToString(&buf, s)
+	var s string
+	var buf [1 << 20]byte
+	for i := 0; i < b.N; i++ {
+		s = SliceToString(buf[:])
+	}
+	_ = s // avoid gc to optimize away the loop body
 }

input_raw.go

@@ -126,14 +126,19 @@ func (i *RAWInput) Read(data []byte) (n int, err error) {
 	if len(data) > len(header) {
 		n += copy(data[len(header):], buf)
 	}
-	dis := len(header) + len(buf) - n
-	if dis > 0 {
-		go Debug(2, "[INPUT-RAW] discarded", dis, "bytes increase copy buffer size")
+	// to be removed....
+	if msg.Truncated || len(header)+len(buf)-n > 0 {
+		go Debug(2, "[INPUT-RAW] message truncated, increase copy-buffer-size")
 	}
-	if msg.Truncated {
-		go Debug(2, "[INPUT-RAW] message truncated, copy-buffer-size")
+	// to be removed...
+	if msg.TimedOut {
+		go Debug(2, "[INPUT-RAW] message timeout reached, increase input-raw-expire")
 	}
-	go i.addStats(msg.Stats)
+	if i.Stats {
+		stat := msg.Stats
+		go i.addStats(stat)
+	}
+	msg = nil
 	return n, nil
 }
 
@@ -188,19 +193,20 @@ func (i *RAWInput) Close() error {
 }
 
 func (i *RAWInput) addStats(mStats tcp.Stats) {
-	if i.Stats {
-		i.Lock()
-		if len(i.messageStats) >= 10000 {
-			i.messageStats = []tcp.Stats{}
-		}
-		i.messageStats = append(i.messageStats, mStats)
-
-		i.Unlock()
+	i.Lock()
+	if len(i.messageStats) >= 10000 {
+		i.messageStats = []tcp.Stats{}
 	}
+	i.messageStats = append(i.messageStats, mStats)
+	i.Unlock()
 }
 
 func startHint(pckt *tcp.Packet) (isIncoming, isOutgoing bool) {
-	return proto.HasRequestTitle(pckt.Payload), proto.HasResponseTitle(pckt.Payload)
+	isIncoming = proto.HasRequestTitle(pckt.Payload)
+	if isIncoming {
+		return
+	}
+	return false, proto.HasResponseTitle(pckt.Payload)
 }
 
 func endHint(m *tcp.Message) bool {

output_dummy.go

@@ -1,7 +1,7 @@
 package main
 
 import (
-	"fmt"
+	"os"
 )
 
 // DummyOutput used for debugging, prints all incoming requests
@@ -16,9 +16,9 @@ func NewDummyOutput() (di *DummyOutput) {
 }
 
 func (i *DummyOutput) Write(data []byte) (int, error) {
-	fmt.Println(string(data))
-
-	return len(data), nil
+	n, err := os.Stdout.Write(data)
+	os.Stdout.Write([]byte{'\n'})
+	return n, err
 }
 
 func (i *DummyOutput) String() string {

proto/proto.go

@@ -332,8 +332,7 @@ const (
 
 // HasResponseTitle reports whether this payload has an HTTP/1 response title
 func HasResponseTitle(payload []byte) bool {
-	var s string
-	byteutils.SliceToString(&payload, &s)
+	s := byteutils.SliceToString(payload)
 	if len(s) < MinResponseCount {
 		return false
 	}
@@ -361,8 +360,7 @@ func HasResponseTitle(payload []byte) bool {
 
 // HasRequestTitle reports whether this payload has an HTTP/1 request title
 func HasRequestTitle(payload []byte) bool {
-	var s string
-	byteutils.SliceToString(&payload, &s)
+	s := byteutils.SliceToString(payload)
 	if len(s) < MinRequestCount {
 		return false
 	}

proto/proto_test.go

@@ -2,10 +2,8 @@ package proto
 
 import (
 	"bytes"
-	"fmt"
 	"reflect"
 	"testing"
-	"time"
 )
 
 func TestHeader(t *testing.T) {
@@ -458,28 +456,21 @@ func TestHasFullPayload(t *testing.T) {
 }
 
 func BenchmarkHasFullPayload(b *testing.B) {
-	now := time.Now()
-	payload := make([]byte, 0xfc00)
-	for i := 0; i < 0xfc00; i++ {
-		payload[i] = '1'
+	var buf bytes.Buffer
+	buf.Write([]byte("HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nTransfer-Encoding: chunked\r\n\r\n"))
+	var chunk = []byte("1e\r\n111111111111111111111111111111\r\n")
+	for i := 0; i < 5000; i++ {
+		buf.Write(chunk)
 	}
+	buf.Write([]byte("0\r\n\r\n"))
+	data := buf.Bytes()
 	var ok bool
-	data := []byte("HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nTransfer-Encoding: chunked\r\n\r\n")
-	if ok = HasFullPayload(data); ok {
-		b.Error("HasFullPayload should fail")
-		return
-	}
+	b.ResetTimer() // ignores the upper initialization
+	b.ReportMetric(float64(5000), "chunks/op")
 	for i := 0; i < b.N; i++ {
-		data = append(data, []byte(fmt.Sprintf("fc00\r\n%s\r\n", payload))...)
-		if ok = HasFullPayload(data); ok {
-			b.Error("HasFullPayload should fail")
-			return
-		}
+		ok = HasFullPayload(data)
 	}
-	data = append(data, []byte("0\r\n\r\n")...)
-	if ok = HasFullPayload(data); !ok {
-		b.Error("HasFullPayload should pass")
-		return
+	if !ok {
+		b.Fail()
 	}
-	b.Logf("%dKB chunks in %s", b.N*64, time.Since(now))
 }

tcp/tcp_message.go

@@ -1,6 +1,7 @@
 package tcp
 
 import (
+	"bytes"
 	"crypto/sha1"
 	"encoding/hex"
 	"fmt"
@@ -30,7 +31,7 @@ type Stats struct {
 type Message struct {
 	packets []*Packet
 	done    chan bool
-	data    []byte
+	buf     bytes.Buffer
 	Stats
 }
 
@@ -44,7 +45,7 @@ func NewMessage(srcAddr, dstAddr string, ipVersion uint8) (m *Message) {
 	return
 }
 
-// UUID the unique id of a TCP session it is not granted to be unique overtime
+// UUID the unique id of a TCP session it is not granted to be unique!
 func (m *Message) UUID() []byte {
 	var src, dst string
 	if m.IsIncoming {
@@ -70,7 +71,9 @@ func (m *Message) add(pckt *Packet) {
 	m.Length += len(pckt.Payload)
 	m.LostData += int(pckt.Lost)
 	m.packets = append(m.packets, pckt)
-	m.data = append(m.data, pckt.Payload...)
+	if len(pckt.Payload) > 0 {
+		m.buf.Write(pckt.Payload)
+	}
 	m.End = pckt.Timestamp
 }
 
@@ -81,7 +84,7 @@ func (m *Message) Packets() []*Packet {
 
 // Data returns data in this message
 func (m *Message) Data() []byte {
-	return m.data
+	return m.buf.Bytes()
 }
 
 // Sort a helper to sort packets
@@ -105,13 +108,12 @@ type HintEnd func(*Message) bool
 type HintStart func(*Packet) (IsIncoming, IsOutgoing bool)
 
 // MessagePool holds data of all tcp messages in progress(still receiving/sending packets).
-// Incoming message is identified by its source port and address e.g: 127.0.0.1:45785.
-// Outgoing message is identified by  server.addr and dst.addr e.g: localhost:80=internet:45785.
+// message is identified by its source port and dst port, and last 4bytes of src IP.
 type MessagePool struct {
 	sync.Mutex
 	debug         Debugger
 	maxSize       size.Size // maximum message size, default 5mb
-	pool          map[string]*Message
+	pool          map[uint64]*Message
 	handler       Handler
 	messageExpire time.Duration // the maximum time to wait for the final packet, minimum is 100ms
 	End           HintEnd
@@ -131,7 +133,7 @@ func NewMessagePool(maxSize size.Size, messageExpire time.Duration, debugger Deb
 	if pool.maxSize < 1 {
 		pool.maxSize = 5 << 20
 	}
-	pool.pool = make(map[string]*Message)
+	pool.pool = make(map[uint64]*Message)
 	return pool
 }
 
@@ -145,20 +147,23 @@ func (pool *MessagePool) Handler(packet gopacket.Packet) {
 	}
 	pool.Lock()
 	defer pool.Unlock()
-	srcKey := pckt.Src()
-	dstKey := srcKey + "=" + pckt.Dst()
-	m, ok := pool.pool[srcKey]
-	if !ok {
-		m, ok = pool.pool[dstKey]
+	lst := 3
+	if pckt.Version == 6 {
+		lst = 15
 	}
+	key := uint64(pckt.SrcPort)<<48 | uint64(pckt.DstPort)<<32 |
+		uint64(pckt.SrcIP[lst])<<24 | uint64(pckt.SrcIP[lst-1])<<16 |
+		uint64(pckt.SrcIP[lst-2])<<8 | uint64(pckt.SrcIP[lst-3])
+	m, ok := pool.pool[key]
 	if pckt.RST {
 		if ok {
 			m.done <- true
 			<-m.done
 		}
-		if m, ok = pool.pool[pckt.Dst()]; !ok {
-			m, ok = pool.pool[pckt.Dst()+"="+srcKey]
-		}
+		key = uint64(pckt.DstPort)<<48 | uint64(pckt.SrcPort)<<32 |
+			uint64(pckt.DstIP[lst])<<24 | uint64(pckt.DstIP[lst-1])<<16 |
+			uint64(pckt.DstIP[lst-2])<<8 | uint64(pckt.DstIP[lst-3])
+		m, ok = pool.pool[key]
 		if ok {
 			m.done <- true
 			<-m.done
@@ -179,25 +184,25 @@ func (pool *MessagePool) Handler(packet gopacket.Packet) {
 	default:
 		return
 	}
-	m = NewMessage(srcKey, pckt.Dst(), pckt.Version)
+	m = NewMessage(pckt.Src(), pckt.Dst(), pckt.Version)
 	m.IsIncoming = in
-	key := srcKey
-	if !m.IsIncoming {
-		key = dstKey
-	}
 	pool.pool[key] = m
 	m.Start = pckt.Timestamp
 	go pool.dispatch(key, m)
 	pool.addPacket(m, pckt)
 }
 
-func (pool *MessagePool) dispatch(key string, m *Message) {
+func (pool *MessagePool) dispatch(key uint64, m *Message) {
 	select {
 	case <-m.done:
-		defer func() { m.done <- true }()
+		defer func() { m.done <- true }() // signal that message was dispatched
 	case <-time.After(pool.messageExpire):
 		pool.Lock()
 		defer pool.Unlock()
+		// avoid dispathing message twice
+		if _, ok := pool.pool[key]; !ok {
+			return
+		}
 		m.TimedOut = true
 	}
 	delete(pool.pool, key)

tcp/tcp_packet.go

@@ -21,8 +21,11 @@ type Packet struct {
 	gopacket.LinkLayer
 
 	// IP Header
-	gopacket.NetworkLayer
 	Version uint8 // Ip version
+	SrcIP   net.IP
+	DstIP   net.IP
+	IHL     uint8
+	Length  uint16
 
 	// TCP Segment Header
 	*layers.TCP
@@ -35,10 +38,12 @@ type Packet struct {
 // ParsePacket parse raw packets
 func ParsePacket(packet gopacket.Packet) (pckt *Packet, err error) {
 	// early check of error
+	if packet == nil {
+		return
+	}
 	defer func() {
 		if packet.ErrorLayer() != nil {
 			err = packet.ErrorLayer().Error()
-			println(err.Error())
 			return
 		}
 	}()
@@ -55,11 +60,17 @@ func ParsePacket(packet gopacket.Packet) (pckt *Packet, err error) {
 
 	// parsing network layer
 	if net4, ok := packet.NetworkLayer().(*layers.IPv4); ok {
-		pckt.NetworkLayer = net4
 		pckt.Version = 4
+		pckt.SrcIP = net4.SrcIP
+		pckt.DstIP = net4.DstIP
+		pckt.IHL = net4.IHL * 4
+		pckt.Length = net4.Length
 	} else if net6, ok := packet.NetworkLayer().(*layers.IPv6); ok {
-		pckt.NetworkLayer = net6
 		pckt.Version = 6
+		pckt.SrcIP = net6.SrcIP
+		pckt.DstIP = net6.DstIP
+		pckt.IHL = 40
+		pckt.Length = net6.Length
 	} else {
 		pckt = nil
 		return
@@ -75,56 +86,23 @@ func ParsePacket(packet gopacket.Packet) (pckt *Packet, err error) {
 	pckt.DataOffset *= 4
 
 	// calculating lost data
-	headerSize := int(uint32(pckt.DataOffset) + uint32(pckt.IHL()))
+	headerSize := int(uint32(pckt.DataOffset) + uint32(pckt.IHL))
 	if pckt.Version == 6 {
 		headerSize -= 40 // in ipv6 the length of payload doesn't include the IPheader size
 	}
-	pckt.Lost = pckt.Length() - uint16(headerSize+len(pckt.Payload))
+	pckt.Lost = pckt.Length - uint16(headerSize+len(pckt.Payload))
 
 	return
 }
 
 // Src returns the source socket of a packet
 func (pckt *Packet) Src() string {
-	return fmt.Sprintf("%s:%d", pckt.SrcIP(), pckt.SrcPort)
+	return fmt.Sprintf("%s:%d", pckt.SrcIP, pckt.SrcPort)
 }
 
 // Dst returns destination socket
 func (pckt *Packet) Dst() string {
-	return fmt.Sprintf("%s:%d", pckt.DstIP(), pckt.DstPort)
-}
-
-// SrcIP returns source IP address
-func (pckt *Packet) SrcIP() net.IP {
-	if pckt.Version == 4 {
-		return pckt.NetworkLayer.(*layers.IPv4).SrcIP
-	}
-	return pckt.NetworkLayer.(*layers.IPv6).SrcIP
-}
-
-// DstIP returns destination IP address
-func (pckt *Packet) DstIP() net.IP {
-	if pckt.Version == 4 {
-		return pckt.NetworkLayer.(*layers.IPv4).DstIP
-	}
-	return pckt.NetworkLayer.(*layers.IPv6).DstIP
-}
-
-// IHL returns IP header length in bytes
-func (pckt *Packet) IHL() uint8 {
-	if l, ok := pckt.NetworkLayer.(*layers.IPv4); ok {
-		return l.IHL * 4
-	}
-	// on IPV6 it's constant, https://en.wikipedia.org/wiki/IPv6_packet#Fixed_header
-	return 40
-}
-
-// Length returns the total length of the packet(IP header, TCP header and the actual data)
-func (pckt *Packet) Length() uint16 {
-	if l, ok := pckt.NetworkLayer.(*layers.IPv4); ok {
-		return l.Length
-	}
-	return pckt.NetworkLayer.(*layers.IPv6).Length
+	return fmt.Sprintf("%s:%d", pckt.DstIP, pckt.DstPort)
 }
 
 // SYNOptions returns MSS and windowscale of syn packets
@@ -205,8 +183,8 @@ Lost Data: %d`,
 		pckt.LinkInfo(),
 		pckt.Src(),
 		pckt.Dst(),
-		pckt.IHL(),
-		pckt.Length(),
+		pckt.IHL,
+		pckt.Length,
 		pckt.Seq,
 		pckt.Ack,
 		pckt.DataOffset,

tcp/tcp_test.go

@@ -3,7 +3,6 @@ package tcp
 import (
 	"bytes"
 	"encoding/binary"
-	"fmt"
 	"testing"
 	"time"
 
@@ -205,75 +204,63 @@ func TestMessageUUID(t *testing.T) {
 }
 
 func BenchmarkPacketParseAndSort(b *testing.B) {
-	if b.N < 3 {
-		return
-	}
-	now := time.Now()
 	m := new(Message)
-	m.packets = make([]*Packet, b.N)
-	for i, v := range GetPackets(1, b.N, nil) {
+	m.packets = make([]*Packet, 100)
+	for i, v := range GetPackets(1, 100, nil) {
 		m.packets[i], _ = ParsePacket(v)
 	}
-	m.Sort()
-	b.Logf("%d packets in %s", b.N, time.Since(now))
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		m.Sort()
+	}
 }
 
 func BenchmarkMessageParserWithoutHint(b *testing.B) {
 	var mssg = make(chan *Message, 1)
-	if b.N < 3 {
-		return
-	}
-	now := time.Now()
-	n := b.N
-	packets := GetPackets(1, n, nil)
-	packets[0].Data()[14:][20:][13] = 2     // SYN flag
-	packets[b.N-1].Data()[14:][20:][13] = 1 // FIN flag
+	var chunk = []byte("111111111111111111111111111111")
+	packets := GetPackets(1, 1000, chunk)
+	packets[0].Data()[14:][20:][13] = 2      // SYN flag
+	packets[1000-1].Data()[14:][20:][13] = 1 // FIN flag
 	p := NewMessagePool(1<<20, time.Second*2, nil, func(m *Message) {
-		b.Logf("%d/%d packets in %s", len(m.packets), n, time.Since(now))
 		mssg <- m
 	})
-	for _, v := range packets {
-		p.Handler(v)
+	b.ResetTimer()
+	b.ReportMetric(float64(1000), "packets/op")
+	for i := 0; i < b.N; i++ {
+		for _, v := range packets {
+			p.Handler(v)
+		}
+		<-mssg
 	}
-	<-mssg
 }
 
 func BenchmarkMessageParserWithHint(b *testing.B) {
-	if b.N < 3 {
-		return
+	var buf [1002][]byte
+	var chunk = []byte("1e\r\n111111111111111111111111111111\r\n")
+	buf[0] = []byte("HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nTransfer-Encoding: chunked\r\n\r\n")
+	for i := 1; i < 1000; i++ {
+		buf[i] = chunk
+	}
+	buf[1001] = []byte("0\r\n\r\n")
+	packets := make([]gopacket.Packet, len(buf))
+	for i := 0; i < len(buf); i++ {
+		packets[i] = GetPackets(uint32(i+10), 1, buf[i])[0]
 	}
-	now := time.Now()
-	n := b.N
 	var mssg = make(chan *Message, 1)
-	payload := make([]byte, 0xfc00)
-	for i := 0; i < 0xfc00; i++ {
-		payload[i] = '1'
-	}
 	pool := NewMessagePool(1<<30, time.Second*10, nil, func(m *Message) { mssg <- m })
 	pool.Start = func(pckt *Packet) (bool, bool) {
-		return proto.HasRequestTitle(pckt.Payload), proto.HasResponseTitle(pckt.Payload)
+		return false, proto.HasResponseTitle(pckt.Payload)
 	}
 	pool.End = func(m *Message) bool {
 		return proto.HasFullPayload(m.Data())
 	}
-	pool.Handler(GetPackets(1, 1, []byte("POST / HTTP/1.1\r\nContent-Type: text/plain\r\nTransfer-Encoding: chunked\r\n\r\n"))[0])
-	i := 0
-	var d []byte
-	for {
-		select {
-		case m := <-mssg:
-			b.Logf("%d/%d packets, %dbytes, truncated: %v, timedout: %v in %s", len(m.packets), n, m.Length, m.Truncated, m.TimedOut, time.Since(now))
-			return
-		default:
-			if i > n-2 {
-				break
-			} else if i < n-2 {
-				d = []byte(fmt.Sprintf("fc00\r\n%s\r\n", payload))
-			} else {
-				d = []byte("0\r\n\r\n")
-			}
-			pool.Handler(GetPackets(1, i+2, d)[0])
-			i++
+	b.ResetTimer()
+	b.ReportMetric(float64(len(packets)), "packets/op")
+	b.ReportMetric(float64(1000), "chunks/op")
+	for i := 0; i < b.N; i++ {
+		for j := range packets {
+			pool.Handler(packets[j])
 		}
+		<-mssg
 	}
 }