store, cluster: add unittests for service

Author: jschwinger233

cluster/calcium/service.go

@@ -109,11 +109,11 @@ func (c *Calcium) RegisterService(ctx context.Context) (unregister func(), err e
 	done := make(chan struct{})
 	ctx, cancel := context.WithCancel(ctx)
 	go func() {
+		defer close(done)
 		defer func() {
 			if err := c.store.UnregisterService(context.Background(), serviceAddress); err != nil {
 				log.Errorf("[RegisterService] failed to unregister service: %v", err)
 			}
-			close(done)
 		}()
 
 		timer := time.NewTicker(c.config.GRPCConfig.ServiceHeartbeatInterval / 2)

cluster/calcium/service_test.go

@@ -0,0 +1,101 @@
+package calcium
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+
+	storemocks "github.com/projecteru2/core/store/mocks"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+)
+
+func TestServiceStatusStream(t *testing.T) {
+	c := NewTestCluster()
+	c.config.Bind = ":5001"
+	c.config.GRPCConfig.ServiceHeartbeatInterval = 100 * time.Millisecond
+	c.config.GRPCConfig.ServiceDiscoveryPushInterval = 10 * time.Second
+	store := &storemocks.Store{}
+	c.store = store
+
+	registered := map[string]int{}
+	store.On("RegisterService", mock.AnythingOfType("*context.cancelCtx"), mock.AnythingOfType("string"), mock.AnythingOfType("time.Duration")).Return(
+		func(_ context.Context, addr string, _ time.Duration) error {
+			if v, ok := registered[addr]; ok {
+				registered[addr] = v + 1
+			} else {
+				registered[addr] = 1
+			}
+			return nil
+		},
+	)
+	store.On("UnregisterService", mock.AnythingOfType("*context.emptyCtx"), mock.AnythingOfType("string")).Return(
+		func(_ context.Context, addr string) error {
+			delete(registered, addr)
+			return nil
+		},
+	)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	unregister, err := c.RegisterService(ctx)
+	assert.NoError(t, err)
+	assert.Equal(t, len(registered), 1)
+	for _, v := range registered {
+		assert.Equal(t, v, 1)
+	}
+	time.Sleep(100 * time.Millisecond)
+	for _, v := range registered {
+		assert.Equal(t, v, 2)
+	}
+	unregister()
+	assert.Equal(t, len(registered), 0)
+}
+
+func TestWatchServiceStatus(t *testing.T) {
+	c := NewTestCluster()
+	c.config.GRPCConfig.ServiceDiscoveryPushInterval = 500 * time.Millisecond
+	store := &storemocks.Store{}
+	c.store = store
+	c.watcher = &serviceWatcher{}
+
+	store.On("ServiceStatusStream", mock.AnythingOfType("*context.emptyCtx")).Return(
+		func(_ context.Context) chan []string {
+			ch := make(chan []string)
+			go func() {
+				ticker := time.NewTicker(50 * time.Millisecond)
+				cnt := 0
+				for range ticker.C {
+					if cnt == 2 {
+						break
+					}
+					ch <- []string{fmt.Sprintf("127.0.0.1:500%d", cnt)}
+					cnt++
+				}
+			}()
+			return ch
+		}, nil,
+	)
+
+	ch, err := c.WatchServiceStatus(context.Background())
+	assert.NoError(t, err)
+	ch2, err := c.WatchServiceStatus(context.Background())
+	assert.NoError(t, err)
+	wg := sync.WaitGroup{}
+	wg.Add(2)
+	go func() {
+		defer wg.Done()
+		assert.Equal(t, (<-ch).Addresses, []string{"127.0.0.1:5000"})
+		assert.Equal(t, (<-ch).Addresses, []string{"127.0.0.1:5001"})
+		assert.Equal(t, (<-ch).Addresses, []string{"127.0.0.1:5001"})
+	}()
+	go func() {
+		defer wg.Done()
+		assert.Equal(t, (<-ch2).Addresses, []string{"127.0.0.1:5000"})
+		assert.Equal(t, (<-ch2).Addresses, []string{"127.0.0.1:5001"})
+		assert.Equal(t, (<-ch2).Addresses, []string{"127.0.0.1:5001"})
+	}()
+	wg.Wait()
+}

core.yaml.sample

@@ -1,6 +1,5 @@
 log_level: "DEBUG"
 bind: ":5001"
-service_address: 10.22.12.87:5001
 statsd: "127.0.0.1:8125"
 profile: ":12346"
 global_timeout: 300s
@@ -14,8 +13,8 @@ auth:
 grpc:
     max_concurrent_streams: 100
     max_recv_msg_size: 30 # will covert to MBytes
-    service_discovery_interval: 10s # WatchServiceStatus push interval
-    service_heartbeat_interval: 10s # RegisterService heartbeat
+    service_discovery_interval: 5s  # WatchServiceStatus push interval
+    service_heartbeat_interval: 5s  # RegisterService heartbeat
 
 etcd:
     machines:

store/etcdv3/mercury.go

@@ -22,8 +22,8 @@ import (
 )
 
 const (
-	podInfoKey          = "/pod/info/%s" // /pod/info/{podname}
-	serviceStatusPrefix = "/services/"
+	podInfoKey       = "/pod/info/%s" // /pod/info/{podname}
+	serviceStatusKey = "/services/%s" // /service/{ipv4:port}
 
 	nodeInfoKey       = "/node/%s"               // /node/{nodename}
 	nodePodKey        = "/node/%s:pod/%s"        // /node/{podname}:pod/{nodename}

store/etcdv3/service.go

@@ -2,6 +2,7 @@ package etcdv3
 
 import (
 	"context"
+	"fmt"
 	"strings"
 	"time"
 
@@ -41,7 +42,7 @@ func (m *Mercury) ServiceStatusStream(ctx context.Context) (chan []string, error
 	go func() {
 		defer close(ch)
 		log.Info("[ServiceStatusStream] start watching service status")
-		resp, err := m.Get(ctx, serviceStatusPrefix, clientv3.WithPrefix())
+		resp, err := m.Get(ctx, fmt.Sprintf(serviceStatusKey, ""), clientv3.WithPrefix())
 		if err != nil {
 			log.Errorf("[ServiceStatusStream] failed to get current services: %v", err)
 			return
@@ -52,7 +53,7 @@ func (m *Mercury) ServiceStatusStream(ctx context.Context) (chan []string, error
 		}
 		ch <- eps.ToSlice()
 
-		for resp := range m.watch(ctx, serviceStatusPrefix, clientv3.WithPrefix()) {
+		for resp := range m.watch(ctx, fmt.Sprintf(serviceStatusKey, ""), clientv3.WithPrefix()) {
 			if resp.Err() != nil {
 				if !resp.Canceled {
 					log.Errorf("[ServiceStatusStream] watch failed %v", resp.Err())
@@ -84,7 +85,7 @@ func (m *Mercury) ServiceStatusStream(ctx context.Context) (chan []string, error
 
 // RegisterService put /services/{address}
 func (m *Mercury) RegisterService(ctx context.Context, serviceAddress string, expire time.Duration) error {
-	key := serviceStatusPrefix + serviceAddress
+	key := fmt.Sprintf(serviceStatusKey, serviceAddress)
 	lease, err := m.cliv3.Grant(ctx, int64(expire/time.Second))
 	if err != nil {
 		return err
@@ -96,7 +97,7 @@ func (m *Mercury) RegisterService(ctx context.Context, serviceAddress string, ex
 
 // UnregisterService del /services/{address}
 func (m *Mercury) UnregisterService(ctx context.Context, serviceAddress string) error {
-	key := serviceStatusPrefix + serviceAddress
+	key := fmt.Sprintf(serviceStatusKey, serviceAddress)
 	_, err := m.Delete(ctx, key)
 	return err
 }

store/etcdv3/service_test.go

@@ -0,0 +1,51 @@
+package etcdv3
+
+import (
+	"context"
+	"sort"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRegisterService(t *testing.T) {
+	m := NewMercury(t)
+	defer m.TerminateEmbededStorage()
+	ctx := context.Background()
+	expire := 100 * time.Millisecond
+	err := m.RegisterService(ctx, "127.0.0.1:5002", expire)
+	assert.NoError(t, err)
+	kv, err := m.GetOne(ctx, "/services/127.0.0.1:5002")
+	assert.NoError(t, err)
+	assert.True(t, strings.HasSuffix(string(kv.Key), "127.0.0.1:5002"))
+}
+
+func TestUnregisterService(t *testing.T) {
+	m := NewMercury(t)
+	defer m.TerminateEmbededStorage()
+	ctx := context.Background()
+	addr := "127.0.0.1:5002"
+	assert.NoError(t, m.RegisterService(ctx, addr, time.Second))
+	assert.NoError(t, m.UnregisterService(ctx, addr))
+	_, err := m.GetOne(ctx, "/services/127.0.0.1:5002")
+	assert.Error(t, err, "bad `Count` value")
+}
+
+func TestServiceStatusStream(t *testing.T) {
+	m := NewMercury(t)
+	defer m.TerminateEmbededStorage()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	assert.NoError(t, m.RegisterService(ctx, "127.0.0.1:5001", time.Second))
+	ch, err := m.ServiceStatusStream(ctx)
+	assert.NoError(t, err)
+	assert.Equal(t, <-ch, []string{"127.0.0.1:5001"})
+	assert.NoError(t, m.RegisterService(ctx, "127.0.0.1:5002", time.Second))
+	endpoints := <-ch
+	sort.Strings(endpoints)
+	assert.Equal(t, endpoints, []string{"127.0.0.1:5001", "127.0.0.1:5002"})
+	assert.NoError(t, m.UnregisterService(ctx, "127.0.0.1:5001"))
+	assert.Equal(t, <-ch, []string{"127.0.0.1:5002"})
+}