feat: Add metrics endpoint

pojntfx · Aug 3, 2022 · ab9656c · ab9656c
1 parent 33740be
commit ab9656c
Show file tree

Hide file tree

Showing 3 changed files with 76 additions and 19 deletions.
diff --git a/cmd/htorrent/cmd/gateway.go b/cmd/htorrent/cmd/gateway.go
@@ -9,6 +9,7 @@ import (
 	"strconv"
 	"syscall"
 
+	v1 "github.com/pojntfx/htorrent/pkg/api/http/v1"
 	"github.com/pojntfx/htorrent/pkg/server"
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
@@ -60,12 +61,13 @@ var gatewayCmd = &cobra.Command{
 			viper.GetString(oidcIssuerFlag),
 			viper.GetString(oidcClientIDFlag),
 			viper.GetInt(verboseFlag) > 5,
-			func(peers int, total, completed int64, path string) {
+			func(torrentMetrics v1.TorrentMetrics, fileMetrics v1.FileMetrics) {
 				log.Debug().
-					Int("peers", peers).
-					Int64("total", total).
-					Int64("completed", completed).
-					Str("path", path).
+					Str("magnet", torrentMetrics.Magnet).
+					Int("peers", torrentMetrics.Peers).
+					Str("path", fileMetrics.Path).
+					Int64("length", fileMetrics.Length).
+					Int64("completed", fileMetrics.Completed).
 					Msg("Streaming")
 			},
 			ctx,

diff --git a/pkg/api/http/v1/info.go b/pkg/api/http/v1/info.go
@@ -11,3 +11,15 @@ type File struct {
 	Path   string `json:"path"`
 	Length int64  `json:"length"`
 }
+
+type TorrentMetrics struct {
+	Magnet string        `json:"magnet"`
+	Peers  int           `json:"peers"`
+	Files  []FileMetrics `json:"files"`
+}
+
+type FileMetrics struct {
+	Path      string `json:"path"`
+	Length    int64  `json:"length"`
+	Completed int64  `json:"completed"`
+}
diff --git a/pkg/server/gateway.go b/pkg/server/gateway.go
@@ -36,7 +36,7 @@ type Gateway struct {
 	oidcClientID string
 	debug        bool
 
-	onDownloadProgress func(peers int, total, completed int64, path string)
+	onDownloadProgress func(torrentMetrics v1.TorrentMetrics, fileMetrics v1.FileMetrics)
 
 	torrentClient *torrent.Client
 	srv           *http.Server
@@ -55,7 +55,7 @@ func NewGateway(
 	oidcClientID string,
 	debug bool,
 
-	onDownloadProgress func(peers int, total, completed int64, path string),
+	onDownloadProgress func(torrentMetrics v1.TorrentMetrics, fileMetrics v1.FileMetrics),
 
 	ctx context.Context,
 ) *Gateway {
@@ -170,6 +170,50 @@ func (g *Gateway) Open() error {
 		}
 	})
 
+	mux.HandleFunc("/metrics", func(w http.ResponseWriter, r *http.Request) {
+		u, p, ok := r.BasicAuth()
+		if err := auth.Validate(u, p); !ok || err != nil {
+			w.WriteHeader(http.StatusUnauthorized)
+
+			panic(fmt.Errorf("%v", http.StatusUnauthorized))
+		}
+
+		log.Debug().
+			Msg("Getting metrics")
+
+		metrics := []v1.TorrentMetrics{}
+		for _, t := range g.torrentClient.Torrents() {
+			mi := t.Metainfo()
+
+			info, err := mi.UnmarshalInfo()
+			if err != nil {
+				panic(err)
+			}
+
+			fileMetrics := []v1.FileMetrics{}
+			for _, f := range t.Files() {
+				fileMetrics = append(fileMetrics, v1.FileMetrics{
+					Path:      f.Path(),
+					Length:    f.Length(),
+					Completed: f.BytesCompleted(),
+				})
+			}
+
+			torrentMetrics := v1.TorrentMetrics{
+				Magnet: mi.Magnet(nil, &info).String(),
+				Peers:  len(t.PeerConns()),
+				Files:  fileMetrics,
+			}
+
+			metrics = append(metrics, torrentMetrics)
+		}
+
+		enc := json.NewEncoder(w)
+		if err := enc.Encode(metrics); err != nil {
+			panic(err)
+		}
+	})
+
 	mux.HandleFunc("/stream", func(w http.ResponseWriter, r *http.Request) {
 		defer func() {
 			err := recover()
@@ -238,20 +282,19 @@ func (g *Gateway) Open() error {
 
 				lastCompleted := int64(0)
 				for range tick.C {
-					if completed, total := f.BytesCompleted(), f.Length(); completed < total {
+					if completed, length := f.BytesCompleted(), f.Length(); completed < length {
 						if completed != lastCompleted {
-							log.Debug().
-								Int("peers", len(f.Torrent().PeerConns())).
-								Int64("total", total).
-								Int64("completed", completed).
-								Str("path", f.Path()).
-								Msg("Streaming")
-
 							g.onDownloadProgress(
-								len(f.Torrent().PeerConns()),
-								total,
-								completed,
-								f.Path(),
+								v1.TorrentMetrics{
+									Magnet: magnetLink,
+									Peers:  len(f.Torrent().PeerConns()),
+									Files:  []v1.FileMetrics{},
+								},
+								v1.FileMetrics{
+									Path:      f.Path(),
+									Length:    length,
+									Completed: completed,
+								},
 							)
 						}