Repository: go-kratos/kratos
Pull request #3511 (2025-01-02 03:16:04)

fix(registry/consul): fix concurrency issues and improve performance
修复了一些并发相关的问题：

1. `Registry` 锁的范围太大，比如 `Watch` 对整个函数加锁，不能同时处理多个（不同的）Service；

https://github.com/go-kratos/kratos/blob/5087366d2f90daf8a96bfcab292a433569d46302/contrib/registry/consul/registry.go#L183-L185

2. 相同的 Service 没有共享 `resolve`，只要 `Watch` 就会新开一个 `resolve` 协程；

https://github.com/go-kratos/kratos/blob/5087366d2f90daf8a96bfcab292a433569d46302/contrib/registry/consul/registry.go#L211-L213

3. `Client` 注销任一 Service 都会直接 `cancel` 掉 `Client.ctx`，导致所有的 `heartbeat` 协程退出；

https://github.com/go-kratos/kratos/blob/5087366d2f90daf8a96bfcab292a433569d46302/contrib/registry/consul/client.go#L246-L249

4. 在 `Watch` 时，由于 `w` 先被添加到 `set` 中，然后才会写入 `services event`，如果在两步中间正好遇到之前的 `resolve` 在 `broadcast`，209行的写入则会被卡住（`w` 还没有返回，不会有调用者主动调用 `Next` 消费）。

https://github.com/go-kratos/kratos/blob/5087366d2f90daf8a96bfcab292a433569d46302/contrib/registry/consul/registry.go#L196-L214

5. 其他，提前判断 `ctx.Err()` 避免做冗长的无用功、`sleepCtx` 等。


