Implements prefix lookups for DNS TTL

agent/catalog_endpoint_test.go

@@ -409,6 +409,7 @@ func TestCatalogServices(t *testing.T) {
 	t.Parallel()
 	a := NewTestAgent(t.Name(), "")
 	defer a.Shutdown()
+	testrpc.WaitForTestAgent(t, a.RPC, "dc1")
 
 	// Register node
 	args := &structs.RegisterRequest{

agent/consul/catalog_endpoint_test.go

@@ -1307,7 +1307,7 @@ func TestCatalog_ListServices_NodeMetaFilter(t *testing.T) {
 	codec := rpcClient(t, s1)
 	defer codec.Close()
 
-	testrpc.WaitForLeader(t, s1.RPC, "dc1")
+	testrpc.WaitForTestAgent(t, s1.RPC, "dc1")
 
 	// Add a new node with the right meta k/v pair
 	node := &structs.Node{Node: "foo", Address: "127.0.0.1", Meta: map[string]string{"somekey": "somevalue"}}

agent/consul/client_test.go

@@ -450,6 +450,7 @@ func TestClient_SnapshotRPC(t *testing.T) {
 
 	// Try to join.
 	joinLAN(t, c1, s1)
+	testrpc.WaitForLeader(t, c1.RPC, "dc1")
 
 	// Wait until we've got a healthy server.
 	retry.Run(t, func(r *retry.R) {

agent/dns.go

@@ -12,6 +12,7 @@ import (
 	"regexp"
 
 	"github.com/armon/go-metrics"
+	"github.com/armon/go-radix"
 	"github.com/coredns/coredns/plugin/pkg/dnsutil"
 	"github.com/hashicorp/consul/agent/config"
 	"github.com/hashicorp/consul/agent/consul"
@@ -64,6 +65,9 @@ type DNSServer struct {
 	domain    string
 	recursors []string
 	logger    *log.Logger
+	// Those are handling prefix lookups
+	ttlRadix  *radix.Tree
+	ttlStrict map[string]time.Duration
 
 	// disableCompression is the config.DisableCompression flag that can
 	// be safely changed at runtime. It always contains a bool and is
@@ -91,7 +95,21 @@ func NewDNSServer(a *Agent) (*DNSServer, error) {
 		domain:    domain,
 		logger:    a.logger,
 		recursors: recursors,
+		ttlRadix:  radix.New(),
+		ttlStrict: make(map[string]time.Duration),
+	}
+	if dnscfg.ServiceTTL != nil {
+		for key, ttl := range dnscfg.ServiceTTL {
+			// All suffix with '*' are put in radix
+			// This include '*' that will match anything
+			if strings.HasSuffix(key, "*") {
+				srv.ttlRadix.Insert(key[:len(key)-1], ttl)
+			} else {
+				srv.ttlStrict[key] = ttl
+			}
+		}
 	}
+
 	srv.disableCompression.Store(a.config.DNSDisableCompression)
 
 	return srv, nil
@@ -115,6 +133,22 @@ func GetDNSConfig(conf *config.RuntimeConfig) *dnsConfig {
 	}
 }
 
+// GetTTLForService Find the TTL for a given service.
+// return ttl, true if found, 0, false otherwise
+func (d *DNSServer) GetTTLForService(service string) (time.Duration, bool) {
+	if d.config.ServiceTTL != nil {
+		ttl, ok := d.ttlStrict[service]
+		if ok {
+			return ttl, true
+		}
+		_, ttlRaw, ok := d.ttlRadix.LongestPrefix(service)
+		if ok {
+			return ttlRaw.(time.Duration), true
+		}
+	}
+	return time.Duration(0), false
+}
+
 func (d *DNSServer) ListenAndServe(network, addr string, notif func()) error {
 	mux := dns.NewServeMux()
 	mux.HandleFunc("arpa.", d.handlePtr)
@@ -1013,14 +1047,7 @@ func (d *DNSServer) serviceLookup(network, datacenter, service, tag string, conn
 	out.Nodes.Shuffle()
 
 	// Determine the TTL
-	var ttl time.Duration
-	if d.config.ServiceTTL != nil {
-		var ok bool
-		ttl, ok = d.config.ServiceTTL[service]
-		if !ok {
-			ttl = d.config.ServiceTTL["*"]
-		}
-	}
+	ttl, _ := d.GetTTLForService(service)
 
 	// Add various responses depending on the request
 	qType := req.Question[0].Qtype
@@ -1141,11 +1168,7 @@ RPC:
 			d.logger.Printf("[WARN] dns: Failed to parse TTL '%s' for prepared query '%s', ignoring", out.DNS.TTL, query)
 		}
 	} else if d.config.ServiceTTL != nil {
-		var ok bool
-		ttl, ok = d.config.ServiceTTL[out.Service]
-		if !ok {
-			ttl = d.config.ServiceTTL["*"]
-		}
+		ttl, _ = d.GetTTLForService(out.Service)
 	}
 
 	// If we have no nodes, return not found!