Setup local cache with TTL for queries

Author: oz123

go.mod

@@ -5,6 +5,7 @@ go 1.14
 require (
 	github.com/caddyserver/caddy v1.0.5
 	github.com/coredns/coredns v1.7.0
+	github.com/imkira/go-ttlmap v2.0.0+incompatible
 	github.com/miekg/dns v1.1.29
 	github.com/prometheus/client_golang v1.7.1
 	github.com/stretchr/testify v1.5.1

go.sum

@@ -235,6 +235,9 @@ github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:
 github.com/iij/doapi v0.0.0-20190504054126-0bbf12d6d7df/go.mod h1:QMZY7/J/KSQEhKWFeDesPjMj+wCHReeknARU3wqlyN4=
 github.com/imdario/mergo v0.3.5/go.mod h1:2EnlNZ0deacrJVfApfmtdGgDfMuh/nq6Ok1EcJh5FfA=
 github.com/imdario/mergo v0.3.9/go.mod h1:2EnlNZ0deacrJVfApfmtdGgDfMuh/nq6Ok1EcJh5FfA=
+github.com/imkira/go-ttlmap v1.0.1 h1:6inPHYxX7xVBaxEicVlfrlWHZfBKz1Rf2Y3ah+h67BA=
+github.com/imkira/go-ttlmap v2.0.0+incompatible h1:U0HKc010vF8qEQY4pvmZ7QLRSOpV4i87DKRngNHsqb0=
+github.com/imkira/go-ttlmap v2.0.0+incompatible/go.mod h1:O+nQkEb7vlapbVn5HZ3m4FItjV+qNG3GY9W/sb3l7aw=
 github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
 github.com/infobloxopen/go-trees v0.0.0-20190313150506-2af4e13f9062/go.mod h1:PcNJqIlcX/dj3DTG/+QQnRvSgTMG6CLpRMjWcv4+J6w=
 github.com/jimstudt/http-authentication v0.0.0-20140401203705-3eca13d6893a/go.mod h1:wK6yTYYcgjHE1Z1QtXACPDjcFJyBskHEdagmnq3vsP8=

netbox.go

@@ -45,7 +45,7 @@ func (n Netbox) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg)
 	answers := []dns.RR{}
 	state := request.Request{W: w, Req: r}
 
