plugin/metrics: fix failed reload (#2816)

Fix metrics endpoint on a failed reload, follows the same lines as the
previous PRs, see for e.g. 076b8d4f. Test with a Corefile with 2 server
blocks and metrics enabled and then introducing a syntax error:

~~~
[ERROR] Restart failed: Corefile:5 - Error during parsing: Unknown directive 'jfkdjk'
[ERROR] SIGUSR1: starting with listener file descriptors: Corefile:5 - Error during parsing: Unknown directive 'jfkdjk'
~~~

And then curl-ing the metrics endpoint.

See #2659 and as this is the last one.

Fixes: #2659

Getting this all right turns out to be tricky, also it's not easy
testable which is something I should fix.

Signed-off-by: Miek Gieben <miek@miek.nl>

3 files changed, 25 insertions, 30 deletions

diff --git a/plugin/metrics/metrics.go b/plugin/metrics/metrics.go
index 89c948aa6..acf31c0a8 100644
--- a/plugin/metrics/metrics.go
+++ b/plugin/metrics/metrics.go
@@ -137,9 +137,7 @@ func (m *Metrics) stopServer() error {
 }
 
 // OnFinalShutdown tears down the metrics listener on shutdown and restart.
-func (m *Metrics) OnFinalShutdown() error {
-	return m.stopServer()
-}
+func (m *Metrics) OnFinalShutdown() error { return m.stopServer() }
 
 func keys(m map[string]struct{}) []string {
 	sx := []string{}
diff --git a/plugin/metrics/setup.go b/plugin/metrics/setup.go
index b50960211..d72ba7ec1 100644
--- a/plugin/metrics/setup.go
+++ b/plugin/metrics/setup.go
@@ -27,57 +27,54 @@ func init() {
 }
 
 func setup(c *caddy.Controller) error {
-	m, err := prometheusParse(c)
+	m, err := parse(c)
 	if err != nil {
 		return plugin.Error("prometheus", err)
 	}
 
-	// register the metrics to its address (ensure only one active metrics per address)
-	obj := uniqAddr.Set(m.Addr, m.OnStartup, m)
-	//propagate the real active Registry to current metrics
-	if om, ok := obj.(*Metrics); ok {
-		m.Reg = om.Reg
-	}
-
-	dnsserver.GetConfig(c).AddPlugin(func(next plugin.Handler) plugin.Handler {
-		m.Next = next
-		return m
-	})
+	c.OnStartup(func() error { m.Reg = uniqAddr.Set(m.Addr, m.OnStartup, m).(*Metrics).Reg; return nil })
+	c.OnRestartFailed(func() error { m.Reg = uniqAddr.Set(m.Addr, m.OnStartup, m).(*Metrics).Reg; return nil })
 
-	c.OncePerServerBlock(func() error {
-		c.OnStartup(func() error {
-			return uniqAddr.ForEach()
-		})
-		return nil
-	})
+	c.OnStartup(func() error { return uniqAddr.ForEach() })
+	c.OnRestartFailed(func() error { return uniqAddr.ForEach() })
 
-	c.OnRestart(func() error {
-		vars.PluginEnabled.Reset()
+	c.OnStartup(func() error {
+		conf := dnsserver.GetConfig(c)
+		for _, h := range conf.ListenHosts {
+			addrstr := conf.Transport + "://" + net.JoinHostPort(h, conf.Port)
+			for _, p := range conf.Handlers() {
+				vars.PluginEnabled.WithLabelValues(addrstr, conf.Zone, p.Name()).Set(1)
+			}
+		}
 		return nil
 	})
-
-	c.OnStartup(func() error {
+	c.OnRestartFailed(func() error {
 		conf := dnsserver.GetConfig(c)
-		plugins := conf.Handlers()
 		for _, h := range conf.ListenHosts {
 			addrstr := conf.Transport + "://" + net.JoinHostPort(h, conf.Port)
-			for _, p := range plugins {
+			for _, p := range conf.Handlers() {
 				vars.PluginEnabled.WithLabelValues(addrstr, conf.Zone, p.Name()).Set(1)
 			}
 		}
 		return nil
-
 	})
+
 	c.OnRestart(m.OnRestart)
+	c.OnRestart(func() error { vars.PluginEnabled.Reset(); return nil })
 	c.OnFinalShutdown(m.OnFinalShutdown)
 
 	// Initialize metrics.
 	buildInfo.WithLabelValues(coremain.CoreVersion, coremain.GitCommit, runtime.Version()).Set(1)
 
+	dnsserver.GetConfig(c).AddPlugin(func(next plugin.Handler) plugin.Handler {
+		m.Next = next
+		return m
+	})
+
 	return nil
 }
 
-func prometheusParse(c *caddy.Controller) (*Metrics, error) {
+func parse(c *caddy.Controller) (*Metrics, error) {
 	var met = New(defaultAddr)
 
 	i := 0
diff --git a/plugin/metrics/setup_test.go b/plugin/metrics/setup_test.go
index 73555427e..43949fdd3 100644
--- a/plugin/metrics/setup_test.go
+++ b/plugin/metrics/setup_test.go
@@ -22,7 +22,7 @@ func TestPrometheusParse(t *testing.T) {
 	}
 	for i, test := range tests {
 		c := caddy.NewTestController("dns", test.input)
-		m, err := prometheusParse(c)
+		m, err := parse(c)
 		if test.shouldErr && err == nil {
 			t.Errorf("Test %v: Expected error but found nil", i)
 			continue