From 0c28c9ea599e462a941c155ecdd4a67f32a9814e Mon Sep 17 00:00:00 2001
From: Paul Lhussiez <paul.lhussiez@schibsted.com>
Date: Tue, 31 Oct 2017 11:05:56 +0100
Subject: [PATCH] Fixing concurrency access to map

---
 README.md |  6 +++---
 prom.go   | 38 +++++++++++++++++++++++++++-----------
 2 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/README.md b/README.md
index 0f95cf4..e96498c 100644
--- a/README.md
+++ b/README.md
@@ -17,7 +17,6 @@ Simply run :
 ## Differences with go-gin-prometheus
 
 - No support for Prometheus' Push Gateway
-- Need to call middleware on each route
 - Options on constructor
 - Adds a `path` label to get the matched route
 
@@ -34,9 +33,10 @@ import (
 func main() {
 	r := gin.Default()
 	p := ginprom.New(ginprom.Subsystem("gin"), ginprom.Path("/metrics"), ginprom.Engine(r))
+	r.Use(p.Instrument())
 
-	r.GET("/hello/:id", p.Instrument("/hello/:id"), func(c *gin.Context) {})
-	r.GET("/world/:id", p.Instrument("/world/:id"), func(c *gin.Context) {})
+	r.GET("/hello/:id", func(c *gin.Context) {})
+	r.GET("/world/:id", func(c *gin.Context) {})
 	r.Run("127.0.0.1:8080")
 }
 ```
diff --git a/prom.go b/prom.go
index de605db..c475bb6 100644
--- a/prom.go
+++ b/prom.go
@@ -2,6 +2,7 @@ package ginprom
 
 import (
 	"strconv"
+	"sync"
 	"time"
 
 	"github.com/gin-gonic/gin"
@@ -12,6 +13,11 @@ import (
 var defaultPath = "/metrics"
 var defaultSys = "gin"
 
+type pmap struct {
+	sync.RWMutex
+	values map[string]string
+}
+
 // Prometheus contains the metrics gathered by the instance and its path
 type Prometheus struct {
 	reqCnt               *prometheus.CounterVec
@@ -20,7 +26,7 @@ type Prometheus struct {
 	MetricsPath string
 	Subsystem   string
 	Engine      *gin.Engine
-	PathMap     map[string]string
+	PathMap     pmap
 }
 
 // Path is an option allowing to set the metrics path when intializing with New.
@@ -66,14 +72,31 @@ func New(options ...func(*Prometheus)) *Prometheus {
 	if p.Engine != nil {
 		p.Engine.GET(p.MetricsPath, prometheusHandler())
 	}
-	p.PathMap = make(map[string]string)
+	p.PathMap.values = make(map[string]string)
 	return p
 }
 
 func (p *Prometheus) updatePathMap() {
+	p.PathMap.Lock()
+	defer p.PathMap.Unlock()
 	for _, ri := range p.Engine.Routes() {
-		p.PathMap[ri.Handler] = ri.Path
+		p.PathMap.values[ri.Handler] = ri.Path
+	}
+}
+
+func (p *Prometheus) getPathFromHandler(handler string) string {
+	p.PathMap.RLock()
+	defer p.PathMap.RUnlock()
+	if in, ok := p.PathMap.values[handler]; ok {
+		return in
+	}
+	p.PathMap.RUnlock()
+	p.updatePathMap()
+	p.PathMap.RLock()
+	if in, ok := p.PathMap.values[handler]; ok {
+		return in
 	}
+	return ""
 }
 
 func (p *Prometheus) register() {
@@ -129,14 +152,7 @@ func (p *Prometheus) Instrument() gin.HandlerFunc {
 			return
 		}
 
-		if in, ok := p.PathMap[c.HandlerName()]; ok {
-			path = in
-		} else {
-			p.updatePathMap()
-			if in, ok := p.PathMap[c.HandlerName()]; ok {
-				path = in
-			}
-		}
+		path = p.getPathFromHandler(c.HandlerName())
 
 		c.Next()