-	ip_address := query(n.Url, n.Token, strings.TrimRight(state.QName(), "."))
+	ip_address := query(n.Url, n.Token, strings.TrimRight(state.QName(), "."), n.CacheDuration)
 	// no IP is found in netbox pass processing to the next plugin
 	if len(ip_address) == 0 {
 		return plugin.NextOrFailure(n.Name(), n.Next, ctx, w, r)

query.go

@@ -22,6 +22,7 @@ import (
 	"time"
 
 	clog "github.com/coredns/coredns/plugin/pkg/log"
+	"github.com/imkira/go-ttlmap"
 )
 
 type Record struct {
@@ -33,52 +34,53 @@ type RecordsList struct {
 	Records []Record `json:"results"`
 }
 
-var stupidCache = make(map[string]string)
-
-func query(url, token, dns_name string) string {
-
-	clog.Debug(stupidCache)
-
-	records := RecordsList{}
-	client := &http.Client{}
-	var resp *http.Response
-	clog.Debug("Querying ", fmt.Sprintf("%s/?dns_name=%s", url, dns_name))
-	req, err := http.NewRequest("GET", fmt.Sprintf("%s/?dns_name=%s", url, dns_name), nil)
-	req.Header.Set("Authorization", fmt.Sprintf("Token %s", token))
-
-	for i := 1; i <= 10; i++ {
-		resp, err = client.Do(req)
-
+var localCache = ttlmap.New(nil)
+
+func query(url, token, dns_name string, duration time.Duration) string {
+	item, err := localCache.Get(dns_name)
+	if err == nil {
+		clog.Debug(fmt.Sprintf("Found in local cache %s", dns_name))
+		return item.Value().(string)
+	} else {
+		records := RecordsList{}
+		client := &http.Client{}
+		var resp *http.Response
+		clog.Debug("Querying ", fmt.Sprintf("%s/?dns_name=%s", url, dns_name))
+		req, err := http.NewRequest("GET", fmt.Sprintf("%s/?dns_name=%s", url, dns_name), nil)
+		req.Header.Set("Authorization", fmt.Sprintf("Token %s", token))
+
+		for i := 1; i <= 10; i++ {
+			resp, err = client.Do(req)
+
+			if err != nil {
+				clog.Fatalf("HTTP Error %v", err)
+			}
+
+			if resp.StatusCode == http.StatusOK {
+				break
+			}
+
+			time.Sleep(1 * time.Second)
+		}
+		// TODO: check that we got status code 200
+		body, err := ioutil.ReadAll(resp.Body)
 		if err != nil {
-			clog.Fatalf("HTTP Error %v", err)
+			clog.Fatalf("Error reading body %v", err)
 		}
 
-		if resp.StatusCode == http.StatusOK {
-			break
+		jsonAns := string(body)
+		err = json.Unmarshal([]byte(jsonAns), &records)
+		if err != nil {
+			clog.Fatalf("could not unmarshal response %v", err)
 		}
 
-		time.Sleep(1 * time.Second)
-	}
-	// TODO: check that we got status code 200
-	body, err := ioutil.ReadAll(resp.Body)
-	if err != nil {
-		clog.Fatalf("Error reading body %v", err)
-	}
-
-	jsonAns := string(body)
-	err = json.Unmarshal([]byte(jsonAns), &records)
-	if err != nil {
-		clog.Fatalf("could not unmarshal response %v", err)
-	}
+		if len(records.Records) == 0 {
+			clog.Info("Recored not found in", jsonAns)
+			return ""
+		}
 
-	if len(records.Records) == 0 {
-		clog.Info("Recored not found in", jsonAns)
-		return ""
+		ip_address := strings.Split(records.Records[0].Address, "/")[0]
+		localCache.Set(dns_name, ttlmap.NewItem(ip_address, ttlmap.WithTTL(duration)), nil)
+		return ip_address
 	}
-
-	ip_address := strings.Split(records.Records[0].Address, "/")[0]
-
-	stupidCache[dns_name] = ip_address
-
-	return ip_address
 }

query_test.go

@@ -15,8 +15,10 @@
 package netbox
 
 import (
+	"github.com/stretchr/testify/assert"
 	"gopkg.in/h2non/gock.v1"
 	"testing"
+	"time"
 )
 
 func TestQuery(t *testing.T) {
@@ -27,7 +29,7 @@ func TestQuery(t *testing.T) {
 		`{"count":1, "results":[{"address": "10.0.0.2/25", "dns_name": "my_host"}]}`)
 
 	want := "10.0.0.2"
-	got := query("https://example.org/api/ipam/ip-addresses", "mytoken", "my_host")
+	got := query("https://example.org/api/ipam/ip-addresses", "mytoken", "my_host", time.Millisecond*100)
 	if got != want {
 		t.Fatalf("Expected %s but got %s", want, got)
 	}
@@ -42,9 +44,44 @@ func TestNoSuchHost(t *testing.T) {
 		200).BodyString(`{"count":0,"next":null,"previous":null,"results":[]}`)
 
 	want := ""
-	got := query("https://example.org/api/ipam/ip-addresses", "mytoken", "NoSuchHost")
+	got := query("https://example.org/api/ipam/ip-addresses", "mytoken", "NoSuchHost", time.Millisecond*100)
 	if got != want {
 		t.Fatalf("Expected empty string but got %s", got)
 	}
 
 }
+
+func TestLocalCache(t *testing.T) {
+	defer gock.Off() // Flush pending mocks after test execution
+	gock.New("https://example.org/api/ipam/ip-addresses/").MatchParams(
+		map[string]string{"dns_name": "my_host"}).Reply(
+		200).BodyString(
+		`{"count":1, "results":[{"address": "10.0.0.2/25", "dns_name": "my_host"}]}`)
+
+	ip_address := ""
+
+	got := query("https://example.org/api/ipam/ip-addresses", "mytoken", "my_host", time.Millisecond*100)
+
+	item, err := localCache.Get("my_host")
+	if err == nil {
+		ip_address = item.Value().(string)
+	}
+
+	assert.Equal(t, got, ip_address, "local cache item didn't match")
+
+}
+
+func TestLocalCacheExpiration(t *testing.T) {
+	defer gock.Off() // Flush pending mocks after test execution
+	gock.New("https://example.org/api/ipam/ip-addresses/").MatchParams(
+		map[string]string{"dns_name": "my_host"}).Reply(
+		200).BodyString(
+		`{"count":1, "results":[{"address": "10.0.0.2/25", "dns_name": "my_host"}]}`)
+
+	query("https://example.org/api/ipam/ip-addresses", "mytoken", "my_host", time.Millisecond*100)
+	<-time.After(101 * time.Millisecond)
+	item, err := localCache.Get("my_host")
+	if err != nil {
+		t.Fatalf("Expected errors, but got: %v", item)
+	}
+}