mirror of
				https://source.quilibrium.com/quilibrium/ceremonyclient.git
				synced 2025-11-04 07:07:26 +00:00 
			
		
		
		
	
		
			
				
	
	
		
			269 lines
		
	
	
		
			6.5 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
			
		
		
	
	
			269 lines
		
	
	
		
			6.5 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
//go:build !windows
 | 
						|
 | 
						|
package tcp
 | 
						|
 | 
						|
import (
 | 
						|
	"strings"
 | 
						|
	"sync"
 | 
						|
	"time"
 | 
						|
 | 
						|
	"github.com/marten-seemann/tcp"
 | 
						|
	"github.com/mikioh/tcpinfo"
 | 
						|
	manet "github.com/multiformats/go-multiaddr/net"
 | 
						|
	"github.com/prometheus/client_golang/prometheus"
 | 
						|
)
 | 
						|
 | 
						|
var (
 | 
						|
	newConns      *prometheus.CounterVec
 | 
						|
	closedConns   *prometheus.CounterVec
 | 
						|
	segsSentDesc  *prometheus.Desc
 | 
						|
	segsRcvdDesc  *prometheus.Desc
 | 
						|
	bytesSentDesc *prometheus.Desc
 | 
						|
	bytesRcvdDesc *prometheus.Desc
 | 
						|
)
 | 
						|
 | 
						|
const collectFrequency = 10 * time.Second
 | 
						|
 | 
						|
var collector *aggregatingCollector
 | 
						|
 | 
						|
var initMetricsOnce sync.Once
 | 
						|
 | 
						|
func initMetrics() {
 | 
						|
	segsSentDesc = prometheus.NewDesc("tcp_sent_segments_total", "TCP segments sent", nil, nil)
 | 
						|
	segsRcvdDesc = prometheus.NewDesc("tcp_rcvd_segments_total", "TCP segments received", nil, nil)
 | 
						|
	bytesSentDesc = prometheus.NewDesc("tcp_sent_bytes", "TCP bytes sent", nil, nil)
 | 
						|
	bytesRcvdDesc = prometheus.NewDesc("tcp_rcvd_bytes", "TCP bytes received", nil, nil)
 | 
						|
 | 
						|
	collector = newAggregatingCollector()
 | 
						|
	prometheus.MustRegister(collector)
 | 
						|
 | 
						|
	const direction = "direction"
 | 
						|
 | 
						|
	newConns = prometheus.NewCounterVec(
 | 
						|
		prometheus.CounterOpts{
 | 
						|
			Name: "tcp_connections_new_total",
 | 
						|
			Help: "TCP new connections",
 | 
						|
		},
 | 
						|
		[]string{direction},
 | 
						|
	)
 | 
						|
	prometheus.MustRegister(newConns)
 | 
						|
	closedConns = prometheus.NewCounterVec(
 | 
						|
		prometheus.CounterOpts{
 | 
						|
			Name: "tcp_connections_closed_total",
 | 
						|
			Help: "TCP connections closed",
 | 
						|
		},
 | 
						|
		[]string{direction},
 | 
						|
	)
 | 
						|
	prometheus.MustRegister(closedConns)
 | 
						|
}
 | 
						|
 | 
						|
type aggregatingCollector struct {
 | 
						|
	cronOnce sync.Once
 | 
						|
 | 
						|
	mutex                sync.Mutex
 | 
						|
	highestID            uint64
 | 
						|
	conns                map[uint64] /* id */ *tracingConn
 | 
						|
	rtts                 prometheus.Histogram
 | 
						|
	connDurations        prometheus.Histogram
 | 
						|
	segsSent, segsRcvd   uint64
 | 
						|
	bytesSent, bytesRcvd uint64
 | 
						|
}
 | 
						|
 | 
						|
var _ prometheus.Collector = &aggregatingCollector{}
 | 
						|
 | 
						|
func newAggregatingCollector() *aggregatingCollector {
 | 
						|
	c := &aggregatingCollector{
 | 
						|
		conns: make(map[uint64]*tracingConn),
 | 
						|
		rtts: prometheus.NewHistogram(prometheus.HistogramOpts{
 | 
						|
			Name:    "tcp_rtt",
 | 
						|
			Help:    "TCP round trip time",
 | 
						|
			Buckets: prometheus.ExponentialBuckets(0.001, 1.25, 40), // 1ms to ~6000ms
 | 
						|
		}),
 | 
						|
		connDurations: prometheus.NewHistogram(prometheus.HistogramOpts{
 | 
						|
			Name:    "tcp_connection_duration",
 | 
						|
			Help:    "TCP Connection Duration",
 | 
						|
			Buckets: prometheus.ExponentialBuckets(1, 1.5, 40), // 1s to ~12 weeks
 | 
						|
		}),
 | 
						|
	}
 | 
						|
	return c
 | 
						|
}
 | 
						|
 | 
						|
