Move treecache to its own package, add unit test

pkg/dns/dns_test.go

@@ -35,6 +35,7 @@ import (
 	"k8s.io/kubernetes/pkg/api/unversioned"
 	"k8s.io/kubernetes/pkg/client/cache"
 	fake "k8s.io/kubernetes/pkg/client/clientset_generated/internalclientset/fake"
+	"k8s.io/kubernetes/pkg/dns/treecache"
 	"k8s.io/kubernetes/pkg/dns/util"
 	"k8s.io/kubernetes/pkg/util/sets"
 )
@@ -51,7 +52,7 @@ func newKubeDNS() *KubeDNS {
 		domain:              testDomain,
 		endpointsStore:      cache.NewStore(cache.MetaNamespaceKeyFunc),
 		servicesStore:       cache.NewStore(cache.MetaNamespaceKeyFunc),
-		cache:               NewTreeCache(),
+		cache:               treecache.NewTreeCache(),
 		reverseRecordMap:    make(map[string]*skymsg.Service),
 		clusterIPServiceMap: make(map[string]*kapi.Service),
 		cacheLock:           sync.RWMutex{},

pkg/dns/treecache/BUILD

+package(default_visibility = ["//visibility:public"])
+
+licenses(["notice"])
+
+load(
+    "@io_bazel_rules_go//go:def.bzl",
+    "go_binary",
+    "go_library",
+    "go_test",
+    "cgo_library",
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = ["treecache.go"],
+    tags = ["automanaged"],
+    deps = ["//vendor:github.com/skynetservices/skydns/msg"],
+)
+
+go_test(
+    name = "go_default_test",
+    srcs = ["treecache_test.go"],
+    library = "go_default_library",
+    tags = ["automanaged"],
+    deps = ["//vendor:github.com/skynetservices/skydns/msg"],
+)

pkg/dns/treecache.go → pkg/dns/treecache/treecache.go

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package dns
+package treecache
 
 import (
 	"encoding/json"
@@ -23,19 +23,48 @@ import (
 	skymsg "github.com/skynetservices/skydns/msg"
 )
 
-type TreeCache struct {
-	ChildNodes map[string]*TreeCache
+type TreeCache interface {
+	// GetEntry with the given key for the given path.
+	GetEntry(key string, path ...string) (interface{}, bool)
+
+	// Get a list of values including wildcards labels (e.g. "*").
+	GetValuesForPathWithWildcards(path ...string) []*skymsg.Service
+
+	// SetEntry creates the entire path if it doesn't already exist in
+	// the cache, then sets the given service record under the given
+	// key. The path this entry would have occupied in an etcd datastore
+	// is computed from the given fqdn and stored as the "Key" of the
+	// skydns service; this is only required because skydns expects the
+	// service record to contain a key in a specific format (presumably
+	// for legacy compatibility). Note that the fqnd string typically
+	// contains both the key and all elements in the path.
+	SetEntry(key string, val *skymsg.Service, fqdn string, path ...string)
+
+	// SetSubCache inserts the given subtree under the given
+	// path:key. Usually the key is the name of a Kubernetes Service,
+	// and the path maps to the cluster subdomains matching the Service.
+	SetSubCache(key string, subCache TreeCache, path ...string)
+
+	// DeletePath removes all entries associated with a given path.
+	DeletePath(path ...string) bool
+
+	// Serialize dumps a JSON representation of the cache.
+	Serialize() (string, error)
+}
+
+type treeCache struct {
+	ChildNodes map[string]*treeCache
 	Entries    map[string]interface{}
 }
 
-func NewTreeCache() *TreeCache {
-	return &TreeCache{
-		ChildNodes: make(map[string]*TreeCache),
+func NewTreeCache() TreeCache {
+	return &treeCache{
+		ChildNodes: make(map[string]*treeCache),
 		Entries:    make(map[string]interface{}),
 	}
 }
 
-func (cache *TreeCache) Serialize() (string, error) {
+func (cache *treeCache) Serialize() (string, error) {
 	prettyJSON, err := json.MarshalIndent(cache, "", "\t")
 	if err != nil {
 		return "", err
@@ -43,14 +72,7 @@ func (cache *TreeCache) Serialize() (string, error) {
 	return string(prettyJSON), nil
 }
 
-// setEntry creates the entire path if it doesn't already exist in the cache,
-// then sets the given service record under the given key. The path this entry
-// would have occupied in an etcd datastore is computed from the given fqdn and
-// stored as the "Key" of the skydns service; this is only required because
-// skydns expects the service record to contain a key in a specific format
-// (presumably for legacy compatibility). Note that the fqnd string typically
-// contains both the key and all elements in the path.
-func (cache *TreeCache) setEntry(key string, val *skymsg.Service, fqdn string, path ...string) {
+func (cache *treeCache) SetEntry(key string, val *skymsg.Service, fqdn string, path ...string) {
 	// TODO: Consolidate setEntry and setSubCache into a single method with a
 	// type switch.
 	// TODO: Instead of passing the fqdn as an argument, we can reconstruct
@@ -70,7 +92,7 @@ func (cache *TreeCache) setEntry(key string, val *skymsg.Service, fqdn string, p
 	node.Entries[key] = val
 }
 
-func (cache *TreeCache) getSubCache(path ...string) *TreeCache {
+func (cache *treeCache) getSubCache(path ...string) *treeCache {
 	childCache := cache
 	for _, subpath := range path {
 		childCache = childCache.ChildNodes[subpath]
@@ -81,15 +103,12 @@ func (cache *TreeCache) getSubCache(path ...string) *TreeCache {
 	return childCache
 }
 
-// setSubCache inserts the given subtree under the given path:key. Usually the
-// key is the name of a Kubernetes Service, and the path maps to the cluster
-// subdomains matching the Service.
-func (cache *TreeCache) setSubCache(key string, subCache *TreeCache, path ...string) {
+func (cache *treeCache) SetSubCache(key string, subCache TreeCache, path ...string) {
 	node := cache.ensureChildNode(path...)
-	node.ChildNodes[key] = subCache
+	node.ChildNodes[key] = subCache.(*treeCache)
 }
 
-func (cache *TreeCache) getEntry(key string, path ...string) (interface{}, bool) {
+func (cache *treeCache) GetEntry(key string, path ...string) (interface{}, bool) {
 	childNode := cache.getSubCache(path...)
 	if childNode == nil {
 		return nil, false
@@ -98,11 +117,11 @@ func (cache *TreeCache) getEntry(key string, path ...string) (interface{}, bool)
 	return val, ok
 }
 
-func (cache *TreeCache) getValuesForPathWithWildcards(path ...string) []*skymsg.Service {
+func (cache *treeCache) GetValuesForPathWithWildcards(path ...string) []*skymsg.Service {
 	retval := []*skymsg.Service{}
-	nodesToExplore := []*TreeCache{cache}
+	nodesToExplore := []*treeCache{cache}
 	for idx, subpath := range path {
-		nextNodesToExplore := []*TreeCache{}
+		nextNodesToExplore := []*treeCache{}
 		if idx == len(path)-1 {
 			// if path ends on an entry, instead of a child node, add the entry
 			for _, node := range nodesToExplore {
@@ -150,7 +169,7 @@ func (cache *TreeCache) getValuesForPathWithWildcards(path ...string) []*skymsg.
 	return retval
 }
 
-func (cache *TreeCache) deletePath(path ...string) bool {
+func (cache *treeCache) DeletePath(path ...string) bool {
 	if len(path) == 0 {
 		return false
 	}
@@ -169,19 +188,7 @@ func (cache *TreeCache) deletePath(path ...string) bool {
 	return false
 }
 
-func (cache *TreeCache) deleteEntry(key string, path ...string) bool {
-	childNode := cache.getSubCache(path...)
-	if childNode == nil {
-		return false
-	}
-	if _, ok := childNode.Entries[key]; ok {
-		delete(childNode.Entries, key)
-		return true
-	}
-	return false
-}
-
-func (cache *TreeCache) appendValues(recursive bool, ref [][]interface{}) {
+func (cache *treeCache) appendValues(recursive bool, ref [][]interface{}) {
 	for _, value := range cache.Entries {
 		ref[0] = append(ref[0], value)
 	}
@@ -192,83 +199,15 @@ func (cache *TreeCache) appendValues(recursive bool, ref [][]interface{}) {
 	}
 }
 
-func (cache *TreeCache) ensureChildNode(path ...string) *TreeCache {
+func (cache *treeCache) ensureChildNode(path ...string) *treeCache {
 	childNode := cache
 	for _, subpath := range path {
 		newNode, ok := childNode.ChildNodes[subpath]
 		if !ok {
-			newNode = NewTreeCache()
+			newNode = NewTreeCache().(*treeCache)
 			childNode.ChildNodes[subpath] = newNode
 		}
 		childNode = newNode
 	}
 	return childNode
 }
-
-// unused function. keeping it around in commented-fashion
-// in the future, we might need some form of this function so that
-// we can serialize to a file in a mounted empty dir..
-//const (
-//	dataFile = "data.dat"
-//	crcFile  = "data.crc"
-//)
-//func (cache *TreeCache) Serialize(dir string) (string, error) {
-//	cache.m.RLock()
-//	defer cache.m.RUnlock()
-//	b, err := json.Marshal(cache)
-//	if err != nil {
-//		return "", err
-//	}
-//
-//	if err := ensureDir(dir, os.FileMode(0755)); err != nil {
-//		return "", err
-//	}
-//	if err := ioutil.WriteFile(path.Join(dir, dataFile), b, 0644); err != nil {
-//		return "", err
-//	}
-//	if err := ioutil.WriteFile(path.Join(dir, crcFile), getMD5(b), 0644); err != nil {
-//		return "", err
-//	}
-//	return string(b), nil
-//}
-
-//func ensureDir(path string, perm os.FileMode) error {
-//	s, err := os.Stat(path)
-//	if err != nil || !s.IsDir() {
-//		return os.Mkdir(path, perm)
-//	}
-//	return nil
-//}
-
-//func getMD5(b []byte) []byte {
-//	h := md5.New()
-//	h.Write(b)
-//	return []byte(fmt.Sprintf("%x", h.Sum(nil)))
-//}
-
-// unused function. keeping it around in commented-fashion
-// in the future, we might need some form of this function so that
-// we can restart kube-dns, deserialize the tree and have a cache
-// without having to wait for kube-dns to reach out to API server.
-//func Deserialize(dir string) (*TreeCache, error) {
-//	b, err := ioutil.ReadFile(path.Join(dir, dataFile))
-//	if err != nil {
-//		return nil, err
-//	}
-//
-//	hash, err := ioutil.ReadFile(path.Join(dir, crcFile))
-//	if err != nil {
-//		return nil, err
-//	}
-//	if !reflect.DeepEqual(hash, getMD5(b)) {
-//		return nil, fmt.Errorf("Checksum failed")
-//	}
-//
-//	var cache TreeCache
-//	err = json.Unmarshal(b, &cache)
-//	if err != nil {
-//		return nil, err
-//	}
-//	cache.m = &sync.RWMutex{}
-//	return &cache, nil
-//}

pkg/dns/treecache/treecache_test.go

+/*
+Copyright 2016 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package treecache
+
+import (
+	"testing"
+
+	"github.com/skynetservices/skydns/msg"
+)
+
+func TestTreeCache(t *testing.T) {
+	tc := NewTreeCache()
+
+	{
+		_, ok := tc.GetEntry("key1", "p1", "p2")
+		if ok {
+			t.Errorf("key should not exist")
+		}
+	}
+
+	checkExists := func(key string, expectedSvc *msg.Service, path ...string) {
+		svc, ok := tc.GetEntry(key, path...)
+		if !ok {
+			t.Fatalf("key %v should exist", key)
+		}
+		if svc := svc.(*msg.Service); svc != nil {
+			if svc != expectedSvc {
+				t.Errorf("value is not correct (%v != %v)", svc, expectedSvc)
+			}
+		} else {
+			t.Errorf("entry is not of the right type: %T", svc)
+		}
+	}
+	setEntryTC := []struct {
+		key  string
+		svc  *msg.Service
+		fqdn string
+		path []string
+	}{
+		{"key1", &msg.Service{}, "key1.p2.p1.", []string{"p1", "p2"}},
+		{"key2", &msg.Service{}, "key2.p2.p1.", []string{"p1", "p2"}},
+		{"key3", &msg.Service{}, "key3.p2.p1.", []string{"p1", "p3"}},
+	}
+
+	for _, testCase := range setEntryTC {
+		tc.SetEntry(testCase.key, testCase.svc, testCase.fqdn, testCase.path...)
+		checkExists(testCase.key, testCase.svc, testCase.path...)
+	}
+
+	wildcardTC := []struct {
+		path  []string
+		count int
+	}{
+		{[]string{"p1"}, 0},
+		{[]string{"p1", "p2"}, 2},
+		{[]string{"p1", "p3"}, 1},
+		{[]string{"p1", "p2", "key1"}, 1},
+		{[]string{"p1", "p2", "key2"}, 1},
+		{[]string{"p1", "p2", "key3"}, 0},
+		{[]string{"p1", "p3", "key3"}, 1},
+		{[]string{"p1", "p2", "*"}, 2},
+		{[]string{"p1", "*", "*"}, 3},
+	}
+
+	for _, testCase := range wildcardTC {
+		services := tc.GetValuesForPathWithWildcards(testCase.path...)
+		if len(services) != testCase.count {
+			t.Fatalf("Expected %v services for path %v, got %v",
+				testCase.count, testCase.path, len(services))
+		}
+	}
+
+	// Delete some paths
+	if !tc.DeletePath("p1", "p2") {
+		t.Fatal("should delete path p2.p1.")
+	}
+	if _, ok := tc.GetEntry("key3", "p1", "p3"); !ok {
+		t.Error("should not affect p3.p1.")
+	}
+	if tc.DeletePath("p1", "p2") {
+		t.Fatalf("should not be able to delete p2.p1")
+	}
+	if !tc.DeletePath("p1", "p3") {
+		t.Fatalf("should be able to delete p3.p1")
+	}
+	if tc.DeletePath("p1", "p3") {
+		t.Fatalf("should not be able to delete p3.t1")
+	}
+
+	for _, testCase := range []struct {
+		k string
+		p []string
+	}{
+		{"key1", []string{"p1", "p2"}},
+		{"key2", []string{"p1", "p2"}},
+		{"key3", []string{"p1", "p3"}},
+	} {
+		if _, ok := tc.GetEntry(testCase.k, testCase.p...); ok {
+			t.Error()
+		}
+	}
+}
+
+func TestTreeCacheSetSubCache(t *testing.T) {
+	tc := NewTreeCache()
+
+	m := &msg.Service{}
+
+	branch := NewTreeCache()
+	branch.SetEntry("key1", m, "key", "p2")
+
+	tc.SetSubCache("p1", branch, "p0")
+
+	if _, ok := tc.GetEntry("key1", "p0", "p1", "p2"); !ok {
+		t.Errorf("should be able to get entry p0.p1.p2.key1")
+	}
+}
+
+func TestTreeCacheSerialize(t *testing.T) {
+	tc := NewTreeCache()
+	tc.SetEntry("key1", &msg.Service{}, "key1.p2.p1.", "p1", "p2")
+
+	const expected = `{
+	"ChildNodes": {
+		"p1": {
+			"ChildNodes": {
+				"p2": {
+					"ChildNodes": {},
+					"Entries": {
+						"key1": {}
+					}
+				}
+			},
+			"Entries": {}
+		}
+	},
+	"Entries": {}
+}`
+
+	actual, err := tc.Serialize()
+	if err != nil {
+	}
+
+	if actual != expected {
+		t.Errorf("expected %q, got %q", expected, actual)
+	}
+}