Merge pull request #13 from 9seconds/block-rwc

Stats server

Gopkg.lock

@@ -22,6 +22,12 @@
   revision = "346938d642f2ec3594ed81d874461961cd0faa76"
   version = "v1.1.0"
 
+[[projects]]
+  branch = "master"
+  name = "github.com/dustin/go-humanize"
+  packages = ["."]
+  revision = "02af3965c54e8cacf948b97fef38925c4120652c"
+
 [[projects]]
   branch = "master"
   name = "github.com/juju/errors"
@@ -80,6 +86,6 @@
 [solve-meta]
   analyzer-name = "dep"
   analyzer-version = 1
-  inputs-digest = "24afdd6b64331aeba47fed75918d04032e13e404612cac107bad1d68a5038b72"
+  inputs-digest = "312c9fb15085cbe9660443b15a07981990e1f70ec3ddfcce1b7e6cd5902307da"
   solver-name = "gps-cdcl"
   solver-version = 1

Gopkg.toml

@@ -44,3 +44,7 @@
 [[constraint]]
   name = "github.com/satori/go.uuid"
   version = "1.2.0"
+
+[[constraint]]
+  branch = "master"
+  name = "github.com/dustin/go-humanize"

main.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/proxy"
+	"github.com/9seconds/mtg/stats"
 	"github.com/juju/errors"
 )
 
