From c01e3b3b882b29ac48a2fe4172abb4a28d9dce4f Mon Sep 17 00:00:00 2001
From: noe <git@kat.cafe>
Date: Wed, 26 Jun 2024 20:18:54 -0400
Subject: [PATCH] update prom logic

---
 hook_handler.go |  4 ++--
 main.go         |  2 ++
 pluralkit.go    | 49 +++++++++++++++++++++++++++++++++++++++++++++++--
 prometheus.go   | 20 +++++++++-----------
 4 files changed, 60 insertions(+), 15 deletions(-)

diff --git a/hook_handler.go b/hook_handler.go
index 9fa7502..3bc4525 100644
--- a/hook_handler.go
+++ b/hook_handler.go
@@ -75,11 +75,11 @@ func handleHookCreateSwitch(h HookPayload) {
 	}
 
 	if len(sw.Members) != 0 {
-		front = &sw.Members[0]
+		front = sw.Members[0]
 	}
 
 	go dsiApi.UpdateFrontField(front.Name)
-	go promCountSwitches(h, sw.Members)
+	go promCountSwitches(h.SystemID, sw.Members)
 	// TODO: discord nickname updates
 }
 
diff --git a/main.go b/main.go
index 9a3bf96..f85dea3 100644
--- a/main.go
+++ b/main.go
@@ -24,6 +24,8 @@ func main() {
 	mux.HandleFunc("/hook/{token}", postGetHookToken)
 	mux.HandleFunc("/healthcheck", getHealthcheck)
 
+	go populateSwitches()
+
 	// prometheus
 	mux.Handle("/metrics", promhttp.Handler())
 
diff --git a/pluralkit.go b/pluralkit.go
index 0a45248..0ede26f 100644
--- a/pluralkit.go
+++ b/pluralkit.go
@@ -3,7 +3,9 @@ package main
 import (
 	"encoding/json"
 	"fmt"
+	"log"
 	"net/http"
+	"os"
 	"time"
 
 	"github.com/hashicorp/golang-lru/v2/expirable"
@@ -14,7 +16,11 @@ var (
 )
 
 type PluralKitSwitch struct {
-	Members []PluralKitMember
+	Members []*PluralKitMember
+}
+
+type PluralKitSwitches struct {
+	Members []string
 }
 
 type PluralKitMember struct {
@@ -79,8 +85,47 @@ func (pk *PluralKit) GetSwitch(systemID string, switchID string) (*PluralKitSwit
 	}
 
 	for _, member := range sw.Members {
-		pk.Cache.Add(member.UUID, &member)
+		pk.Cache.Add(member.UUID, member)
 	}
 
 	return &sw, nil
 }
+
+func (pk *PluralKit) GetSwitches(systemID string) (*PluralKitSwitches, error) {
+	resp, err := pk.Client.Get(fmt.Sprintf("%s/systems/%s/switches", pk.BaseURL, systemID))
+	if err != nil {
+		return nil, fmt.Errorf("fetching pluralkit system %s failed: %w", systemID, err)
+	}
+
+	if resp.StatusCode != 200 {
+		return nil, fmt.Errorf("pluralkit did not find system %s", systemID)
+	}
+
+	var sw []PluralKitSwitches
+	err = json.NewDecoder(resp.Body).Decode(&sw)
+	if err != nil {
+		return nil, fmt.Errorf("malformed json: %w", err)
+	}
+
+	return &sw[0], nil
+}
+
+func populateSwitches() {
+	sysID := os.Getenv("PK_SYSTEM_ID")
+	switches, err := pkApi.GetSwitches(sysID)
+	if err != nil {
+		log.Fatalln("failed to prepopulate switches.", err)
+	}
+
+	members := make([]*PluralKitMember, len(switches.Members))
+	for idx, id := range switches.Members {
+		m, err := pkApi.GetMember(id)
+		if err != nil {
+			log.Fatalln("failed to get member.", err)
+		}
+
+		members[idx] = m
+	}
+
+	promCountSwitches(sysID, members)
+}
diff --git a/prometheus.go b/prometheus.go
index cbf6bd3..df6dfb8 100644
--- a/prometheus.go
+++ b/prometheus.go
@@ -15,7 +15,7 @@ var (
 	switchGauge *prometheus.GaugeVec = promauto.NewGaugeVec(prometheus.GaugeOpts{
 		Namespace: "plapkit",
 		Name:      "switches",
-	}, []string{"system", "member", "member_display", "role"})
+	}, []string{"system", "member", "member_display"})
 
 	messageCounter *prometheus.CounterVec = promauto.NewCounterVec(prometheus.CounterOpts{
 		Namespace: "plapkit",
@@ -29,24 +29,22 @@ func clearHighSwitches() {
 	}
 }
 
-func promCountSwitches(h HookPayload, members []PluralKitMember) {
+func promCountSwitches(systemID string, members []*PluralKitMember) {
 	clearHighSwitches()
 
 	for idx, member := range members {
-		role := "cofront"
-		if idx == 0 {
-			role = "front"
-		}
-
 		labels := prometheus.Labels{
-			"system":         h.SystemID,
+			"system":         systemID,
 			"member":         member.UUID,
 			"member_display": member.Name,
-			"role":           role,
 		}
 
-		switchGauge.With(labels).Set(1)
-		knownHighSwitches[member.UUID+":"+role] = labels
+		gv := 1.0
+		if idx != 0 {
+			gv = -1.0
+		}
+		switchGauge.With(labels).Set(gv)
+		knownHighSwitches[member.UUID] = labels
 	}
 }