func (c *aggregatingCollector) AddConn(t *tracingConn) uint64 {
 | 
						|
	c.mutex.Lock()
 | 
						|
	defer c.mutex.Unlock()
 | 
						|
	c.highestID++
 | 
						|
	c.conns[c.highestID] = t
 | 
						|
	return c.highestID
 | 
						|
}
 | 
						|
 | 
						|
func (c *aggregatingCollector) removeConn(id uint64) {
 | 
						|
	delete(c.conns, id)
 | 
						|
}
 | 
						|
 | 
						|
func (c *aggregatingCollector) Describe(descs chan<- *prometheus.Desc) {
 | 
						|
	descs <- c.rtts.Desc()
 | 
						|
	descs <- c.connDurations.Desc()
 | 
						|
	if hasSegmentCounter {
 | 
						|
		descs <- segsSentDesc
 | 
						|
		descs <- segsRcvdDesc
 | 
						|
	}
 | 
						|
	if hasByteCounter {
 | 
						|
		descs <- bytesSentDesc
 | 
						|
		descs <- bytesRcvdDesc
 | 
						|
	}
 | 
						|
}
 | 
						|
 | 
						|
func (c *aggregatingCollector) cron() {
 | 
						|
	ticker := time.NewTicker(collectFrequency)
 | 
						|
	defer ticker.Stop()
 | 
						|
 | 
						|
	for now := range ticker.C {
 | 
						|
		c.gatherMetrics(now)
 | 
						|
	}
 | 
						|
}
 | 
						|
 | 
						|
func (c *aggregatingCollector) gatherMetrics(now time.Time) {
 | 
						|
	c.mutex.Lock()
 | 
						|
	defer c.mutex.Unlock()
 | 
						|
 | 
						|
	c.segsSent = 0
 | 
						|
	c.segsRcvd = 0
 | 
						|
	c.bytesSent = 0
 | 
						|
	c.bytesRcvd = 0
 | 
						|
	for _, conn := range c.conns {
 | 
						|
		info, err := conn.getTCPInfo()
 | 
						|
		if err != nil {
 | 
						|
			if strings.Contains(err.Error(), "use of closed network connection") {
 | 
						|
				continue
 | 
						|
			}
 | 
						|
			log.Errorf("Failed to get TCP info: %s", err)
 | 
						|
			continue
 | 
						|
		}
 | 
						|
		if hasSegmentCounter {
 | 
						|
			c.segsSent += getSegmentsSent(info)
 | 
						|
			c.segsRcvd += getSegmentsRcvd(info)
 | 
						|
		}
 | 
						|
		if hasByteCounter {
 | 
						|
			c.bytesSent += getBytesSent(info)
 | 
						|
			c.bytesRcvd += getBytesRcvd(info)
 | 
						|
		}
 | 
						|
		c.rtts.Observe(info.RTT.Seconds())
 | 
						|
		c.connDurations.Observe(now.Sub(conn.startTime).Seconds())
 | 
						|
	}
 | 
						|
}
 | 
						|
 | 
						|