@@ -115,13 +116,15 @@ func main() {
 	zap.ReplaceGlobals(logger)
 	defer logger.Sync()
 
+	printURLs(conf.GetURLs())
+
 	if conf.UseMiddleProxy() {
 		zap.S().Infow("Use middle proxy connection to Telegram")
 	} else {
 		zap.S().Infow("Use direct connection to Telegram")
 	}
 
-	printURLs(conf.GetURLs())
+	go stats.Start(conf)
 
 	server := proxy.NewProxy(conf)
 	if err := server.Serve(); err != nil {

proxy/proxy.go

@@ -12,6 +12,7 @@ import (
 	"github.com/9seconds/mtg/client"
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/stats"
 	"github.com/9seconds/mtg/telegram"
 	"github.com/9seconds/mtg/wrappers"
 )
@@ -39,12 +40,13 @@ func (p *Proxy) Serve() error {
 
 func (p *Proxy) accept(conn net.Conn) {
 	connID := uuid.NewV4().String()
-	log := zap.S().With("connection_id", connID)
+	log := zap.S().With("connection_id", connID).Named("main")
 
 	defer func() {
 		conn.Close()
 
 		if err := recover(); err != nil {
+			stats.NewCrash()
 			log.Errorw("Crash of accept handler", "error", err)
 		}
 	}()
@@ -58,6 +60,9 @@ func (p *Proxy) accept(conn net.Conn) {
 	}
 	defer client.(io.Closer).Close()
 
+	stats.ClientConnected(opts.ConnectionType, client.RemoteAddr())
+	defer stats.ClientDisconnected(opts.ConnectionType, client.RemoteAddr())
+
 	server, err := p.getTelegramConn(opts, connID)
 	if err != nil {
 		log.Errorw("Cannot initialize server connection", "error", err)

stats/channels.go

@@ -0,0 +1,147 @@
+package stats
+
+import (
+	"net"
+	"sync/atomic"
+	"time"
+
+	"github.com/9seconds/mtg/mtproto"
+)
+
+const (
+	crashesChanLength     = 1
+	connectionsChanLength = 20
+	trafficChanLength     = 5000
+)
+
+var (
+	CrashesChan     = make(chan struct{}, crashesChanLength)
+	ConnectionsChan = make(chan *connectionData, connectionsChanLength)
+	TrafficChan     = make(chan *trafficData, trafficChanLength)
+)
+
+type connectionData struct {
+	connectionType mtproto.ConnectionType
+	addr           *net.TCPAddr
+	connected      bool
+}
+
+type trafficData struct {
+	traffic int
+	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 {
+				atomic.AddUint32(&instance.ActiveConnections.Abridged.IPv4, inc)
+				if event.connected {
+					atomic.AddUint32(&instance.AllConnections.Abridged.IPv4, inc)
+				}
+			} else {
+				atomic.AddUint32(&instance.ActiveConnections.Abridged.IPv6, inc)
+				if event.connected {
+					atomic.AddUint32(&instance.AllConnections.Abridged.IPv6, inc)
+				}
+			}
+		default:
+			if isIPv4 {
+				atomic.AddUint32(&instance.ActiveConnections.Intermediate.IPv4, inc)
+				if event.connected {
+					atomic.AddUint32(&instance.AllConnections.Intermediate.IPv4, inc)
+				}
+			} else {
+				atomic.AddUint32(&instance.ActiveConnections.Intermediate.IPv6, inc)
+				if event.connected {
+					atomic.AddUint32(&instance.AllConnections.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()
+		}
+	}
+}
+
+func NewCrash() {
+	CrashesChan <- struct{}{}
+}
+
+func ClientConnected(connectionType mtproto.ConnectionType, addr *net.TCPAddr) {
+	ConnectionsChan <- &connectionData{
+		connectionType: connectionType,
+		addr:           addr,
+		connected:      true,
+	}
+}
+
+func ClientDisconnected(connectionType mtproto.ConnectionType, addr *net.TCPAddr) {
+	ConnectionsChan <- &connectionData{
+		connectionType: connectionType,
+		addr:           addr,
+		connected:      false,
+	}
+}
+
+func IngressTraffic(traffic int) {
+	TrafficChan <- &trafficData{
+		traffic: traffic,
+		ingress: true,
+	}
+}
+
+func EgressTraffic(traffic int) {
+	TrafficChan <- &trafficData{
+		traffic: traffic,
+		ingress: false,
+	}
+}

stats/server.go

@@ -0,0 +1,56 @@
+package stats
+
+import (
+	"encoding/json"
+	"net/http"
+	"sync"
+	"time"
+
+	"go.uber.org/zap"
+
+	"github.com/9seconds/mtg/config"
+)
+
+var instance *stats
+
+func Start(conf *config.Config) {
+	log := zap.S().Named("stats")
+
+	instance = &stats{
+		URLs:   conf.GetURLs(),
+		Uptime: uptime(time.Now()),
+		mutex:  &sync.RWMutex{},
+	}
+
+	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()
+
+		if err != nil {
+			log.Errorw("Cannot encode json", "error", err)
+			http.Error(w, "Internal server error", 500)
+			return
+		}
+
+		interm := map[string]interface{}{}
+		json.Unmarshal(first, &interm)
+
+		encoder := json.NewEncoder(w)
+		encoder.SetEscapeHTML(false)
+		encoder.SetIndent("", "  ")
+		if err = encoder.Encode(interm); err != nil {
+			log.Errorw("Cannot encode json", "error", err)
+		}
+	})
+
+	if err := http.ListenAndServe(conf.StatAddr(), nil); err != nil {
+		log.Fatalw("Stats server has been stopped", "error", err)
+	}
+}

stats/stats.go

@@ -0,0 +1,100 @@
+package stats
+
+import (
+	"encoding/json"
+	"fmt"
+	"strconv"
+	"sync"
+	"time"
+
+	humanize "github.com/dustin/go-humanize"
+
+	"github.com/9seconds/mtg/config"
+)
+
+type uptime time.Time
+
+func (u uptime) MarshalJSON() ([]byte, error) {
+	duration := time.Since(time.Time(u))
+	value := map[string]string{
+		"seconds": strconv.Itoa(int(duration.Seconds())),
+		"human":   humanize.Time(time.Time(u)),
+	}
+
+	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 connections struct {
+	All          connectionType `json:"all"`
+	Abridged     connectionType `json:"abridged"`
+	Intermediate connectionType `json:"intermediate"`
+}
+
+func (c connections) MarshalJSON() ([]byte, error) {
+	c.All.IPv4 = c.Abridged.IPv4 + c.Intermediate.IPv4
+	c.All.IPv6 = c.Abridged.IPv6 + c.Intermediate.IPv6
+
+	value := struct {
+		All          connectionType `json:"all"`
+		Abridged     connectionType `json:"abridged"`
+		Intermediate connectionType `json:"intermediate"`
+	}{
+		All:          c.All,
+		Abridged:     c.Abridged,
+		Intermediate: c.Intermediate,
+	}
+
+	return json.Marshal(value)
+}
+
+type connectionType struct {
+	IPv6 uint32 `json:"ipv6"`
+	IPv4 uint32 `json:"ipv4"`
+}
+
+type traffic struct {
+	Ingress trafficValue `json:"ingress"`
+	Egress  trafficValue `json:"egress"`
+}
+
+type speed struct {
+	Ingress trafficSpeedValue `json:"ingress"`
+	Egress  trafficSpeedValue `json:"egress"`
+}
+
+type stats struct {
+	URLs              config.IPURLs `json:"urls"`
+	ActiveConnections connections   `json:"active_connections"`
+	AllConnections    connections   `json:"all_connections"`
+	Traffic           traffic       `json:"traffic"`
+	Speed             speed         `json:"speed"`
+	Uptime            uptime        `json:"uptime"`
+	Crashes           uint32        `json:"crashes"`
+
+	speedCurrent speed
+	mutex        *sync.RWMutex
+}

wrappers/conn.go

@@ -5,6 +5,8 @@ import (
 	"time"
 
 	"go.uber.org/zap"
+
+	"github.com/9seconds/mtg/stats"
 )
 
 type ConnPurpose uint8
@@ -31,9 +33,10 @@ const (
 )
 
 type Conn struct {
-	connID     string
+	connID string
-	conn       net.Conn
+	conn   net.Conn
-	logger     *zap.SugaredLogger
+	logger *zap.SugaredLogger
+
 	publicIPv4 net.IP
 	publicIPv6 net.IP
 }
@@ -43,6 +46,7 @@ func (c *Conn) Write(p []byte) (int, error) {
 	n, err := c.conn.Write(p)
 
 	c.logger.Debugw("Write to stream", "bytes", n, "error", err)
+	stats.EgressTraffic(n)
 
 	return n, err
 }
@@ -52,6 +56,7 @@ func (c *Conn) Read(p []byte) (int, error) {
 	n, err := c.conn.Read(p)
 
 	c.logger.Debugw("Read from stream", "bytes", n, "error", err)
+	stats.IngressTraffic(n)
 
 	return n, err
 }