From a65d7a516b3836168577561d5334b665fef9fbad Mon Sep 17 00:00:00 2001
From: pg9182 <96569817+pg9182@users.noreply.github.com>
Date: Sun, 27 Nov 2022 01:19:59 -0500
Subject: pkg/atlas: Implement geo metrics

---
 pkg/atlas/config.go | 3 ++-
 pkg/atlas/server.go | 7 ++++++-
 2 files changed, 8 insertions(+), 2 deletions(-)

(limited to 'pkg')

diff --git a/pkg/atlas/config.go b/pkg/atlas/config.go
index 149c9ab..af70817 100644
--- a/pkg/atlas/config.go
+++ b/pkg/atlas/config.go
@@ -172,7 +172,8 @@ type Config struct {
 	// country and region fields. The database must not be modified while atlas
 	// is running, but it can be replaced (and a reload can be triggered with
 	// SIGHUP). If not provided, geolocation-dependent features like server
-	// regions will not be enabled.
+	// regions and geo metrics will not be enabled. If it doesn't include latlon
+	// info, geo metrics will be disabled too.
 	IP2Location string `env:"ATLAS_IP2LOCATION"`
 
 	// For sd-notify.
diff --git a/pkg/atlas/server.go b/pkg/atlas/server.go
index 55f1b5c..3c881dd 100644
--- a/pkg/atlas/server.go
+++ b/pkg/atlas/server.go
@@ -763,12 +763,13 @@ func (s *Server) HandleSIGHUP() {
 // serveRest handles endpoints not handled by the API.
 func (s *Server) serveRest(w http.ResponseWriter, r *http.Request) {
 	if r.URL.Path == "/metrics" {
-		var internal bool
+		var internal, geo bool
 		if s := s.MetricsSecret; s != "" {
 			if r.URL.Query().Get("secret") == s {
 				internal = true
 			}
 		}
+		geo = r.URL.Query().Has("geo")
 
 		var ms []func(io.Writer)
 		if internal {
@@ -776,6 +777,10 @@ func (s *Server) serveRest(w http.ResponseWriter, r *http.Request) {
 			ms = append(ms, s.API0.WritePrometheus)
 		}
 		ms = append(ms, s.API0.ServerList.WritePrometheus)
+		if internal && geo {
+			ms = append(ms, s.API0.WritePrometheusGeo)
+			ms = append(ms, s.API0.ServerList.WritePrometheusGeo)
+		}
 
 		var b bytes.Buffer
 		for i, m := range ms {
-- 
cgit v1.2.3