From 7b544f370dfc86cdeec18fcd1c4fb6cc852759f5 Mon Sep 17 00:00:00 2001
From: LittleSheep <littlesheep.code@hotmail.com>
Date: Wed, 31 Jan 2024 13:17:52 +0800
Subject: [PATCH] :sparkles: More detailed metrics

---
 pkg/hypertext/proxies.go |  2 +-
 pkg/navi/config.go       |  9 ++++++---
 pkg/navi/metrics.go      | 30 +++++++++++++++++++++++++++++-
 pkg/navi/route.go        |  4 ++--
 pkg/sideload/metrics.go  |  8 ++++++--
 pkg/sideload/server.go   |  1 +
 6 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/pkg/hypertext/proxies.go b/pkg/hypertext/proxies.go
index 8b0c15a..3b12566 100644
--- a/pkg/hypertext/proxies.go
+++ b/pkg/hypertext/proxies.go
@@ -121,7 +121,7 @@ func makeResponse(c *fiber.Ctx, region *navi.Region, location *navi.Location, de
 			message = err.Error()
 		}
 
-		go navi.R.AddTrace(navi.RoadTrace{
+		go navi.R.Metrics.AddTrace(navi.RoadTrace{
 			Region:      region.ID,
 			Location:    location.ID,
 			Destination: dest.ID,
diff --git a/pkg/navi/config.go b/pkg/navi/config.go
index f165901..733080a 100644
--- a/pkg/navi/config.go
+++ b/pkg/navi/config.go
@@ -6,8 +6,6 @@ import (
 	"path/filepath"
 	"strings"
 
-	"github.com/spf13/viper"
-
 	"github.com/pelletier/go-toml/v2"
 )
 
@@ -16,7 +14,12 @@ var R *RoadApp
 func ReadInConfig(root string) error {
 	instance := &RoadApp{
 		Regions: make([]*Region, 0),
-		Traces:  make([]RoadTrace, 0, viper.GetInt("performance.traces_limit")),
+		Metrics: &RoadMetrics{
+			Traces:       make([]RoadTrace, 0),
+			Traffic:      make(map[string]int64),
+			TrafficFrom:  make(map[string]int64),
+			TotalTraffic: 0,
+		},
 	}
 
 	if err := filepath.Walk(root, func(fp string, info os.FileInfo, _ error) error {
diff --git a/pkg/navi/metrics.go b/pkg/navi/metrics.go
index 550c2c7..e5f915e 100644
--- a/pkg/navi/metrics.go
+++ b/pkg/navi/metrics.go
@@ -2,6 +2,14 @@ package navi
 
 import "github.com/spf13/viper"
 
+type RoadMetrics struct {
+	Traces []RoadTrace `json:"-"`
+
+	Traffic      map[string]int64 `json:"traffic"`
+	TrafficFrom  map[string]int64 `json:"traffic_from"`
+	TotalTraffic int64            `json:"total_traffic"`
+}
+
 type RoadTrace struct {
 	Region      string         `json:"region"`
 	Location    string         `json:"location"`
@@ -17,8 +25,28 @@ type RoadTraceError struct {
 	Message string `json:"message"`
 }
 
-func (v *RoadApp) AddTrace(trace RoadTrace) {
+func (v *RoadMetrics) AddTrace(trace RoadTrace) {
+	v.TotalTraffic++
+	if _, ok := v.Traffic[trace.Region]; !ok {
+		v.Traffic[trace.Region] = 0
+	} else {
+		v.Traffic[trace.Region]++
+	}
+	if _, ok := v.TrafficFrom[trace.IpAddress]; !ok {
+		v.TrafficFrom[trace.IpAddress] = 0
+	} else {
+		v.TrafficFrom[trace.IpAddress]++
+	}
+
 	v.Traces = append(v.Traces, trace)
+
+	// Garbage recycle
+	if len(v.Traffic) > viper.GetInt("performance.traces_limit") {
+		v.Traffic = make(map[string]int64)
+	}
+	if len(v.TrafficFrom) > viper.GetInt("performance.traces_limit") {
+		v.TrafficFrom = make(map[string]int64)
+	}
 	if len(v.Traces) > viper.GetInt("performance.traces_limit") {
 		v.Traces = v.Traces[1:]
 	}
diff --git a/pkg/navi/route.go b/pkg/navi/route.go
index f23aa33..dd9a7c4 100644
--- a/pkg/navi/route.go
+++ b/pkg/navi/route.go
@@ -10,8 +10,8 @@ import (
 )
 
 type RoadApp struct {
-	Regions []*Region   `json:"regions"`
-	Traces  []RoadTrace `json:"traces"`
+	Regions []*Region    `json:"regions"`
+	Metrics *RoadMetrics `json:"metrics"`
 }
 
 func (v *RoadApp) Forward(c *fiber.Ctx, dest *Destination) error {
diff --git a/pkg/sideload/metrics.go b/pkg/sideload/metrics.go
index 25d42ba..4fda5e3 100644
--- a/pkg/sideload/metrics.go
+++ b/pkg/sideload/metrics.go
@@ -5,6 +5,10 @@ import (
 	"github.com/gofiber/fiber/v2"
 )
 
+func getTraffic(c *fiber.Ctx) error {
+	return c.JSON(navi.R.Metrics)
+}
+
 func getTraces(c *fiber.Ctx) error {
-	return c.JSON(navi.R.Traces)
-}
\ No newline at end of file
+	return c.JSON(navi.R.Metrics.Traces)
+}
diff --git a/pkg/sideload/server.go b/pkg/sideload/server.go
index 6240b94..1ecbc95 100644
--- a/pkg/sideload/server.go
+++ b/pkg/sideload/server.go
@@ -45,6 +45,7 @@ func InitSideload() *fiber.App {
 	cgi := app.Group("/cgi").Name("CGI")
 	{
 		cgi.Get("/metadata", getMetadata)
+		cgi.Get("/traffic", getTraffic)
 		cgi.Get("/traces", getTraces)
 		cgi.Get("/stats", getStats)
 		cgi.Get("/regions", getRegions)