tailscale/net/tunstats/stats.go
Joe Tsai 2934c5114c
net/tunstats: new package to track per-connection counters (#5818)
High-level API:

	type Statistics struct { ... }
	type Counts struct { TxPackets, TxBytes, RxPackets, RxBytes uint64 }
	func (*Statistics) UpdateTx([]byte)
	func (*Statistics) UpdateRx([]byte)
	func (*Statistics) Extract() map[flowtrack.Tuple]Counts

The API accepts a []byte instead of a packet.Parsed so that a future
implementation can directly hash the address and port bytes,
which are contiguous in most IP packets.
This will be useful for a custom concurrent-safe hashmap implementation.

Signed-off-by: Joe Tsai <joetsai@digital-static.net>
2022-10-04 15:10:33 -07:00

79 lines
2.1 KiB
Go

// Copyright (c) 2022 Tailscale Inc & AUTHORS All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package tunstats maintains statistics about connections
// flowing through a TUN device (which operate at the IP layer).
package tunstats
import (
"sync"
"tailscale.com/net/flowtrack"
"tailscale.com/net/packet"
)
// Statistics maintains counters for every connection.
// All methods are safe for concurrent use.
// The zero value is ready for use.
type Statistics struct {
mu sync.Mutex
m map[flowtrack.Tuple]Counts
}
// Counts are statistics about a particular connection.
type Counts struct {
TxPackets uint64 `json:"txPkts,omitempty"`
TxBytes uint64 `json:"txBytes,omitempty"`
RxPackets uint64 `json:"rxPkts,omitempty"`
RxBytes uint64 `json:"rxBytes,omitempty"`
}
// UpdateTx updates the counters for a transmitted IP packet
// The source and destination of the packet directly correspond with
// the source and destination in flowtrack.Tuple.
func (s *Statistics) UpdateTx(b []byte) {
s.update(b, false)
}
// UpdateRx updates the counters for a received IP packet.
// The source and destination of the packet are inverted with respect to
// the source and destination in flowtrack.Tuple.
func (s *Statistics) UpdateRx(b []byte) {
s.update(b, true)
}
func (s *Statistics) update(b []byte, receive bool) {
var p packet.Parsed
p.Decode(b)
tuple := flowtrack.Tuple{Proto: p.IPProto, Src: p.Src, Dst: p.Dst}
if receive {
tuple.Src, tuple.Dst = tuple.Dst, tuple.Src
}
s.mu.Lock()
defer s.mu.Unlock()
if s.m == nil {
s.m = make(map[flowtrack.Tuple]Counts)
}
cnts := s.m[tuple]
if receive {
cnts.RxPackets++
cnts.RxBytes += uint64(len(b))
} else {
cnts.TxPackets++
cnts.TxBytes += uint64(len(b))
}
s.m[tuple] = cnts
}
// Extract extracts and resets the counters for all active connections.
// It must be called periodically otherwise the memory used is unbounded.
func (s *Statistics) Extract() map[flowtrack.Tuple]Counts {
s.mu.Lock()
defer s.mu.Unlock()
m := s.m
s.m = make(map[flowtrack.Tuple]Counts)
return m
}