From d12e25a8866eb189bc775bd1e791c586e6c5e943 Mon Sep 17 00:00:00 2001
From: kevin <wanjunfeng@gmail.com>
Date: Thu, 30 Jul 2020 22:56:39 +0800
Subject: [PATCH] add more tests

---
 core/discov/subscriber.go      |  48 ++++++------
 core/discov/subscriber_test.go | 132 +++++++++++++++++++++++++++++++++
 example/etcd/sub/sub.go        |   5 +-
 3 files changed, 159 insertions(+), 26 deletions(-)
 create mode 100644 core/discov/subscriber_test.go

diff --git a/core/discov/subscriber.go b/core/discov/subscriber.go
index d5a295fc..1e1bf508 100644
--- a/core/discov/subscriber.go
+++ b/core/discov/subscriber.go
@@ -2,8 +2,10 @@ package discov
 
 import (
 	"sync"
+	"sync/atomic"
 
 	"zero/core/discov/internal"
+	"zero/core/syncx"
 )
 
 type (
@@ -18,18 +20,20 @@ type (
 	}
 )
 
-func NewSubscriber(endpoints []string, key string, opts ...SubOption) *Subscriber {
+func NewSubscriber(endpoints []string, key string, opts ...SubOption) (*Subscriber, error) {
 	var subOpts subOptions
 	for _, opt := range opts {
 		opt(&subOpts)
 	}
 
-	subscriber := &Subscriber{
+	sub := &Subscriber{
 		items: newContainer(subOpts.exclusive),
 	}
-	internal.GetRegistry().Monitor(endpoints, key, subscriber.items)
+	if err := internal.GetRegistry().Monitor(endpoints, key, sub.items); err != nil {
+		return nil, err
+	}
 
-	return subscriber
+	return sub, nil
 }
 
 func (s *Subscriber) Values() []string {
@@ -48,6 +52,8 @@ type container struct {
 	exclusive bool
 	values    map[string][]string
 	mapping   map[string]string
+	snapshot  atomic.Value
+	dirty     *syncx.AtomicBool
 	lock      sync.Mutex
 }
 
@@ -56,6 +62,7 @@ func newContainer(exclusive bool) *container {
 		exclusive: exclusive,
 		values:    make(map[string][]string),
 		mapping:   make(map[string]string),
+		dirty:     syncx.ForAtomicBool(true),
 	}
 }
 
@@ -72,6 +79,7 @@ func (c *container) addKv(key, value string) ([]string, bool) {
 	c.lock.Lock()
 	defer c.lock.Unlock()
 
+	c.dirty.Set(true)
 	keys := c.values[value]
 	previous := append([]string(nil), keys...)
 	early := len(keys) > 0
@@ -114,14 +122,21 @@ func (c *container) doRemoveKey(key string) {
 }
 
 func (c *container) getValues() []string {
+	if !c.dirty.True() {
+		return c.snapshot.Load().([]string)
+	}
+
 	c.lock.Lock()
 	defer c.lock.Unlock()
 
-	var vs []string
+	var vals []string
 	for each := range c.values {
-		vs = append(vs, each)
+		vals = append(vals, each)
 	}
-	return vs
+	c.snapshot.Store(vals)
+	c.dirty.Set(false)
+
+	return vals
 }
 
 // removeKey removes the kv, returns true if there are still other keys associate with the value
@@ -129,23 +144,6 @@ func (c *container) removeKey(key string) {
 	c.lock.Lock()
 	defer c.lock.Unlock()
 
+	c.dirty.Set(true)
 	c.doRemoveKey(key)
 }
-
-func (c *container) removeVal(val string) (empty bool) {
-	c.lock.Lock()
-	defer c.lock.Unlock()
-
-	for k := range c.values {
-		if k == val {
-			delete(c.values, k)
-		}
-	}
-	for k, v := range c.mapping {
-		if v == val {
-			delete(c.mapping, k)
-		}
-	}
-
-	return len(c.values) == 0
-}
diff --git a/core/discov/subscriber_test.go b/core/discov/subscriber_test.go
new file mode 100644
index 00000000..64e3c2a2
--- /dev/null
+++ b/core/discov/subscriber_test.go
@@ -0,0 +1,132 @@
+package discov
+
+import (
+	"testing"
+
+	"zero/core/discov/internal"
+
+	"github.com/stretchr/testify/assert"
+)
+
+const (
+	actionAdd = iota
+	actionDel
+)
+
+func TestContainer(t *testing.T) {
+	type action struct {
+		act int
+		key string
+		val string
+	}
+	tests := []struct {
+		name   string
+		do     []action
+		expect []string
+	}{
+		{
+			name: "add one",
+			do: []action{
+				{
+					act: actionAdd,
+					key: "first",
+					val: "a",
+				},
+			},
+			expect: []string{
+				"a",
+			},
+		},
+		{
+			name: "add two",
+			do: []action{
+				{
+					act: actionAdd,
+					key: "first",
+					val: "a",
+				},
+				{
+					act: actionAdd,
+					key: "second",
+					val: "b",
+				},
+			},
+			expect: []string{
+				"a",
+				"b",
+			},
+		},
+		{
+			name: "add two, delete one",
+			do: []action{
+				{
+					act: actionAdd,
+					key: "first",
+					val: "a",
+				},
+				{
+					act: actionAdd,
+					key: "second",
+					val: "b",
+				},
+				{
+					act: actionDel,
+					key: "first",
+					val: "a",
+				},
+			},
+			expect: []string{
+				"b",
+			},
+		},
+		{
+			name: "add two, delete two",
+			do: []action{
+				{
+					act: actionAdd,
+					key: "first",
+					val: "a",
+				},
+				{
+					act: actionAdd,
+					key: "second",
+					val: "b",
+				},
+				{
+					act: actionDel,
+					key: "first",
+					val: "a",
+				},
+				{
+					act: actionDel,
+					key: "second",
+					val: "b",
+				},
+			},
+			expect: []string{},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			c := newContainer(false)
+			for _, order := range test.do {
+				if order.act == actionAdd {
+					c.OnAdd(internal.KV{
+						Key: order.key,
+						Val: order.val,
+					})
+				} else {
+					c.OnDelete(internal.KV{
+						Key: order.key,
+						Val: order.val,
+					})
+				}
+			}
+			assert.True(t, c.dirty.True())
+			assert.ElementsMatch(t, test.expect, c.getValues())
+			assert.False(t, c.dirty.True())
+			assert.ElementsMatch(t, test.expect, c.getValues())
+		})
+	}
+}
diff --git a/example/etcd/sub/sub.go b/example/etcd/sub/sub.go
index 13cbcd72..c6cd03c2 100644
--- a/example/etcd/sub/sub.go
+++ b/example/etcd/sub/sub.go
@@ -5,10 +5,13 @@ import (
 	"time"
 
 	"zero/core/discov"
+	"zero/core/lang"
 )
 
 func main() {
-	sub := discov.NewSubscriber([]string{"etcd.discovery:2379"}, "028F2C35852D", discov.Exclusive())
+	sub, err := discov.NewSubscriber([]string{"etcd.discovery:2379"}, "028F2C35852D", discov.Exclusive())
+	lang.Must(err)
+
 	ticker := time.NewTicker(time.Second * 3)
 	defer ticker.Stop()