6 files changed, 77 insertions, 75 deletions

diff --git a/plugin/metrics/addr.go b/plugin/metrics/addr.go
deleted file mode 100644
index fe8e5e5fe..000000000
--- a/plugin/metrics/addr.go
+++ /dev/null
@@ -1,52 +0,0 @@
-package metrics
-
-// addrs keeps track on which addrs we listen, so we only start one listener, is
-// prometheus is used in multiple Server Blocks.
-type addrs struct {
-	a map[string]value
-}
-
-type value struct {
-	state int
-	f     func() error
-}
-
-var uniqAddr addrs
-
-func newAddress() addrs {
-	return addrs{a: make(map[string]value)}
-}
-
-func (a addrs) setAddress(addr string, f func() error) {
-	if a.a[addr].state == done {
-		return
-	}
-	a.a[addr] = value{todo, f}
-}
-
-// setAddressTodo sets addr to 'todo' again.
-func (a addrs) setAddressTodo(addr string) {
-	v, ok := a.a[addr]
-	if !ok {
-		return
-	}
-	v.state = todo
-	a.a[addr] = v
-}
-
-// forEachTodo iterates for a and executes f for each element that is 'todo' and sets it to 'done'.
-func (a addrs) forEachTodo() error {
-	for k, v := range a.a {
-		if v.state == todo {
-			v.f()
-		}
-		v.state = done
-		a.a[k] = v
-	}
-	return nil
-}
-
-const (
-	todo = 1
-	done = 2
-)
diff --git a/plugin/metrics/addr_test.go b/plugin/metrics/addr_test.go
deleted file mode 100644
index d7a08656b..000000000
--- a/plugin/metrics/addr_test.go
+++ /dev/null
@@ -1,17 +0,0 @@
-package metrics
-
-import "testing"
-
-func TestForEachTodo(t *testing.T) {
-	a, i := newAddress(), 0
-	a.setAddress("test", func() error { i++; return nil })
-
-	a.forEachTodo()
-	if i != 1 {
-		t.Errorf("Failed to executed f for %s", "test")
-	}
-	a.forEachTodo()
-	if i != 1 {
-		t.Errorf("Executed f twice instead of once")
-	}
-}
diff --git a/plugin/metrics/metrics.go b/plugin/metrics/metrics.go
index e183d33d5..5816288a4 100644
--- a/plugin/metrics/metrics.go
+++ b/plugin/metrics/metrics.go
@@ -106,7 +106,7 @@ func (m *Metrics) OnRestart() error {
 		return nil
 	}
 
-	uniqAddr.setAddressTodo(m.Addr)
+	uniqAddr.SetTodo(m.Addr)
 
 	m.ln.Close()
 	m.lnSetup = false
diff --git a/plugin/metrics/setup.go b/plugin/metrics/setup.go
index 52d5775c1..c00f44a83 100644
--- a/plugin/metrics/setup.go
+++ b/plugin/metrics/setup.go
@@ -8,19 +8,21 @@ import (
 	"github.com/coredns/coredns/coremain"
 	"github.com/coredns/coredns/plugin"
 	clog "github.com/coredns/coredns/plugin/pkg/log"
+	"github.com/coredns/coredns/plugin/pkg/uniq"
 
 	"github.com/mholt/caddy"
 )
 
-var log = clog.NewWithPlugin("prometheus")
+var (
+	log      = clog.NewWithPlugin("prometheus")
+	uniqAddr = uniq.New()
+)
 
 func init() {
 	caddy.RegisterPlugin("prometheus", caddy.Plugin{
 		ServerType: "dns",
 		Action:     setup,
 	})
-
-	uniqAddr = newAddress()
 }
 
 func setup(c *caddy.Controller) error {
@@ -36,7 +38,7 @@ func setup(c *caddy.Controller) error {
 
 	c.OncePerServerBlock(func() error {
 		c.OnStartup(func() error {
-			return uniqAddr.forEachTodo()
+			return uniqAddr.ForEach()
 		})
 		return nil
 	})
@@ -54,7 +56,7 @@ func prometheusParse(c *caddy.Controller) (*Metrics, error) {
 	var met = New(defaultAddr)
 
 	defer func() {
-		uniqAddr.setAddress(met.Addr, met.OnStartup)
+		uniqAddr.Set(met.Addr, met.OnStartup)
 	}()
 
 	i := 0
diff --git a/plugin/pkg/uniq/uniq.go b/plugin/pkg/uniq/uniq.go
new file mode 100644
index 000000000..3e50d64b5
--- /dev/null
+++ b/plugin/pkg/uniq/uniq.go
@@ -0,0 +1,52 @@
+// Package uniq keeps track of "thing" that are either "todo" or "done". Multiple
+// identical events will only be processed once.
+package uniq
+
+// U keeps track of item to be done.
+type U struct {
+	u map[string]item
+}
+
+type item struct {
+	state int          // either todo or done
+	f     func() error // function to be executed.
+}
+
+// New returns a new initialized U.
+func New() U { return U{u: make(map[string]item)} }
+
+// Set sets function f in U under key. If the key already exists
+// it is not overwritten.
+func (u U) Set(key string, f func() error) {
+	if _, ok := u.u[key]; ok {
+		return
+	}
+	u.u[key] = item{todo, f}
+}
+
+// SetTodo sets key to 'todo' again.
+func (u U) SetTodo(key string) {
+	v, ok := u.u[key]
+	if !ok {
+		return
+	}
+	v.state = todo
+	u.u[key] = v
+}
+
+// ForEach iterates for u executes f for each element that is 'todo' and sets it to 'done'.
+func (u U) ForEach() error {
+	for k, v := range u.u {
+		if v.state == todo {
+			v.f()
+		}
+		v.state = done
+		u.u[k] = v
+	}
+	return nil
+}
+
+const (
+	todo = 1
+	done = 2
+)
diff --git a/plugin/pkg/uniq/uniq_test.go b/plugin/pkg/uniq/uniq_test.go
new file mode 100644
index 000000000..5d58c924b
--- /dev/null
+++ b/plugin/pkg/uniq/uniq_test.go
@@ -0,0 +1,17 @@
+package uniq
+
+import "testing"
+
+func TestForEach(t *testing.T) {
+	u, i := New(), 0
+	u.Set("test", func() error { i++; return nil })
+
+	u.ForEach()
+	if i != 1 {
+		t.Errorf("Failed to executed f for %s", "test")
+	}
+	u.ForEach()
+	if i != 1 {
+		t.Errorf("Executed f twice instead of once")
+	}
+}