func (c *aggregatingCollector) Collect(metrics chan<- prometheus.Metric) {
 | 
						|
	// Start collecting the metrics collection the first time Collect is called.
 | 
						|
	c.cronOnce.Do(func() {
 | 
						|
		c.gatherMetrics(time.Now())
 | 
						|
		go c.cron()
 | 
						|
	})
 | 
						|
 | 
						|
	c.mutex.Lock()
 | 
						|
	defer c.mutex.Unlock()
 | 
						|
 | 
						|
	metrics <- c.rtts
 | 
						|
	metrics <- c.connDurations
 | 
						|
	if hasSegmentCounter {
 | 
						|
		segsSentMetric, err := prometheus.NewConstMetric(segsSentDesc, prometheus.CounterValue, float64(c.segsSent))
 | 
						|
		if err != nil {
 | 
						|
			log.Errorf("creating tcp_sent_segments_total metric failed: %v", err)
 | 
						|
			return
 | 
						|
		}
 | 
						|
		segsRcvdMetric, err := prometheus.NewConstMetric(segsRcvdDesc, prometheus.CounterValue, float64(c.segsRcvd))
 | 
						|
		if err != nil {
 | 
						|
			log.Errorf("creating tcp_rcvd_segments_total metric failed: %v", err)
 | 
						|
			return
 | 
						|
		}
 | 
						|
		metrics <- segsSentMetric
 | 
						|
		metrics <- segsRcvdMetric
 | 
						|
	}
 | 
						|
	if hasByteCounter {
 | 
						|
		bytesSentMetric, err := prometheus.NewConstMetric(bytesSentDesc, prometheus.CounterValue, float64(c.bytesSent))
 | 
						|
		if err != nil {
 | 
						|
			log.Errorf("creating tcp_sent_bytes metric failed: %v", err)
 | 
						|
			return
 | 
						|
		}
 | 
						|
		bytesRcvdMetric, err := prometheus.NewConstMetric(bytesRcvdDesc, prometheus.CounterValue, float64(c.bytesRcvd))
 | 
						|
		if err != nil {
 | 
						|
			log.Errorf("creating tcp_rcvd_bytes metric failed: %v", err)
 | 
						|
			return
 | 
						|
		}
 | 
						|
		metrics <- bytesSentMetric
 | 
						|
		metrics <- bytesRcvdMetric
 | 
						|
	}
 | 
						|
}
 | 
						|
 | 
						|
func (c *aggregatingCollector) ClosedConn(conn *tracingConn, direction string) {
 | 
						|
	c.mutex.Lock()
 | 
						|
	collector.removeConn(conn.id)
 | 
						|
	c.mutex.Unlock()
 | 
						|
	closedConns.WithLabelValues(direction).Inc()
 | 
						|
}
 | 
						|
 | 
						|
type tracingConn struct {
 | 
						|
	id uint64
 | 
						|
 | 
						|
	startTime time.Time
 | 
						|
	isClient  bool
 | 
						|
 | 
						|
	manet.Conn
 | 
						|
	tcpConn *tcp.Conn
 | 
						|
}
 | 
						|
 | 
						|
func newTracingConn(c manet.Conn, isClient bool) (*tracingConn, error) {
 | 
						|
	initMetricsOnce.Do(func() { initMetrics() })
 | 
						|
	conn, err := tcp.NewConn(c)
 | 
						|
	if err != nil {
 | 
						|
		return nil, err
 | 
						|
	}
 | 
						|
	tc := &tracingConn{
 | 
						|
		startTime: time.Now(),
 | 
						|
		isClient:  isClient,
 | 
						|
		Conn:      c,
 | 
						|
		tcpConn:   conn,
 | 
						|
	}
 | 
						|
	tc.id = collector.AddConn(tc)
 | 
						|
	newConns.WithLabelValues(tc.getDirection()).Inc()
 | 
						|
	return tc, nil
 | 
						|
}
 | 
						|
 | 
						|
func (c *tracingConn) getDirection() string {
 | 
						|
	if c.isClient {
 | 
						|
		return "outgoing"
 | 
						|
	}
 | 
						|
	return "incoming"
 | 
						|
}
 | 
						|
 | 
						|
func (c *tracingConn) Close() error {
 | 
						|
	collector.ClosedConn(c, c.getDirection())
 | 
						|
	return c.Conn.Close()
 | 
						|
}
 | 
						|
 | 
						|
func (c *tracingConn) getTCPInfo() (*tcpinfo.Info, error) {
 | 
						|
	var o tcpinfo.Info
 | 
						|
	var b [256]byte
 | 
						|
	i, err := c.tcpConn.Option(o.Level(), o.Name(), b[:])
 | 
						|
	if err != nil {
 | 
						|
		return nil, err
 | 
						|
	}
 | 
						|
	info := i.(*tcpinfo.Info)
 | 
						|
	return info, nil
 | 
						|
}
 | 
						|
 | 
						|
type tracingListener struct {
 | 
						|
	manet.Listener
 | 
						|
}
 | 
						|
 | 
						|
func newTracingListener(l manet.Listener) *tracingListener {
 | 
						|
	return &tracingListener{Listener: l}
 | 
						|
}
 | 
						|
 | 
						|
func (l *tracingListener) Accept() (manet.Conn, error) {
 | 
						|
	conn, err := l.Listener.Accept()
 | 
						|
	if err != nil {
 | 
						|
		return nil, err
 | 
						|
	}
 | 
						|
	return newTracingConn(conn, false)
 | 
						|
}
 |