sing-box/experimental/v2rayapi/stats.go
2022-09-29 08:49:34 +08:00

194 lines
5.4 KiB
Go

package v2rayapi
import (
"context"
"net"
"regexp"
"runtime"
"strings"
"sync"
"time"
"github.com/sagernet/sing-box/adapter"
"github.com/sagernet/sing-box/experimental/trackerconn"
"github.com/sagernet/sing-box/option"
E "github.com/sagernet/sing/common/exceptions"
N "github.com/sagernet/sing/common/network"
"go.uber.org/atomic"
)
func init() {
StatsService_ServiceDesc.ServiceName = "v2ray.core.app.stats.command.StatsService"
}
var (
_ adapter.V2RayStatsService = (*StatsService)(nil)
_ StatsServiceServer = (*StatsService)(nil)
)
type StatsService struct {
createdAt time.Time
directIO bool
inbounds map[string]bool
outbounds map[string]bool
access sync.Mutex
counters map[string]*atomic.Int64
}
func NewStatsService(options option.V2RayStatsServiceOptions) *StatsService {
if !options.Enabled {
return nil
}
inbounds := make(map[string]bool)
outbounds := make(map[string]bool)
for _, inbound := range options.Inbounds {
inbounds[inbound] = true
}
for _, outbound := range options.Outbounds {
outbounds[outbound] = true
}
return &StatsService{
createdAt: time.Now(),
directIO: options.DirectIO,
inbounds: inbounds,
outbounds: outbounds,
counters: make(map[string]*atomic.Int64),
}
}
func (s *StatsService) RoutedConnection(inbound string, outbound string, conn net.Conn) net.Conn {
var readCounter *atomic.Int64
var writeCounter *atomic.Int64
countInbound := inbound != "" && s.inbounds[inbound]
countOutbound := outbound != "" && s.outbounds[outbound]
if !countInbound && !countOutbound {
return conn
}
s.access.Lock()
if countInbound {
readCounter = s.loadOrCreateCounter("inbound>>>"+inbound+">>>traffic>>>uplink", readCounter)
writeCounter = s.loadOrCreateCounter("inbound>>>"+inbound+">>>traffic>>>downlink", writeCounter)
}
if countOutbound {
readCounter = s.loadOrCreateCounter("outbound>>>"+outbound+">>>traffic>>>uplink", readCounter)
writeCounter = s.loadOrCreateCounter("outbound>>>"+outbound+">>>traffic>>>downlink", writeCounter)
}
s.access.Unlock()
return trackerconn.New(conn, readCounter, writeCounter, s.directIO)
}
func (s *StatsService) RoutedPacketConnection(inbound string, outbound string, conn N.PacketConn) N.PacketConn {
var readCounter *atomic.Int64
var writeCounter *atomic.Int64
countInbound := inbound != "" && s.inbounds[inbound]
countOutbound := outbound != "" && s.outbounds[outbound]
if !countInbound && !countOutbound {
return conn
}
s.access.Lock()
if countInbound {
readCounter = s.loadOrCreateCounter("inbound>>>"+inbound+">>>traffic>>>uplink", readCounter)
writeCounter = s.loadOrCreateCounter("inbound>>>"+inbound+">>>traffic>>>downlink", writeCounter)
}
if countOutbound {
readCounter = s.loadOrCreateCounter("outbound>>>"+outbound+">>>traffic>>>uplink", readCounter)
writeCounter = s.loadOrCreateCounter("outbound>>>"+outbound+">>>traffic>>>downlink", writeCounter)
}
s.access.Unlock()
return trackerconn.NewPacket(conn, readCounter, writeCounter)
}
func (s *StatsService) GetStats(ctx context.Context, request *GetStatsRequest) (*GetStatsResponse, error) {
s.access.Lock()
counter, loaded := s.counters[request.Name]
s.access.Unlock()
if !loaded {
return nil, E.New(request.Name, " not found.")
}
var value int64
if request.Reset_ {
value = counter.Swap(0)
} else {
value = counter.Load()
}
return &GetStatsResponse{Stat: &Stat{Name: request.Name, Value: value}}, nil
}
func (s *StatsService) QueryStats(ctx context.Context, request *QueryStatsRequest) (*QueryStatsResponse, error) {
var response QueryStatsResponse
s.access.Lock()
defer s.access.Unlock()
if request.Regexp {
matchers := make([]*regexp.Regexp, 0, len(request.Patterns))
for _, pattern := range request.Patterns {
matcher, err := regexp.Compile(pattern)
if err != nil {
return nil, err
}
matchers = append(matchers, matcher)
}
for name, counter := range s.counters {
for _, matcher := range matchers {
if matcher.MatchString(name) {
var value int64
if request.Reset_ {
value = counter.Swap(0)
} else {
value = counter.Load()
}
response.Stat = append(response.Stat, &Stat{Name: name, Value: value})
}
}
}
} else {
for name, counter := range s.counters {
for _, matcher := range request.Patterns {
if strings.Contains(name, matcher) {
var value int64
if request.Reset_ {
value = counter.Swap(0)
} else {
value = counter.Load()
}
response.Stat = append(response.Stat, &Stat{Name: name, Value: value})
}
}
}
}
return &response, nil
}
func (s *StatsService) GetSysStats(ctx context.Context, request *SysStatsRequest) (*SysStatsResponse, error) {
var rtm runtime.MemStats
runtime.ReadMemStats(&rtm)
response := &SysStatsResponse{
Uptime: uint32(time.Now().Sub(s.createdAt).Seconds()),
NumGoroutine: uint32(runtime.NumGoroutine()),
Alloc: rtm.Alloc,
TotalAlloc: rtm.TotalAlloc,
Sys: rtm.Sys,
Mallocs: rtm.Mallocs,
Frees: rtm.Frees,
LiveObjects: rtm.Mallocs - rtm.Frees,
NumGC: rtm.NumGC,
PauseTotalNs: rtm.PauseTotalNs,
}
return response, nil
}
func (s *StatsService) mustEmbedUnimplementedStatsServiceServer() {
}
func (s *StatsService) loadOrCreateCounter(name string, counter *atomic.Int64) *atomic.Int64 {
counter, loaded := s.counters[name]
if !loaded {
if counter == nil {
counter = atomic.NewInt64(0)
}
s.counters[name] = counter
}
return counter
}