From cac9cfe3e61fafc78adf8344524b5aa858fef0f8 Mon Sep 17 00:00:00 2001
From: Klaus <32713077+hirosumee@users.noreply.github.com>
Date: Wed, 21 Sep 2022 16:38:39 +0700
Subject: [PATCH] using handler pattern instead of url mapping

---
 prometheus.go | 39 ++++++++++++++++-----------------------
 1 file changed, 16 insertions(+), 23 deletions(-)

diff --git a/prometheus.go b/prometheus.go
index 6a886e5..d96dadd 100644
--- a/prometheus.go
+++ b/prometheus.go
@@ -3,7 +3,6 @@ package fiberprom
 import (
 	"log"
 	"strconv"
-	"strings"
 	"time"
 
 	"github.com/gofiber/adaptor/v2"
@@ -19,8 +18,6 @@ type Prometheus struct {
 	reqDur      *prometheus.HistogramVec
 	router      fiber.Router
 	MetricsPath string
-	// urlMapper is a map of url to be mapped to a different url to avoid too many labels
-	urlMapper map[string]string
 }
 
 // NewPrometheus generates a new set of metrics with a certain subsystem name
@@ -33,10 +30,6 @@ func NewPrometheus(subsystem string) *Prometheus {
 	return p
 }
 
-func (p *Prometheus) SetURLMapper(mapper map[string]string) {
-	p.urlMapper = mapper
-}
-
 func (p *Prometheus) registerMetrics(subsystem string) {
 	p.reqDur = prometheus.NewHistogramVec(
 		prometheus.HistogramOpts{
@@ -69,25 +62,25 @@ func (p *Prometheus) HandlerFunc() fiber.Handler {
 			return ctx.Next()
 		}
 
-		if len(p.urlMapper) > 0 {
-			for prefix, val := range p.urlMapper {
-				if strings.HasPrefix(uri, prefix) {
-					uri = val
-					break
-				}
+		start := time.Now()
+		// next
+		err := ctx.Next()
+
+		status := fiber.StatusInternalServerError
+		if err != nil {
+			if e, ok := err.(*fiber.Error); ok {
+				// Get correct error code from fiber.Error type
+				status = e.Code
 			}
+		} else {
+			status = ctx.Response().StatusCode()
 		}
+		uri = ctx.Route().Path
+		elapsed := float64(time.Since(start)) / float64(time.Second)
+		ep := ctx.Method() + "_" + uri
+		p.reqDur.WithLabelValues(strconv.Itoa(status), ep).Observe(elapsed)
 
-		start := time.Now()
-
-		defer func() {
-			status := strconv.Itoa(ctx.Response().StatusCode())
-			elapsed := float64(time.Since(start)) / float64(time.Second)
-			ep := ctx.Method() + "_" + uri
-			p.reqDur.WithLabelValues(status, ep).Observe(elapsed)
-		}()
-		// next
-		return ctx.Next()
+		return err
 	}
 }