Patch:
diff --git a/contrib/registry/consul/client.go b/contrib/registry/consul/client.go
index ab1551cd..1954a243 100644
--- a/contrib/registry/consul/client.go
+++ b/contrib/registry/consul/client.go
@@ -9,6 +9,7 @@ import (
 	"net/url"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/go-kratos/kratos/v2/log"
@@ -26,10 +27,8 @@ const (
 
 // Client is consul client config
 type Client struct {
-	dc     Datacenter
-	cli    *api.Client
-	ctx    context.Context
-	cancel context.CancelFunc
+	dc  Datacenter
+	cli *api.Client
 
 	// resolve service entry endpoints
 	resolver ServiceResolver
@@ -41,6 +40,16 @@ type Client struct {
 	deregisterCriticalServiceAfter int
 	// serviceChecks  user custom checks
 	serviceChecks api.AgentServiceChecks
+
+	// used to control heartbeat
+	lock      sync.RWMutex
+	cancelers map[string]*canceler
+}
+
+type canceler struct {
+	ctx    context.Context
+	cancel context.CancelFunc
+	done   chan struct{}
 }
 
 func defaultResolver(_ context.Context, entries []*api.ServiceEntry) []*registry.ServiceInstance {
@@ -143,7 +152,7 @@ func (c *Client) singleDCEntries(service, tag string, passingOnly bool, opts *ap
 }
 
 // Register register service instance to consul
-func (c *Client) Register(_ context.Context, svc *registry.ServiceInstance, enableHealthCheck bool) error {
+func (c *Client) Register(ctx context.Context, svc *registry.ServiceInstance, enableHealthCheck bool) error {
 	addresses := make(map[string]api.ServiceAddress, len(svc.Endpoints))
 	checkAddresses := make([]string, 0, len(svc.Endpoints))
 	for _, endpoint := range svc.Endpoints {
@@ -190,13 +199,43 @@ func (c *Client) Register(_ context.Context, svc *registry.ServiceInstance, enab
 		})
 	}
 
-	err := c.cli.Agent().ServiceRegister(asr)
+	c.lock.Lock()
+	if cc, ok := c.cancelers[svc.ID]; ok {
+		cc.cancel()
+		<-cc.done
+	}
+	var cc *canceler
+	if c.heartbeat {
+		cancelCtx, cancel := context.WithCancel(context.Background())
+		cc = &canceler{
+			ctx:    cancelCtx,
+			cancel: cancel,
+			done:   make(chan struct{}),
+		}
+		c.cancelers[svc.ID] = cc
+		go func() {
+			<-cc.done
+			cc.cancel()
+			c.lock.Lock()
+			if c.cancelers[svc.ID] == cc {
+				delete(c.cancelers, svc.ID)
+			}
+			c.lock.Unlock()
+		}()
+	}
+	c.lock.Unlock()
+
+	err := c.cli.Agent().ServiceRegisterOpts(asr, api.ServiceRegisterOpts{}.WithContext(ctx))
 	if err != nil {
+		if c.heartbeat {
+			close(cc.done)
+		}
 		return err
 	}
+
 	if c.heartbeat {
 		go func() {
-			time.Sleep(time.Second)
+			defer close(cc.done)
 			err = c.cli.Agent().UpdateTTL("service:"+svc.ID, "pass", "pass")
 			if err != nil {
 				log.Errorf("[Consul]update ttl heartbeat to consul failed!err:=%v", err)
@@ -205,22 +244,11 @@ func (c *Client) Register(_ context.Context, svc *registry.ServiceInstance, enab
 			defer ticker.Stop()
 			for {
 				select {
-				case <-c.ctx.Done():
-					_ = c.cli.Agent().ServiceDeregister(svc.ID)
-					return
-				default:
-				}
-				select {
-				case <-c.ctx.Done():
+				case <-cc.ctx.Done():
 					_ = c.cli.Agent().ServiceDeregister(svc.ID)
 					return
 				case <-ticker.C:
-					// ensure that unregistered services will not be re-registered by mistake
-					if errors.Is(c.ctx.Err(), context.Canceled) || errors.Is(c.ctx.Err(), context.DeadlineExceeded) {
-						_ = c.cli.Agent().ServiceDeregister(svc.ID)
-						return
-					}
-					err = c.cli.Agent().UpdateTTLOpts("service:"+svc.ID, "pass", "pass", new(api.QueryOptions).WithContext(c.ctx))
+					err = c.cli.Agent().UpdateTTLOpts("service:"+svc.ID, "pass", "pass", new(api.QueryOptions).WithContext(cc.ctx))
 					if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
 						_ = c.cli.Agent().ServiceDeregister(svc.ID)
 						return
@@ -228,8 +256,11 @@ func (c *Client) Register(_ context.Context, svc *registry.ServiceInstance, enab
 					if err != nil {
 						log.Errorf("[Consul] update ttl heartbeat to consul failed! err=%v", err)
 						// when the previous report fails, try to re register the service
-						time.Sleep(time.Duration(rand.Intn(5)) * time.Second)
-						if err := c.cli.Agent().ServiceRegister(asr); err != nil {
+						if err := sleepCtx(cc.ctx, time.Duration(rand.Intn(5))*time.Second); err != nil {
+							_ = c.cli.Agent().ServiceDeregister(svc.ID)
+							return
+						}
+						if err := c.cli.Agent().ServiceRegisterOpts(asr, api.ServiceRegisterOpts{}.WithContext(cc.ctx)); err != nil {
 							log.Errorf("[Consul] re registry service failed!, err=%v", err)
 						} else {
 							log.Warn("[Consul] re registry of service occurred success")
@@ -242,8 +273,32 @@ func (c *Client) Register(_ context.Context, svc *registry.ServiceInstance, enab
 	return nil
 }
 
+func sleepCtx(ctx context.Context, d time.Duration) error {
+	t := time.NewTimer(d)
+	defer t.Stop()
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-t.C:
+		return nil
+	}
+}
+
 // Deregister service by service ID
-func (c *Client) Deregister(_ context.Context, serviceID string) error {
-	defer c.cancel()
-	return c.cli.Agent().ServiceDeregister(serviceID)
+func (c *Client) Deregister(ctx context.Context, serviceID string) error {
+	c.lock.RLock()
+	cc, ok := c.cancelers[serviceID]
+	c.lock.RUnlock()
+	if ok {
+		cc.cancel()
+		<-cc.done
+	}
+
+	err := c.cli.Agent().ServiceDeregisterOpts(serviceID, new(api.QueryOptions).WithContext(ctx))
+	var se api.StatusError
+	if errors.As(err, &se) && se.Code == 404 {
+		// not found
+		err = nil
+	}
+	return err
 }
diff --git a/contrib/registry/consul/registry.go b/contrib/registry/consul/registry.go
index 42fa4a37..da9a899f 100644
--- a/contrib/registry/consul/registry.go
+++ b/contrib/registry/consul/registry.go
@@ -113,12 +113,12 @@ func New(apiClient *api.Client, opts ...Option) *Registry {
 			healthcheckInterval:            10,
 			heartbeat:                      true,
 			deregisterCriticalServiceAfter: 600,
+			cancelers:                      make(map[string]*canceler),
 		},
 	}
 	for _, o := range opts {
 		o(r)
 	}
-	r.cli.ctx, r.cli.cancel = context.WithCancel(context.Background())
 	return r
 }
 
@@ -135,8 +135,8 @@ func (r *Registry) Deregister(ctx context.Context, svc *registry.ServiceInstance
 // GetService return service by name
 func (r *Registry) GetService(ctx context.Context, name string) ([]*registry.ServiceInstance, error) {
 	r.lock.RLock()
-	defer r.lock.RUnlock()
 	set := r.registry[name]
+	r.lock.RUnlock()
 
 	getRemote := func() []*registry.ServiceInstance {
 		services, _, err := r.cli.Service(ctx, name, 0, true)
@@ -181,17 +181,26 @@ func (r *Registry) ListServices() (allServices map[string][]*registry.ServiceIns
 
 // Watch resolve service by name
 func (r *Registry) Watch(ctx context.Context, name string) (registry.Watcher, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+
 	r.lock.Lock()
-	defer r.lock.Unlock()
 	set, ok := r.registry[name]
 	if !ok {
+		cancelCtx, cancel := context.WithCancel(context.Background())
 		set = &serviceSet{
+			registry:    r,
 			watcher:     make(map[*watcher]struct{}),
 			services:    &atomic.Value{},
 			serviceName: name,
+			ctx:         cancelCtx,
+			cancel:      cancel,
 		}
 		r.registry[name] = set
 	}
+	set.ref.Add(1)
+	r.lock.Unlock()
 
 	// init watcher
 	w := &watcher{
@@ -202,14 +211,21 @@ func (r *Registry) Watch(ctx context.Context, name string) (registry.Watcher, er
 	set.lock.Lock()
 	set.watcher[w] = struct{}{}
 	set.lock.Unlock()
+
 	ss, _ := set.services.Load().([]*registry.ServiceInstance)
 	if len(ss) > 0 {
 		// If the service has a value, it needs to be pushed to the watcher,
 		// otherwise the initial data may be blocked forever during the watch.
-		w.event <- struct{}{}
+		select {
+		case w.event <- struct{}{}:
+		default:
+		}
 	}
-	if err := r.resolve(ctx, set); err != nil {
-		return nil, err
+
+	if !ok {
+		if err := r.resolve(ctx, set); err != nil {
+			return nil, err
+		}
 	}
 	return w, nil
 }
@@ -239,9 +255,11 @@ func (r *Registry) resolve(ctx context.Context, ss *serviceSet) error {
 		for {
 			select {
 			case <-ticker.C:
-				tmpService, tmpIdx, err := listServices(context.Background(), ss.serviceName, idx, true)
+				tmpService, tmpIdx, err := listServices(ss.ctx, ss.serviceName, idx, true)
 				if err != nil {
-					time.Sleep(time.Second)
+					if err := sleepCtx(ss.ctx, time.Second); err != nil {
+						return
+					}
 					continue
 				}
 				if len(tmpService) != 0 && tmpIdx != idx {
@@ -249,7 +267,7 @@ func (r *Registry) resolve(ctx context.Context, ss *serviceSet) error {
 					ss.broadcast(services)
 				}
 				idx = tmpIdx
-			case <-ctx.Done():
+			case <-ss.ctx.Done():
 				return
 			}
 		}
@@ -257,3 +275,14 @@ func (r *Registry) resolve(ctx context.Context, ss *serviceSet) error {
 
 	return nil
 }
+
+func (r *Registry) tryDelete(ss *serviceSet) bool {
+	r.lock.Lock()
+	defer r.lock.Unlock()
+	if ss.ref.Add(-1) != 0 {
+		return false
+	}
+	ss.cancel()
+	delete(r.registry, ss.serviceName)
+	return true
+}
diff --git a/contrib/registry/consul/service.go b/contrib/registry/consul/service.go
index a0a16609..5d2ba1a0 100644
--- a/contrib/registry/consul/service.go
+++ b/contrib/registry/consul/service.go
@@ -1,6 +1,7 @@
 package consul
 
 import (
+	"context"
 	"sync"
 	"sync/atomic"
 
@@ -8,10 +9,16 @@ import (
 )
 
 type serviceSet struct {
+	registry    *Registry
 	serviceName string
 	watcher     map[*watcher]struct{}
+	ref         atomic.Int32
 	services    *atomic.Value
 	lock        sync.RWMutex
+
+	// for cancel
+	ctx    context.Context
+	cancel context.CancelFunc
 }
 
 func (s *serviceSet) broadcast(ss []*registry.ServiceInstance) {
@@ -25,3 +32,10 @@ func (s *serviceSet) broadcast(ss []*registry.ServiceInstance) {
 		}
 	}
 }
+
+func (s *serviceSet) delete(w *watcher) {
+	s.lock.Lock()
+	delete(s.watcher, w)
+	s.lock.Unlock()
+	s.registry.tryDelete(s)
+}
diff --git a/contrib/registry/consul/watcher.go b/contrib/registry/consul/watcher.go
index 0a5d3575..72150252 100644
--- a/contrib/registry/consul/watcher.go
+++ b/contrib/registry/consul/watcher.go
@@ -16,6 +16,10 @@ type watcher struct {
 }
 
 func (w *watcher) Next() (services []*registry.ServiceInstance, err error) {
+	if err = w.ctx.Err(); err != nil {
+		return
+	}
+
 	select {
 	case <-w.ctx.Done():
 		err = w.ctx.Err()
@@ -24,7 +28,6 @@ func (w *watcher) Next() (services []*registry.ServiceInstance, err error) {
 	}
 
 	ss, ok := w.set.services.Load().([]*registry.ServiceInstance)
-
 	if ok {
 		services = append(services, ss...)
 	}
@@ -32,9 +35,10 @@ func (w *watcher) Next() (services []*registry.ServiceInstance, err error) {
 }
 
 func (w *watcher) Stop() error {
-	w.cancel()
-	w.set.lock.Lock()
-	defer w.set.lock.Unlock()
-	delete(w.set.watcher, w)
+	if w.cancel != nil {
+		w.cancel()
+		w.cancel = nil
+		w.set.delete(w)
+	}
 	return nil
 }


Test patch:
diff --git a/contrib/registry/consul/registry_test.go b/contrib/registry/consul/registry_test.go
index 743e1dfb..12c95e33 100644
--- a/contrib/registry/consul/registry_test.go
+++ b/contrib/registry/consul/registry_test.go
@@ -2,9 +2,15 @@ package consul
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
+	"io"
 	"net"
+	"net/http"
+	"net/http/httptest"
 	"reflect"
+	"strconv"
+	"sync"
 	"testing"
 	"time"
 
@@ -19,8 +25,10 @@ func tcpServer(lis net.Listener) {
 		if err != nil {
 			return
 		}
-		fmt.Println("get tcp")
-		conn.Close()
+		go func() {
+			_, _ = io.Copy(io.Discard, conn)
+			_ = conn.Close()
+		}()
 	}
 }
 
@@ -112,18 +120,27 @@ func TestRegistry_Register(t *testing.T) {
 			r := New(cli, opts...)
 
 			for _, instance := range tt.args.server {
+				instance := instance
 				err = r.Register(tt.args.ctx, instance)
 				if err != nil {
 					t.Error(err)
 				}
+				defer func() {
+					err = r.Deregister(tt.args.ctx, instance)
+					if err != nil {
+						t.Error(err)
+					}
+				}()
 			}
 			watchCtx, watchCancel := context.WithCancel(context.Background())
 			watch, err := r.Watch(watchCtx, tt.args.serverName)
 			if err != nil {
 				t.Error(err)
+				watchCancel()
+				return
 			}
-			got, err := watch.Next()
 
+			got, err := watch.Next()
 			if (err != nil) != tt.wantErr {
 				t.Errorf("GetService() error = %v, wantErr %v", err, tt.wantErr)
 				t.Errorf("GetService() got = %v", got)
@@ -134,9 +151,6 @@ func TestRegistry_Register(t *testing.T) {
 				t.Errorf("GetService() got = %v, want %v", got, tt.want)
 			}
 
-			for _, instance := range tt.args.server {
-				_ = r.Deregister(tt.args.ctx, instance)
-			}
 			err = watch.Stop()
 			if err != nil {
 				t.Error(err)
@@ -293,8 +307,14 @@ func TestRegistry_GetService(t *testing.T) {
 
 func TestRegistry_Watch(t *testing.T) {
 	addr := fmt.Sprintf("%s:9091", getIntranetIP())
+	lis, err := net.Listen("tcp", addr)
+	if err != nil {
+		t.Errorf("listen tcp %s failed!", addr)
+		return
+	}
+	defer lis.Close()
+	go tcpServer(lis)
 
-	time.Sleep(time.Millisecond * 100)
 	cli, err := api.NewClient(&api.Config{Address: "127.0.0.1:8500", WaitTime: 2 * time.Second})
 	if err != nil {
 		t.Fatalf("create consul client failed: %v", err)
@@ -376,14 +396,7 @@ func TestRegistry_Watch(t *testing.T) {
 			},
 			want:    []*registry.ServiceInstance{instance3},
 			wantErr: false,
-			preFunc: func(t *testing.T) {
-				lis, err := net.Listen("tcp", addr)
-				if err != nil {
-					t.Errorf("listen tcp %s failed!", addr)
-					return
-				}
-				go tcpServer(lis)
-			},
+			preFunc: func(*testing.T) {},
 		},
 	}
 
@@ -398,9 +411,10 @@ func TestRegistry_Watch(t *testing.T) {
 			err := r.Register(tt.args.ctx, tt.args.instance)
 			if err != nil {
 				t.Error(err)
+				return
 			}
 			defer func() {
-				err = r.Deregister(tt.args.ctx, tt.args.instance)
+				err = r.Deregister(context.Background(), tt.args.instance)
 				if err != nil {
 					t.Error(err)
 				}
@@ -465,10 +479,10 @@ func TestRegistry_IdleAndWatch(t *testing.T) {
 	}
 
 	tests := []struct {
-		name    string
-		args    args
-		want    []*registry.ServiceInstance
-		wantErr bool
+		name  string
+		args  args
+		want1 []*registry.ServiceInstance
+		want2 []*registry.ServiceInstance
 	}{
 		{
 			name: "many client, one idle",
@@ -477,77 +491,86 @@ func TestRegistry_IdleAndWatch(t *testing.T) {
 				instance:       instance1,
 				changeInstance: instance2,
 			},
-			want:    []*registry.ServiceInstance{instance1},
-			wantErr: false,
+			want1: []*registry.ServiceInstance{instance1},
+			want2: []*registry.ServiceInstance{instance2},
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
+			var watchs []registry.Watcher
+			for i := 0; i < 10; i++ {
+				watch, err := r.Watch(tt.args.ctx, tt.args.instance.Name) //nolint
+				if err != nil {
+					t.Error(err)
+				}
+				defer func() {
+					_ = watch.Stop()
+				}()
+
+				watchs = append(watchs, watch)
+			}
+
 			err = r.Register(tt.args.ctx, tt.args.instance)
 			if err != nil {
 				t.Error(err)
 			}
 			defer func() {
-				err = r.Deregister(tt.args.ctx, tt.args.instance)
+				err = r.Deregister(context.Background(), tt.args.instance)
 				if err != nil {
 					t.Error(err)
 				}
 			}()
 
-			for i := 0; i < 10; i++ {
-				watchCtx, watchCancel := context.WithCancel(context.Background())
-				watch, err1 := r.Watch(watchCtx, tt.args.instance.Name)
-				if err1 != nil {
-					t.Error(err1)
-				}
-				if i != 9 {
-					go func(i int) {
-						// first
-						service, err2 := watch.Next()
-						if (err2 != nil) != tt.wantErr {
-							t.Errorf("GetService() error = %v, wantErr %v", err, tt.wantErr)
-							t.Errorf("GetService() got = %v", service)
-							watchCancel()
-							return
-						}
-						// instance changes
-						service, err2 = watch.Next()
-						if i == 9 {
-							return
-						}
-						if (err2 != nil) != tt.wantErr {
-							t.Errorf("GetService() error = %v, wantErr %v", err, tt.wantErr)
-							t.Errorf("GetService() got = %v", service)
-							watchCancel()
-							return
-						}
-						if !reflect.DeepEqual(service, tt.want) {
-							t.Errorf("GetService() got = %v, want %v", service, tt.want)
-						}
-						err2 = watch.Stop()
-						if err2 != nil {
-							t.Error(err)
-						}
-						watchCancel()
-						// t.Logf("service:%v, i:%d", service, i)
-					}(i)
-				} else {
-					time.Sleep(time.Second * 3)
-					// become idle, close watcher
-					err1 = watch.Stop()
-					if err1 != nil {
-						t.Errorf("watch stop err:%v", err)
+			var wg1 sync.WaitGroup
+			for _, watch := range watchs {
+				wg1.Add(1)
+				go func(watch registry.Watcher, want []*registry.ServiceInstance) {
+					defer wg1.Done()
+
+					// first
+					service, err := watch.Next() //nolint
+					if err != nil {
+						t.Error(err)
+						return
 					}
-					watchCancel()
-				}
+					if !reflect.DeepEqual(service, want) {
+						t.Errorf("GetService() got = %v, want = %v", service, want)
+						return
+					}
+				}(watch, tt.want1)
 			}
-			time.Sleep(2 * time.Second)
+			wg1.Wait()
+
 			err = r.Register(tt.args.ctx, tt.args.changeInstance)
 			if err != nil {
 				t.Error(err)
 			}
-			time.Sleep(1 * time.Second)
+			defer func() {
+				err := r.Deregister(context.Background(), tt.args.changeInstance)
+				if err != nil {
+					t.Error(err)
+				}
+			}()
+
+			var wg2 sync.WaitGroup
+			for _, watch := range watchs {
+				wg2.Add(1)
+				go func(watch registry.Watcher, want []*registry.ServiceInstance) {
+					defer wg2.Done()
+
+					// instance changes
+					service, err := watch.Next() //nolint
+					if err != nil {
+						t.Error(err)
+						return
+					}
+					if !reflect.DeepEqual(service, want) {
+						t.Errorf("GetService() got = %v, want = %v", service, want)
+					}
+				}(watch, tt.want2)
+			}
+			wg2.Wait()
 		})
 	}
 }
@@ -767,6 +790,12 @@ func TestRegistry_ExitOldResolverAndReWatch(t *testing.T) {
 			if err != nil {
 				t.Error(err)
 			}
+			defer func() {
+				err = r.Deregister(tt.args.ctx, tt.args.instance)
+				if err != nil {
+					t.Error(err)
+				}
+			}()
 
 			time.Sleep(time.Second * 2)
 
@@ -796,6 +825,193 @@ func TestRegistry_ExitOldResolverAndReWatch(t *testing.T) {
 	}
 }
 
+func TestRegistry_ShareServiceSet(t *testing.T) {
+	lastIndex := uint64(0)
+	serviceName := "share-service-set"
+	mux := http.NewServeMux()
+	mux.HandleFunc("/v1/health/service/"+serviceName, func(w http.ResponseWriter, r *http.Request) {
+		var index uint64
+		if s := r.URL.Query().Get("index"); s != "" {
+			val, err := strconv.ParseUint(s, 10, 64)
+			if err != nil {
+				http.Error(w, err.Error(), http.StatusBadRequest)
+				return
+			}
+			index = val
+		}
+
+		if index < lastIndex {
+			msg := "repeated request, not the same ServiceSet"
+			http.Error(w, msg, http.StatusBadRequest)
+			t.Error(msg)
+			t.FailNow()
+			return
+		}
+
+		lastIndex = index + 1
+		w.Header().Set("X-Consul-Index", strconv.FormatUint(lastIndex, 10))
+
+		out := []*api.ServiceEntry{
+			{
+				Service: &api.AgentService{
+					ID:      "1",
+					Service: serviceName,
+				},
+			},
+		}
+		err := json.NewEncoder(w).Encode(out)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+		}
+	})
+
+	ts := httptest.NewServer(mux)
+	defer ts.Close()
+
+	cli, err := api.NewClient(&api.Config{Address: ts.URL, WaitTime: 2 * time.Second})
+	if err != nil {
+		t.Fatalf("create consul client failed: %v", err)
+	}
+
+	var prev registry.Watcher
+	r := New(cli, WithHealthCheck(false), WithHeartbeat(false))
+	for i := 0; i < 100; i++ {
+		w, err := r.Watch(context.Background(), serviceName) //nolint
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		// close previous watcher
+		if prev != nil {
+			if err = prev.Stop(); err != nil {
+				t.Error(err)
+				return
+			}
+		}
+		prev = w
+	}
+
+	time.Sleep(time.Second * 5)
+
+	if prev != nil {
+		if err = prev.Stop(); err != nil {
+			t.Error(err)
+			return
+		}
+	}
+}
+
+func TestRegistry_MultiWatch(t *testing.T) {
+	cli, err := api.NewClient(&api.Config{Address: "127.0.0.1:8500", WaitTime: 2 * time.Second})
+	if err != nil {
+		t.Fatalf("create consul client failed: %v", err)
+	}
+
+	serviceName := "multi-watch"
+	addr := fmt.Sprintf("%s:9091", getIntranetIP())
+	instances := []*registry.ServiceInstance{
+		{
+			ID:        "1",
+			Name:      serviceName,
+			Version:   "v1.0.0",
+			Endpoints: []string{fmt.Sprintf("tcp://%s?isSecure=false", addr)},
+		},
+		{
+			ID:        "2",
+			Name:      serviceName,
+			Version:   "v1.0.0",
+			Endpoints: []string{fmt.Sprintf("tcp://%s?isSecure=false", addr)},
+		},
+	}
+
+	r := New(cli, WithHealthCheck(false), WithHeartbeat(true))
+	err = r.Register(context.Background(), instances[0])
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	defer func() {
+		err = r.Deregister(context.Background(), instances[0])
+		if err != nil {
+			t.Error(err)
+		}
+	}()
+
+	watch1, err := r.Watch(context.Background(), serviceName)
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	defer func() {
+		if err = watch1.Stop(); err != nil {
+			t.Error(err)
+		}
+	}()
+
+	watch2, err := r.Watch(context.Background(), serviceName)
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	defer func() {
+		if err = watch2.Stop(); err != nil {
+			t.Error(err)
+		}
+	}()
+
+	got1, err := watch1.Next()
+	if err != nil {
+		t.Error(err)
+		return
+	}
+
+	got2, err := watch2.Next()
+	if err != nil {
+		t.Error(err)
+		return
+	}
+
+	if !reflect.DeepEqual(got1, instances[:1]) {
+		t.Errorf("got = %v, want = %v", got1, instances[:1])
+		return
+	}
+	if !reflect.DeepEqual(got2, instances[:1]) {
+		t.Errorf("got = %v, want = %v", got2, instances[:1])
+		return
+	}
+
+	// close first watcher
+	if err = watch1.Stop(); err != nil {
+		t.Error(err)
+		return
+	}
+
+	// register a new instance
+	err = r.Register(context.Background(), instances[1])
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	defer func() {
+		err = r.Deregister(context.Background(), instances[1])
+		if err != nil {
+			t.Error(err)
+		}
+	}()
+
+	// second watcher should get the new instance
+	got, err := watch2.Next()
+	if err != nil {
+		t.Error(err)
+		return
+	}
+
+	if !reflect.DeepEqual(got, instances[:2]) {
+		t.Errorf("got = %v, want = %v", got, instances[:2])
+		return
+	}
+}
+
 func getIntranetIP() string {
 	addrs, err := net.InterfaceAddrs()
 	if err != nil {