Rework stats server

main.go

@@ -193,7 +193,7 @@ func main() { // nolint: gocyclo
 		zap.S().Infow("Use direct connection to Telegram")
 	}
 
-	if err := stats.Start(conf); err != nil {
+	if err := stats.Init(conf); err != nil {
 		panic(err)
 	}
 

stats/channels.go

@@ -2,21 +2,20 @@ package stats
 
 import (
 	"net"
-	"time"
 
 	"github.com/9seconds/mtg/mtproto"
 )
 
 const (
-	crashesChanLength     = 1
-	connectionsChanLength = 20
-	trafficChanLength     = 5000
+	connectionsChanLength = 10
+	trafficChanLength     = 10
 )
 
 var (
-	crashesChan     = make(chan struct{}, crashesChanLength)
-	connectionsChan = make(chan *connectionData, connectionsChanLength)
-	trafficChan     = make(chan *trafficData, trafficChanLength)
+	crashesChan     = make(chan struct{})
+	statsChan       = make(chan chan<- Stats)
+	connectionsChan = make(chan connectionData, connectionsChanLength)
+	trafficChan     = make(chan trafficData, trafficChanLength)
 )
 
 type connectionData struct {
@@ -30,81 +29,6 @@ type trafficData struct {
 	ingress bool
 }
 
-func crashManager() {
-	for range crashesChan {
-		instance.mutex.RLock()
-
-		instance.Crashes++
-
-		instance.mutex.RUnlock()
-	}
-}
-
-func connectionManager() {
-	for event := range connectionsChan {
-		instance.mutex.RLock()
-
-		isIPv4 := event.addr.IP.To4() != nil
-		var inc uint32 = 1
-		if !event.connected {
-			inc = ^uint32(0)
-		}
-
-		switch event.connectionType {
-		case mtproto.ConnectionTypeAbridged:
-			if isIPv4 {
-				instance.Connections.Abridged.IPv4 += inc
-			} else {
-				instance.Connections.Abridged.IPv6 += inc
-			}
-		case mtproto.ConnectionTypeSecure:
-			if isIPv4 {
-				instance.Connections.Secure.IPv4 += inc
-			} else {
-				instance.Connections.Secure.IPv6 += inc
-			}
-		default:
-			if isIPv4 {
-				instance.Connections.Intermediate.IPv4 += inc
-			} else {
-				instance.Connections.Intermediate.IPv6 += inc
-			}
-		}
-
-		instance.mutex.RUnlock()
-	}
-}
-
-func trafficManager() {
-	speedChan := time.Tick(time.Second)
-
-	for {
-		select {
-		case event := <-trafficChan:
-			instance.mutex.RLock()
-
-			if event.ingress {
-				instance.Traffic.Ingress += trafficValue(event.traffic)
-				instance.speedCurrent.Ingress += trafficSpeedValue(event.traffic)
-			} else {
-				instance.Traffic.Egress += trafficValue(event.traffic)
-				instance.speedCurrent.Egress += trafficSpeedValue(event.traffic)
-			}
-
-			instance.mutex.RUnlock()
-		case <-speedChan:
-			instance.mutex.RLock()
-
-			instance.Speed.Ingress = instance.speedCurrent.Ingress
-			instance.Speed.Egress = instance.speedCurrent.Egress
-			instance.speedCurrent.Ingress = trafficSpeedValue(0)
-			instance.speedCurrent.Egress = trafficSpeedValue(0)
-
-			instance.mutex.RUnlock()
-		}
-	}
-}
-
 // NewCrash indicates new crash.
 func NewCrash() {
 	crashesChan <- struct{}{}
@@ -112,7 +36,7 @@ func NewCrash() {
 
 // ClientConnected indicates that new client was connected.
 func ClientConnected(connectionType mtproto.ConnectionType, addr *net.TCPAddr) {
-	connectionsChan <- &connectionData{
+	connectionsChan <- connectionData{
 		connectionType: connectionType,
 		addr:           addr,
 		connected:      true,
@@ -121,7 +45,7 @@ func ClientConnected(connectionType mtproto.ConnectionType, addr *net.TCPAddr) {
 
 // ClientDisconnected indicates that client was disconnected.
 func ClientDisconnected(connectionType mtproto.ConnectionType, addr *net.TCPAddr) {
-	connectionsChan <- &connectionData{
+	connectionsChan <- connectionData{
 		connectionType: connectionType,
 		addr:           addr,
 		connected:      false,
@@ -130,7 +54,7 @@ func ClientDisconnected(connectionType mtproto.ConnectionType, addr *net.TCPAddr
 
 // IngressTraffic accounts new ingress traffic.
 func IngressTraffic(traffic int) {
-	trafficChan <- &trafficData{
+	trafficChan <- trafficData{
 		traffic: traffic,
 		ingress: true,
 	}
@@ -138,8 +62,14 @@ func IngressTraffic(traffic int) {
 
 // EgressTraffic accounts new ingress traffic.
 func EgressTraffic(traffic int) {
-	trafficChan <- &trafficData{
+	trafficChan <- trafficData{
 		traffic: traffic,
 		ingress: false,
 	}
 }
+
+func GetStats() Stats {
+	rpcChan := make(chan Stats)
+	statsChan <- rpcChan
+	return <-rpcChan
+}

stats/init.go

@@ -0,0 +1,22 @@
+package stats
+
+import (
+	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/config"
+)
+
+func Init(conf *config.Config) error {
+	if conf.StatsD.Enabled {
+		client, err := newStatsd(conf)
+		if err != nil {
+			return errors.Annotate(err, "Cannot initialize statsd client")
+		}
+		go client.run()
+	}
+
+	go NewStats(conf).start()
+	go startServer(conf)
+
+	return nil
+}

stats/server.go

@@ -3,45 +3,19 @@ package stats
 import (
 	"encoding/json"
 	"net/http"
-	"sync"
-	"time"
 
 	"go.uber.org/zap"
 
 	"github.com/9seconds/mtg/config"
 )
 
-var instance *stats
-
-// Start starts new statistics server.
-func Start(conf *config.Config) error {
+func startServer(conf *config.Config) {
 	log := zap.S().Named("stats")
 
-	instance = &stats{
-		URLs:   conf.GetURLs(),
-		Uptime: uptime(time.Now()),
-		mutex:  &sync.RWMutex{},
-	}
-
-	if conf.StatsD.Enabled {
-		client, err := newStatsd(conf)
-		if err != nil {
-			return err
-		}
-		go client.run()
-	}
-
-	go crashManager()
-	go connectionManager()
-	go trafficManager()
-
 	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Content-Type", "application/json")
 
-		instance.mutex.Lock()
-		first, err := json.Marshal(instance)
-		instance.mutex.Unlock()
-
+		first, err := json.Marshal(GetStats())
 		if err != nil {
 			log.Errorw("Cannot encode json", "error", err)
 			http.Error(w, "Internal server error", 500)
@@ -59,11 +33,7 @@ func Start(conf *config.Config) error {
 		}
 	})
 
-	go func() {
-		if err := http.ListenAndServe(conf.StatAddr(), nil); err != nil {
-			log.Fatalw("Stats server has been stopped", "error", err)
-		}
-	}()
-
-	return nil
+	if err := http.ListenAndServe(conf.StatAddr(), nil); err != nil {
+		log.Fatalw("Stats server has been stopped", "error", err)
+	}
 }

stats/stats.go

@@ -4,12 +4,12 @@ import (
 	"encoding/json"
 	"fmt"
 	"strconv"
-	"sync"
 	"time"
 
 	humanize "github.com/dustin/go-humanize"
 
 	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/mtproto"
 )
 
 type uptime time.Time
@@ -24,72 +24,72 @@ func (u uptime) MarshalJSON() ([]byte, error) {
 	return json.Marshal(value)
 }
 
-type trafficValue uint64
-
-func (t trafficValue) MarshalJSON() ([]byte, error) {
-	tv := uint64(t)
-	value := map[string]interface{}{
-		"bytes": tv,
-		"human": humanize.Bytes(tv),
-	}
-
-	return json.Marshal(value)
-}
-
-type trafficSpeedValue uint64
-
-func (t trafficSpeedValue) MarshalJSON() ([]byte, error) {
-	speed := uint64(t)
-	value := map[string]interface{}{
-		"bytes/s": speed,
-		"human":   fmt.Sprintf("%s/S", humanize.Bytes(speed)),
-	}
-
-	return json.Marshal(value)
+type connectionType struct {
+	IPv6 uint32 `json:"ipv6"`
+	IPv4 uint32 `json:"ipv4"`
 }
 
-type connections struct {
+type baseConnections struct {
 	All          connectionType `json:"all"`
 	Abridged     connectionType `json:"abridged"`
 	Intermediate connectionType `json:"intermediate"`
 	Secure       connectionType `json:"secure"`
 }
 
+type connections struct {
+	baseConnections
+}
+
 func (c connections) MarshalJSON() ([]byte, error) {
 	c.All.IPv4 = c.Abridged.IPv4 + c.Intermediate.IPv4 + c.Secure.IPv4
 	c.All.IPv6 = c.Abridged.IPv6 + c.Intermediate.IPv6 + c.Secure.IPv6
 
-	value := struct {
-		All          connectionType `json:"all"`
-		Abridged     connectionType `json:"abridged"`
-		Intermediate connectionType `json:"intermediate"`
-		Secure       connectionType `json:"secure"`
-	}{
-		All:          c.All,
-		Abridged:     c.Abridged,
-		Intermediate: c.Intermediate,
-		Secure:       c.Secure,
+	return json.Marshal(c.baseConnections)
+}
+
+type traffic struct {
+	ingress uint64
+	egress  uint64
+}
+
+func (t *traffic) dumpValue(value uint64) map[string]interface{} {
+	return map[string]interface{}{
+		"bytes": value,
+		"human": humanize.Bytes(value),
+	}
+}
+
+func (t traffic) MarshalJSON() ([]byte, error) {
+	value := map[string]map[string]interface{}{
+		"ingress": t.dumpValue(t.ingress),
+		"egress":  t.dumpValue(t.egress),
 	}
 
 	return json.Marshal(value)
 }
 
-type connectionType struct {
-	IPv6 uint32 `json:"ipv6"`
-	IPv4 uint32 `json:"ipv4"`
+type speed struct {
+	ingress uint64
+	egress  uint64
 }
 
-type traffic struct {
-	Ingress trafficValue `json:"ingress"`
-	Egress  trafficValue `json:"egress"`
+func (s *speed) dumpValue(value uint64) map[string]interface{} {
+	return map[string]interface{}{
+		"bytes/s": value,
+		"human":   fmt.Sprintf("%s/s", humanize.Bytes(value)),
+	}
 }
 
-type speed struct {
-	Ingress trafficSpeedValue `json:"ingress"`
-	Egress  trafficSpeedValue `json:"egress"`
+func (s speed) MarshalJSON() ([]byte, error) {
+	value := map[string]map[string]interface{}{
+		"ingress": s.dumpValue(s.ingress),
+		"egress":  s.dumpValue(s.egress),
+	}
+
+	return json.Marshal(value)
 }
 
-type stats struct {
+type Stats struct {
 	URLs        config.IPURLs `json:"urls"`
 	Connections connections   `json:"connections"`
 	Traffic     traffic       `json:"traffic"`
@@ -97,6 +97,77 @@ type stats struct {
 	Uptime      uptime        `json:"uptime"`
 	Crashes     uint32        `json:"crashes"`
 
-	speedCurrent speed
-	mutex        *sync.RWMutex
+	previousTraffic traffic
+}
+
+func (s *Stats) start() {
+	speedChan := time.Tick(time.Second)
+
+	for {
+		select {
+		case <-speedChan:
+			s.handleSpeed()
+		case event := <-trafficChan:
+			s.handleTraffic(event)
+		case event := <-connectionsChan:
+			s.handleConnection(event)
+		case getStatsChan := <-statsChan:
+			s.handleGetStats(getStatsChan)
+		case <-crashesChan:
+			s.handleCrash()
+		}
+	}
+}
+
+func (s *Stats) handleTraffic(evt trafficData) {
+	if evt.ingress {
+		s.Traffic.ingress += uint64(evt.traffic)
+	} else {
+		s.Traffic.egress += uint64(evt.traffic)
+	}
+}
+
+func (s *Stats) handleSpeed() {
+	s.Speed.ingress = s.Traffic.ingress - s.previousTraffic.ingress
+	s.Speed.egress = s.Traffic.egress - s.previousTraffic.egress
+	s.previousTraffic.ingress = s.Traffic.ingress
+	s.previousTraffic.egress = s.Traffic.egress
+}
+
+func (s *Stats) handleConnection(evt connectionData) {
+	var inc uint32 = 1
+	if !evt.connected {
+		inc = ^uint32(0)
+	}
+
+	var conn *connectionType
+	switch evt.connectionType {
+	case mtproto.ConnectionTypeAbridged:
+		conn = &s.Connections.Abridged
+	case mtproto.ConnectionTypeSecure:
+		conn = &s.Connections.Secure
+	default:
+		conn = &s.Connections.Intermediate
+	}
+
+	if evt.addr.IP.To4() != nil {
+		conn.IPv4 += inc
+	} else {
+		conn.IPv6 += inc
+	}
+}
+
+func (s *Stats) handleGetStats(getStatsChan chan<- Stats) {
+	getStatsChan <- *s
+}
+
+func (s *Stats) handleCrash() {
+	s.Crashes++
+}
+
+func NewStats(conf *config.Config) *Stats {
+	return &Stats{
+		URLs:   conf.GetURLs(),
+		Uptime: uptime(time.Now()),
+	}
 }

stats/statsd.go

@@ -36,7 +36,7 @@ type statsdExporter struct {
 
 func (s *statsdExporter) run() {
 	for range time.Tick(statsdPollTime) {
-		instance.mutex.Lock()
+		instance := GetStats()
 
 		s.client.Gauge(statsdConnectionsAbridgedV4, instance.Connections.Abridged.IPv4)
 		s.client.Gauge(statsdConnectionsAbridgedV6, instance.Connections.Abridged.IPv6)
@@ -44,13 +44,11 @@ func (s *statsdExporter) run() {
 		s.client.Gauge(statsdConnectionsIntermediateV6, instance.Connections.Intermediate.IPv6)
 		s.client.Gauge(statsdConnectionsSecureV4, instance.Connections.Secure.IPv4)
 		s.client.Gauge(statsdConnectionsSecureV6, instance.Connections.Secure.IPv6)
-		s.client.Gauge(statsdTrafficIngress, uint64(instance.Traffic.Ingress))
-		s.client.Gauge(statsdTrafficEgress, uint64(instance.Traffic.Egress))
-		s.client.Gauge(statsdSpeedIngress, uint64(instance.Speed.Ingress))
-		s.client.Gauge(statsdSpeedEgress, uint64(instance.Speed.Egress))
+		s.client.Gauge(statsdTrafficIngress, instance.Traffic.ingress)
+		s.client.Gauge(statsdTrafficEgress, instance.Traffic.egress)
+		s.client.Gauge(statsdSpeedIngress, instance.Speed.ingress)
+		s.client.Gauge(statsdSpeedEgress, instance.Speed.egress)
 		s.client.Gauge(statsdCrashes, instance.Crashes)
-
-		instance.mutex.Unlock()
 	}
 }