From 9519b9a9f51af01697f6538eb97183ecb71bfa45 Mon Sep 17 00:00:00 2001 From: Urban Ishimwe Date: Tue, 22 Sep 2020 20:14:31 +0200 Subject: [PATCH] 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 | 8 +--- byteutils/byteutils_test.go | 11 ++--- input_raw.go | 36 +++++++++------- output_dummy.go | 8 ++-- proto/proto.go | 6 +-- proto/proto_test.go | 33 ++++++--------- tcp/tcp_message.go | 51 +++++++++++++---------- tcp/tcp_packet.go | 64 ++++++++++------------------ tcp/tcp_test.go | 83 ++++++++++++++++--------------------- 9 files changed, 131 insertions(+), 169 deletions(-) diff --git a/byteutils/byteutils.go b/byteutils/byteutils.go index 81a7933..bee9cd1 100644 --- a/byteutils/byteutils.go +++ b/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)) } diff --git a/byteutils/byteutils_test.go b/byteutils/byteutils_test.go index 5996db3..236ba0f 100644 --- a/byteutils/byteutils_test.go +++ b/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 } diff --git a/input_raw.go b/input_raw.go index 6a1dd51..ab3e7e4 100644 --- a/input_raw.go +++ b/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 { diff --git a/output_dummy.go b/output_dummy.go index c9e51d0..eca0f53 100644 --- a/output_dummy.go +++ b/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 { diff --git a/proto/proto.go b/proto/proto.go index b54be29..4b905db 100644 --- a/proto/proto.go +++ b/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 } diff --git a/proto/proto_test.go b/proto/proto_test.go index 4baba18..9f7fd26 100644 --- a/proto/proto_test.go +++ b/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)) } diff --git a/tcp/tcp_message.go b/tcp/tcp_message.go index c83b765..b7db474 100644 --- a/tcp/tcp_message.go +++ b/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) diff --git a/tcp/tcp_packet.go b/tcp/tcp_packet.go index a1a4968..445fe58 100644 --- a/tcp/tcp_packet.go +++ b/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, diff --git a/tcp/tcp_test.go b/tcp/tcp_test.go index b38e4b7..f2cb16c 100644 --- a/tcp/tcp_test.go +++ b/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 } }