fix: should waiting for the ephemeral thread done (#326)

Co-authored-by: anrs <anders.hu@shopee.com>

Author: anrs

store/etcdv3/meta/ephemeral.go

@@ -2,6 +2,7 @@ package meta
 
 import (
 	"context"
+	"sync"
 	"time"
 
 	"go.etcd.io/etcd/v3/clientv3"
@@ -30,7 +31,10 @@ func (e *ETCD) StartEphemeral(ctx context.Context, path string, heartbeat time.D
 	ctx, cancel := context.WithCancel(context.Background())
 	expiry := make(chan struct{})
 
+	var wg sync.WaitGroup
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		defer close(expiry)
 
 		tick := time.NewTicker(heartbeat / 3)
@@ -58,5 +62,8 @@ func (e *ETCD) StartEphemeral(ctx context.Context, path string, heartbeat time.D
 		}
 	}()
 
-	return expiry, func() { cancel() }, nil
+	return expiry, func() {
+		cancel()
+		wg.Wait()
+	}, nil
 }

store/etcdv3/meta/ephemeral_test.go

@@ -8,6 +8,28 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestEphemeralDeregister(t *testing.T) {
+	m := NewEmbeddedETCD(t)
+	defer m.TerminateEmbededStorage()
+
+	ctx := context.Background()
+	path := "/ident"
+	heartbeat := time.Millisecond
+	expiry, stop, err := m.StartEphemeral(ctx, path, heartbeat)
+	require.NoError(t, err)
+	require.NotNil(t, stop)
+	require.NotNil(t, expiry)
+
+	kv, err := m.GetOne(ctx, path)
+	require.NoError(t, err)
+	require.Equal(t, path, string(kv.Key))
+
+	stop()
+	kv, err = m.GetOne(ctx, path)
+	require.Error(t, err)
+	require.Nil(t, kv)
+}
+
 func TestEphemeral(t *testing.T) {
 	m := NewEmbeddedETCD(t)
 	defer m.TerminateEmbededStorage()

store/etcdv3/service_test.go

@@ -2,13 +2,35 @@ package etcdv3
 
 import (
 	"context"
+	"fmt"
 	"sort"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 )
 
+func TestRegisterServiceWithDeregister(t *testing.T) {
+	m := NewMercury(t)
+	defer m.TerminateEmbededStorage()
+
+	ctx := context.Background()
+	svc := "svc"
+	path := fmt.Sprintf(serviceStatusKey, svc)
+	_, deregister, err := m.RegisterService(ctx, svc, time.Minute)
+	assert.NoError(t, err)
+
+	kv, err := m.GetOne(ctx, path)
+	assert.NoError(t, err)
+	assert.Equal(t, path, string(kv.Key))
+
+	deregister()
+	//time.Sleep(time.Second)
+	kv, err = m.GetOne(ctx, path)
+	assert.Error(t, err)
+	assert.Nil(t, kv)
+}
+
 func TestServiceStatusStream(t *testing.T) {
 	m := NewMercury(t)
 	defer m.TerminateEmbededStorage()