Merge pull request #2795 from weaveworks/issues/2792-refcount-ipset

Reference-count ipset entries

Fix #2792

npc/ipset/ipset.go

@@ -27,10 +27,11 @@ type Interface interface {
 }
 
 type ipset struct {
+	refCount
 }
 
 func New() Interface {
-	return &ipset{}
+	return &ipset{refCount: newRefCount()}
 }
 
 func (i *ipset) Create(ipsetName Name, ipsetType Type) error {
@@ -38,26 +39,36 @@ func (i *ipset) Create(ipsetName Name, ipsetType Type) error {
 }
 
 func (i *ipset) AddEntry(ipsetName Name, entry string) error {
+	if i.inc(ipsetName, entry) > 1 { // already in the set
+		return nil
+	}
 	return doExec("add", string(ipsetName), entry)
 }
 
 func (i *ipset) DelEntry(ipsetName Name, entry string) error {
+	if i.dec(ipsetName, entry) > 0 { // still needed
+		return nil
+	}
 	return doExec("del", string(ipsetName), entry)
 }
 
 func (i *ipset) Flush(ipsetName Name) error {
+	i.removeSet(ipsetName)
 	return doExec("flush", string(ipsetName))
 }
 
 func (i *ipset) FlushAll() error {
+	i.refCount = newRefCount()
 	return doExec("flush")
 }
 
 func (i *ipset) Destroy(ipsetName Name) error {
+	i.removeSet(ipsetName)
 	return doExec("destroy", string(ipsetName))
 }
 
 func (i *ipset) DestroyAll() error {
+	i.refCount = newRefCount()
 	return doExec("destroy")
 }
 
@@ -67,3 +78,38 @@ func doExec(args ...string) error {
 	}
 	return nil
 }
+
+// Reference-counting
+type key struct {
+	ipsetName Name
+	entry     string
+}
+
+// note no locking is required as all operations are serialised in the controller
+type refCount struct {
+	ref map[key]int
+}
+
+func newRefCount() refCount {
+	return refCount{ref: make(map[key]int)}
+}
+
+func (rc *refCount) inc(ipsetName Name, entry string) int {
+	k := key{ipsetName, entry}
+	rc.ref[k]++
+	return rc.ref[k]
+}
+
+func (rc *refCount) dec(ipsetName Name, entry string) int {
+	k := key{ipsetName, entry}
+	rc.ref[k]--
+	return rc.ref[k]
+}
+
+func (rc *refCount) removeSet(ipsetName Name) {
+	for k := range rc.ref {
+		if k.ipsetName == ipsetName {
+			delete(rc.ref, k)
+		}
+	}
+}