optimization: optimized locking and op atomicity for the packets queue and network statistics view

modules/net_recon_show.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 	"sort"
+	"sync/atomic"
 	"time"
 
 	"github.com/evilsocket/bettercap-ng/core"
@@ -112,9 +113,6 @@ func (d *Discovery) showTable(header []string, rows [][]string) {
 }
 
 func (d *Discovery) Show(by string) error {
-	d.Session.Queue.Lock()
-	defer d.Session.Queue.Unlock()
-
 	targets := d.Session.Targets.List()
 	if by == "seen" {
 		sort.Sort(BySeenSorter(targets))
@@ -146,11 +144,11 @@ func (d *Discovery) Show(by string) error {
 
 	fmt.Printf("\n%s %s / %s %s / %d pkts / %d errs\n\n",
 		core.Red("↑"),
-		humanize.Bytes(d.Session.Queue.Sent),
+		humanize.Bytes(atomic.LoadUint64(&d.Session.Queue.Sent)),
 		core.Green("↓"),
-		humanize.Bytes(d.Session.Queue.Received),
+		humanize.Bytes(atomic.LoadUint64(&d.Session.Queue.Received)),
-		d.Session.Queue.PktReceived,
+		atomic.LoadUint64(&d.Session.Queue.PktReceived),
-		d.Session.Queue.Errors)
+		atomic.LoadUint64(&d.Session.Queue.Errors))
 
 	s := EventsStream{}
 	events := d.Session.Events.Sorted()

modules/net_recon_sort.go

@@ -28,6 +28,9 @@ type BySentSorter []*net.Endpoint
 func (a BySentSorter) Len() int      { return len(a) }
 func (a BySentSorter) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
 func (a BySentSorter) Less(i, j int) bool {
+	session.I.Queue.Lock()
+	defer session.I.Queue.Unlock()
+
 	var found bool = false
 	var aTraffic *packets.Traffic = nil
 	var bTraffic *packets.Traffic = nil
@@ -48,6 +51,9 @@ type ByRcvdSorter []*net.Endpoint
 func (a ByRcvdSorter) Len() int      { return len(a) }
 func (a ByRcvdSorter) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
 func (a ByRcvdSorter) Less(i, j int) bool {
+	session.I.Queue.Lock()
+	defer session.I.Queue.Unlock()
+
 	var found bool = false
 	var aTraffic *packets.Traffic = nil
 	var bTraffic *packets.Traffic = nil

packets/queue.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net"
 	"sync"
+	"sync/atomic"
 
 	bnet "github.com/evilsocket/bettercap-ng/net"
 
@@ -75,11 +76,10 @@ func (q *Queue) worker() {
 			return
 		}
 
-		q.Lock()
 		pktSize := uint64(len(pkt.Data()))
 
-		q.PktReceived++
+		atomic.AddUint64(&q.PktReceived, 1)
-		q.Received += pktSize
+		atomic.AddUint64(&q.Received, pktSize)
 
 		// gather protocols stats
 		pktLayers := pkt.Layers()
@@ -89,10 +89,13 @@ func (q *Queue) worker() {
 				continue
 			}
 
+			q.Lock()
 			if _, found := q.Protos[proto]; found == false {
-				q.Protos[proto] = 0
+				q.Protos[proto] = 1
+			} else {
+				q.Protos[proto] += 1
 			}
-			q.Protos[proto] += 1
+			q.Unlock()
 		}
 
 		// check for new ipv4 endpoints
@@ -104,6 +107,7 @@ func (q *Queue) worker() {
 			ip4 := lip4.(*layers.IPv4)
 
 			if bytes.Compare(q.iface.IP, ip4.SrcIP) != 0 && q.iface.Net.Contains(ip4.SrcIP) {
+				q.Lock()
 				q.Activities <- Activity{
 					IP:     ip4.SrcIP,
 					MAC:    eth.SrcMAC,
@@ -112,13 +116,17 @@ func (q *Queue) worker() {
 
 				addr := ip4.SrcIP.String()
 				if _, found := q.Traffic[addr]; found == false {
-					q.Traffic[addr] = &Traffic{}
+					q.Traffic[addr] = &Traffic{
+						Sent: pktSize,
+					}
+				} else {
+					q.Traffic[addr].Sent += pktSize
 				}
-
+				q.Unlock()
-				q.Traffic[addr].Sent += pktSize
 			}
 
 			if bytes.Compare(q.iface.IP, ip4.DstIP) != 0 && q.iface.Net.Contains(ip4.DstIP) {
+				q.Lock()
 				q.Activities <- Activity{
 					IP:     ip4.DstIP,
 					MAC:    eth.SrcMAC,
@@ -127,14 +135,15 @@ func (q *Queue) worker() {
 
 				addr := ip4.DstIP.String()
 				if _, found := q.Traffic[addr]; found == false {
-					q.Traffic[addr] = &Traffic{}
+					q.Traffic[addr] = &Traffic{
+						Received: pktSize,
+					}
+				} else {
+					q.Traffic[addr].Received += pktSize
 				}
-
+				q.Unlock()
-				q.Traffic[addr].Received += pktSize
 			}
 		}
-
-		q.Unlock()
 	}
 }