This is an automated email from the ASF dual-hosted git repository.
thunguo pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-seata-go.git
The following commit(s) were added to refs/heads/master by this push:
new ee76dddc [GSOC] Add discovery provider registry and address store
(#1126)
ee76dddc is described below
commit ee76dddc72c53e9c4e8eda22a2b6170c2b1b2b0f
Author: CAICAII <[email protected]>
AuthorDate: Fri Jul 10 21:23:21 2026 +0800
[GSOC] Add discovery provider registry and address store (#1126)
* refactor: add discovery provider registry
* fix: close etcd client on registry shutdown
* test: cover registry init failure state
---------
Co-authored-by: ThunGuo <[email protected]>
---
pkg/discovery/base.go | 6 ++
pkg/discovery/etcd3.go | 94 ++++++---------------
pkg/discovery/etcd3_test.go | 44 ++++++++--
pkg/discovery/init.go | 48 +++++------
pkg/discovery/init_test.go | 188 +++++++++++++++++++++++++++++++++++++++++
pkg/discovery/provider.go | 61 +++++++++++++
pkg/discovery/provider_test.go | 57 +++++++++++++
pkg/discovery/store.go | 179 +++++++++++++++++++++++++++++++++++++++
pkg/discovery/store_test.go | 142 +++++++++++++++++++++++++++++++
9 files changed, 717 insertions(+), 102 deletions(-)
diff --git a/pkg/discovery/base.go b/pkg/discovery/base.go
index 4fc0f3c6..9a30b04e 100644
--- a/pkg/discovery/base.go
+++ b/pkg/discovery/base.go
@@ -21,6 +21,7 @@ const (
FILE string = "file"
NACOS string = "nacos"
ETCD string = "etcd"
+ ETCD3 string = "etcd3"
EUREKA string = "eureka"
REDIS string = "redis"
ZK string = "zk"
@@ -39,3 +40,8 @@ type RegistryService interface {
Lookup(key string) ([]*ServiceInstance, error)
Close()
}
+
+type RegisterableRegistryService interface {
+ Register(instance *ServiceInstance) error
+ Unregister(instance *ServiceInstance) error
+}
diff --git a/pkg/discovery/etcd3.go b/pkg/discovery/etcd3.go
index fb71c7b7..3997ab88 100644
--- a/pkg/discovery/etcd3.go
+++ b/pkg/discovery/etcd3.go
@@ -19,6 +19,7 @@ package discovery
import (
"context"
+ "errors"
"fmt"
"strings"
"sync"
@@ -39,17 +40,18 @@ type EtcdRegistryService struct {
client *etcd3.Client
cfg etcd3.Config
vgroupMapping map[string]string
- grouplist map[string][]*ServiceInstance
- rwLock sync.RWMutex
+ store *AddressStore
- stopCh chan struct{}
+ stopCh chan struct{}
+ closeOnce sync.Once
}
-func newEtcdRegistryService(config *ServiceConfig, etcd3Config *Etcd3Config)
RegistryService {
-
+func newEtcdRegistryService(config *ServiceConfig, etcd3Config *Etcd3Config)
(RegistryService, error) {
+ if config == nil {
+ return nil, fmt.Errorf("service config is nil")
+ }
if etcd3Config == nil {
- log.Fatalf("etcd config is nil")
- panic("etcd config is nil")
+ return nil, fmt.Errorf("etcd config is nil")
}
cfg := etcd3.Config{
@@ -57,23 +59,21 @@ func newEtcdRegistryService(config *ServiceConfig,
etcd3Config *Etcd3Config) Reg
}
cli, err := etcd3.New(cfg)
if err != nil {
- log.Fatalf("failed to create etcd3 client")
- panic("failed to create etcd3 client")
+ return nil, fmt.Errorf("failed to create etcd3 client: %w", err)
}
vgroupMapping := config.VgroupMapping
- grouplist := make(map[string][]*ServiceInstance, 0)
etcdRegistryService := &EtcdRegistryService{
client: cli,
cfg: cfg,
vgroupMapping: vgroupMapping,
- grouplist: grouplist,
+ store: NewAddressStore(),
stopCh: make(chan struct{}),
}
go etcdRegistryService.watch(etcdClusterPrefix)
- return etcdRegistryService
+ return etcdRegistryService, nil
}
func (s *EtcdRegistryService) watch(key string) {
@@ -100,13 +100,7 @@ func (s *EtcdRegistryService) watch(key string) {
log.Errorf("etcd value has an incorrect format:
%v", err)
return
}
- s.rwLock.Lock()
- if s.grouplist[clusterName] == nil {
- s.grouplist[clusterName] =
[]*ServiceInstance{serverInstance}
- } else {
- s.grouplist[clusterName] =
append(s.grouplist[clusterName], serverInstance)
- }
- s.rwLock.Unlock()
+ s.store.upsert(clusterName, serverInstance)
}
}
@@ -138,18 +132,7 @@ func (s *EtcdRegistryService) watch(key string) {
return
}
- s.rwLock.Lock()
- if s.grouplist[clusterName] == nil {
- s.grouplist[clusterName] =
[]*ServiceInstance{serverInstance}
- s.rwLock.Unlock()
- continue
- }
- if
ifHaveSameServiceInstances(s.grouplist[clusterName], serverInstance) {
- s.rwLock.Unlock()
- continue
- }
- s.grouplist[clusterName] =
append(s.grouplist[clusterName], serverInstance)
- s.rwLock.Unlock()
+ s.store.upsert(clusterName,
serverInstance)
case etcd3.EventTypeDelete:
log.Infof("Key %s deleted.\n",
event.Kv.Key)
@@ -160,15 +143,9 @@ func (s *EtcdRegistryService) watch(key string) {
return
}
- s.rwLock.Lock()
- serviceInstances := s.grouplist[cluster]
- if serviceInstances == nil {
- log.Warnf("etcd doesnt exit
cluster: ", cluster)
- s.rwLock.Unlock()
- continue
+ if !s.store.remove(cluster, ip, port) {
+ log.Warnf("etcd instance not
found. cluster: %s addr: %s:%d", cluster, ip, port)
}
- s.grouplist[cluster] =
removeValueFromList(serviceInstances, ip, port)
- s.rwLock.Unlock()
}
}
case <-s.stopCh:
@@ -218,41 +195,24 @@ func getClusterAndAddress(key []byte) (string, string,
int, error) {
return cluster, ip, port, nil
}
-func ifHaveSameServiceInstances(list []*ServiceInstance, value
*ServiceInstance) bool {
- for _, v := range list {
- if v.Addr == value.Addr && v.Port == value.Port {
- return true
- }
- }
- return false
-}
-
-func removeValueFromList(list []*ServiceInstance, ip string, port int)
[]*ServiceInstance {
- for k, v := range list {
- if v.Addr == ip && v.Port == port {
- result := list[:k]
- if k < len(list)-1 {
- result = append(result, list[k+1:]...)
- }
- return result
- }
- }
-
- return list
-}
-
func (s *EtcdRegistryService) Lookup(key string) ([]*ServiceInstance, error) {
- s.rwLock.RLock()
- defer s.rwLock.RUnlock()
cluster := s.vgroupMapping[key]
if cluster == "" {
return nil, fmt.Errorf("cluster doesnt exit")
}
- list := s.grouplist[cluster]
- return list, nil
+ return s.store.Snapshot(cluster), nil
}
func (s *EtcdRegistryService) Close() {
- s.stopCh <- struct{}{}
+ s.closeOnce.Do(func() {
+ if s.stopCh != nil {
+ close(s.stopCh)
+ }
+ if s.client != nil {
+ if err := s.client.Close(); err != nil &&
!errors.Is(err, context.Canceled) {
+ log.Warnf("close etcd client failed: %v", err)
+ }
+ }
+ })
}
diff --git a/pkg/discovery/etcd3_test.go b/pkg/discovery/etcd3_test.go
index e3b4fa3d..e8322688 100644
--- a/pkg/discovery/etcd3_test.go
+++ b/pkg/discovery/etcd3_test.go
@@ -18,6 +18,7 @@
package discovery
import (
+ "context"
"reflect"
"testing"
"time"
@@ -133,20 +134,18 @@ func TestEtcd3RegistryService_Lookup(t *testing.T) {
ctrl := gomock.NewController(t)
mockEtcdClient := mock.NewMockEtcdClient(ctrl)
etcdRegistryService := &EtcdRegistryService{
- client: &clientv3.Client{
- KV: mockEtcdClient,
- Watcher: mockEtcdClient,
- },
+ client: newTestEtcdClient(mockEtcdClient),
vgroupMapping: map[string]string{
"default_tx_group": "default",
},
- grouplist: make(map[string][]*ServiceInstance, 0),
- stopCh: make(chan struct{}),
+ store: NewAddressStore(),
+ stopCh: make(chan struct{}),
}
mockEtcdClient.EXPECT().Get(gomock.Any(), gomock.Any(),
gomock.Any()).Return(tt.getResp, nil)
ch := make(chan clientv3.WatchResponse)
mockEtcdClient.EXPECT().Watch(gomock.Any(), gomock.Any(),
gomock.Any()).Return(ch)
+ mockEtcdClient.EXPECT().Close().Return(nil)
go func() {
etcdRegistryService.watch("registry-seata")
@@ -176,3 +175,36 @@ func TestEtcd3RegistryService_Lookup(t *testing.T) {
etcdRegistryService.Close()
}
}
+
+func TestEtcd3RegistryService_CloseIsRepeatable(t *testing.T) {
+ client := clientv3.NewCtxClient(context.Background())
+ etcdRegistryService := &EtcdRegistryService{
+ client: client,
+ stopCh: make(chan struct{}),
+ }
+
+ etcdRegistryService.Close()
+ etcdRegistryService.Close()
+
+ select {
+ case <-etcdRegistryService.stopCh:
+ case <-time.After(time.Second):
+ t.Fatal("stop channel was not closed")
+ }
+
+ select {
+ case <-client.Ctx().Done():
+ case <-time.After(time.Second):
+ t.Fatal("etcd client was not closed")
+ }
+}
+
+func newTestEtcdClient(client mock.EtcdClient) *clientv3.Client {
+ return clientv3.NewCtxClient(
+ context.Background(),
+ func(c *clientv3.Client) {
+ c.KV = client
+ c.Watcher = client
+ },
+ )
+}
diff --git a/pkg/discovery/init.go b/pkg/discovery/init.go
index 98855b2f..57daf48e 100644
--- a/pkg/discovery/init.go
+++ b/pkg/discovery/init.go
@@ -26,40 +26,30 @@ var (
)
func InitRegistry(serviceConfig *ServiceConfig, registryConfig
*RegistryConfig) {
- var registryService RegistryService
- var err error
- switch registryConfig.Type {
- case FILE:
- //init file registry
- registryService = newFileRegistryService(serviceConfig)
- case ETCD:
- //init etcd registry
- registryService = newEtcdRegistryService(serviceConfig,
®istryConfig.Etcd3)
- case RAFT:
- registryService = NewRaftRegistryService(serviceConfig,
registryConfig)
- case NACOS:
- //TODO: init nacos registry
- case EUREKA:
- //TODO: init eureka registry
- case REDIS:
- //TODO: init redis registry
- case ZK:
- //TODO: init zk registry
- case CONSUL:
- //TODO: init consul registry
- case SOFA:
- //TODO: init sofa registry
- case NAMINGSERVER:
- // init namingserver registry
- registryService = newNamingServerRegistryService(serviceConfig,
®istryConfig.NamingServer)
- default:
- err = fmt.Errorf("service registry not support registry
type:%s", registryConfig.Type)
+ if err := InitRegistryWithError(serviceConfig, registryConfig); err !=
nil {
+ panic(fmt.Errorf("init service registry err:%v", err))
}
+}
+func InitRegistryWithError(serviceConfig *ServiceConfig, registryConfig
*RegistryConfig) error {
+ if registryConfig == nil {
+ return fmt.Errorf("registry config is nil")
+ }
+
+ provider, ok := registryProviderFor(registryConfig.Type)
+ if !ok {
+ return unsupportedRegistryTypeError(registryConfig.Type)
+ }
+
+ registryService, err := provider(serviceConfig, registryConfig)
if err != nil {
- panic(fmt.Errorf("init service registry err:%v", err))
+ return err
+ }
+ if registryService == nil {
+ return fmt.Errorf("registry provider returned nil for type:%s",
registryConfig.Type)
}
registryServiceInstance = registryService
+ return nil
}
func GetRegistry() RegistryService {
diff --git a/pkg/discovery/init_test.go b/pkg/discovery/init_test.go
index e7835067..a50d11f4 100644
--- a/pkg/discovery/init_test.go
+++ b/pkg/discovery/init_test.go
@@ -18,7 +18,9 @@
package discovery
import (
+ "fmt"
"reflect"
+ "strings"
"testing"
)
@@ -75,6 +77,12 @@ func TestInitRegistry(t *testing.T) {
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
+ t.Cleanup(func() {
+ if registryServiceInstance != nil {
+ registryServiceInstance.Close()
+ registryServiceInstance = nil
+ }
+ })
defer func() {
if r := recover(); r != nil {
if !tt.hasPanic {
@@ -95,3 +103,183 @@ func TestInitRegistry(t *testing.T) {
})
}
}
+
+func TestInitRegistryWithErrorUnsupportedType(t *testing.T) {
+ registryServiceInstance = nil
+ t.Cleanup(func() {
+ registryServiceInstance = nil
+ })
+
+ err := InitRegistryWithError(&ServiceConfig{}, &RegistryConfig{Type:
"unknown"})
+ if err == nil {
+ t.Fatal("expected error")
+ }
+ if err.Error() != "service registry not support registry type:unknown" {
+ t.Fatalf("unexpected error: %v", err)
+ }
+ if GetRegistry() != nil {
+ t.Fatal("registry should not be initialized on error")
+ }
+}
+
+func TestInitRegistryWithErrorNilRegistryConfig(t *testing.T) {
+ err := InitRegistryWithError(&ServiceConfig{}, nil)
+ if err == nil {
+ t.Fatal("expected error")
+ }
+ if err.Error() != "registry config is nil" {
+ t.Fatalf("unexpected error: %v", err)
+ }
+}
+
+func TestInitRegistryWithErrorReturnsProviderError(t *testing.T) {
+ registryServiceInstance = nil
+ t.Cleanup(func() {
+ registryServiceInstance = nil
+ })
+
+ err := InitRegistryWithError(nil, &RegistryConfig{Type: FILE})
+ if err == nil {
+ t.Fatal("expected error")
+ }
+ if err.Error() != "service config is nil" {
+ t.Fatalf("unexpected error: %v", err)
+ }
+ if GetRegistry() != nil {
+ t.Fatal("registry should not be initialized on error")
+ }
+}
+
+func TestInitRegistryWithErrorKeepsExistingRegistryOnError(t *testing.T) {
+ registryServiceInstance = nil
+ t.Cleanup(func() {
+ registryServiceInstance = nil
+ })
+
+ if err := InitRegistryWithError(&ServiceConfig{}, &RegistryConfig{Type:
FILE}); err != nil {
+ t.Fatalf("InitRegistryWithError() error = %v", err)
+ }
+ existing := GetRegistry()
+ if existing == nil {
+ t.Fatal("registry is nil")
+ }
+
+ err := InitRegistryWithError(nil, &RegistryConfig{Type: FILE})
+ if err == nil {
+ t.Fatal("expected error")
+ }
+ if err.Error() != "service config is nil" {
+ t.Fatalf("unexpected error: %v", err)
+ }
+ if GetRegistry() != existing {
+ t.Fatal("existing registry should be kept on error")
+ }
+}
+
+func TestInitRegistryWithErrorNilProviderResult(t *testing.T) {
+ const providerType = "empty"
+ oldProvider, hadOldProvider := registryProviders[providerType]
+ registryProviders[providerType] = func(*ServiceConfig, *RegistryConfig)
(RegistryService, error) {
+ return nil, nil
+ }
+ t.Cleanup(func() {
+ if hadOldProvider {
+ registryProviders[providerType] = oldProvider
+ } else {
+ delete(registryProviders, providerType)
+ }
+ })
+
+ err := InitRegistryWithError(&ServiceConfig{}, &RegistryConfig{Type:
providerType})
+ if err == nil {
+ t.Fatal("expected error")
+ }
+ if err.Error() != "registry provider returned nil for type:empty" {
+ t.Fatalf("unexpected error: %v", err)
+ }
+}
+
+func TestInitRegistryWithErrorSupportedProviders(t *testing.T) {
+ tests := []struct {
+ name string
+ serviceConfig *ServiceConfig
+ registryConfig *RegistryConfig
+ expectedType string
+ cleanup func()
+ }{
+ {
+ name: "raft",
+ serviceConfig: &ServiceConfig{
+ VgroupMapping: map[string]string{
+ "default_tx_group": "default",
+ },
+ },
+ registryConfig: &RegistryConfig{
+ Type: RAFT,
+ Raft: RaftConfig{
+ ServerAddr: "127.0.0.1:7091",
+ },
+ },
+ expectedType: "RaftRegistryService",
+ },
+ {
+ name: "namingserver",
+ serviceConfig: nil,
+ registryConfig: &RegistryConfig{
+ Type: NAMINGSERVER,
+ NamingServer: NamingServerConfig{
+ ServerAddr: "127.0.0.1:8081",
+ HeartbeatPeriod: 5000,
+ },
+ },
+ expectedType: "NamingServerRegistryService",
+ cleanup: resetInstance,
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ registryServiceInstance = nil
+ if tt.cleanup != nil {
+ tt.cleanup()
+ }
+ t.Cleanup(func() {
+ if tt.cleanup != nil {
+ tt.cleanup()
+ registryServiceInstance = nil
+ return
+ }
+ if registryServiceInstance != nil {
+ registryServiceInstance.Close()
+ registryServiceInstance = nil
+ }
+ })
+
+ if err := InitRegistryWithError(tt.serviceConfig,
tt.registryConfig); err != nil {
+ t.Fatalf("InitRegistryWithError() error = %v",
err)
+ }
+ instance := GetRegistry()
+ if instance == nil {
+ t.Fatal("registry is nil")
+ }
+ actualType := reflect.TypeOf(instance).Elem().Name()
+ if actualType != tt.expectedType {
+ t.Fatalf("type = %v, want %v", actualType,
tt.expectedType)
+ }
+ })
+ }
+}
+
+func TestInitRegistryPanicCompatibility(t *testing.T) {
+ defer func() {
+ r := recover()
+ if r == nil {
+ t.Fatal("expected panic")
+ }
+ if !strings.Contains(fmt.Sprint(r), "init service registry
err:service registry not support registry type:unknown") {
+ t.Fatalf("unexpected panic: %v", r)
+ }
+ }()
+
+ InitRegistry(&ServiceConfig{}, &RegistryConfig{Type: "unknown"})
+}
diff --git a/pkg/discovery/provider.go b/pkg/discovery/provider.go
new file mode 100644
index 00000000..caf26866
--- /dev/null
+++ b/pkg/discovery/provider.go
@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package discovery
+
+import (
+ "fmt"
+)
+
+type RegistryProvider func(serviceConfig *ServiceConfig, registryConfig
*RegistryConfig) (RegistryService, error)
+
+var registryProviders = map[string]RegistryProvider{
+ FILE: func(serviceConfig *ServiceConfig, _ *RegistryConfig)
(RegistryService, error) {
+ if serviceConfig == nil {
+ return nil, fmt.Errorf("service config is nil")
+ }
+ return newFileRegistryService(serviceConfig), nil
+ },
+ ETCD3: func(serviceConfig *ServiceConfig, registryConfig
*RegistryConfig) (RegistryService, error) {
+ return newEtcdRegistryService(serviceConfig,
®istryConfig.Etcd3)
+ },
+ RAFT: func(serviceConfig *ServiceConfig, registryConfig
*RegistryConfig) (RegistryService, error) {
+ if serviceConfig == nil {
+ return nil, fmt.Errorf("service config is nil")
+ }
+ return NewRaftRegistryService(serviceConfig, registryConfig),
nil
+ },
+ NAMINGSERVER: func(serviceConfig *ServiceConfig, registryConfig
*RegistryConfig) (RegistryService, error) {
+ return newNamingServerRegistryService(serviceConfig,
®istryConfig.NamingServer), nil
+ },
+}
+
+func registryProviderFor(registryType string) (RegistryProvider, bool) {
+ provider, ok := registryProviders[normalizeRegistryType(registryType)]
+ return provider, ok
+}
+
+func normalizeRegistryType(registryType string) string {
+ if registryType == ETCD {
+ return ETCD3
+ }
+ return registryType
+}
+
+func unsupportedRegistryTypeError(registryType string) error {
+ return fmt.Errorf("service registry not support registry type:%s",
registryType)
+}
diff --git a/pkg/discovery/provider_test.go b/pkg/discovery/provider_test.go
new file mode 100644
index 00000000..7f9e4065
--- /dev/null
+++ b/pkg/discovery/provider_test.go
@@ -0,0 +1,57 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package discovery
+
+import "testing"
+
+func TestRegistryProviderForSupportedTypes(t *testing.T) {
+ tests := []struct {
+ name string
+ registryType string
+ }{
+ {name: "file", registryType: FILE},
+ {name: "etcd alias", registryType: ETCD},
+ {name: "etcd3", registryType: ETCD3},
+ {name: "raft", registryType: RAFT},
+ {name: "namingserver", registryType: NAMINGSERVER},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ provider, ok := registryProviderFor(tt.registryType)
+ if !ok {
+ t.Fatalf("provider not found for type %s",
tt.registryType)
+ }
+ if provider == nil {
+ t.Fatalf("provider is nil for type %s",
tt.registryType)
+ }
+ })
+ }
+}
+
+func TestRegistryProviderForUnsupportedType(t *testing.T) {
+ provider, ok := registryProviderFor("unknown")
+ if ok {
+ t.Fatalf("unexpected provider for unknown type: %v", provider)
+ }
+
+ err := unsupportedRegistryTypeError("unknown")
+ if err.Error() != "service registry not support registry type:unknown" {
+ t.Fatalf("unexpected error: %v", err)
+ }
+}
diff --git a/pkg/discovery/store.go b/pkg/discovery/store.go
new file mode 100644
index 00000000..fd53d194
--- /dev/null
+++ b/pkg/discovery/store.go
@@ -0,0 +1,179 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package discovery
+
+import (
+ "strconv"
+ "sync"
+)
+
+type AddressStoreSubscriber func(cluster string, instances []*ServiceInstance)
+
+type AddressStore struct {
+ mu sync.RWMutex
+ clusters map[string][]*ServiceInstance
+ subscribers map[uint64]AddressStoreSubscriber
+ nextID uint64
+}
+
+func NewAddressStore() *AddressStore {
+ return &AddressStore{
+ clusters: make(map[string][]*ServiceInstance),
+ subscribers: make(map[uint64]AddressStoreSubscriber),
+ }
+}
+
+func (s *AddressStore) Snapshot(cluster string) []*ServiceInstance {
+ s.mu.RLock()
+ defer s.mu.RUnlock()
+
+ return cloneServiceInstances(s.clusters[cluster])
+}
+
+func (s *AddressStore) Update(cluster string, instances []*ServiceInstance) {
+ s.saveAndNotify(cluster, instances)
+}
+
+func (s *AddressStore) Subscribe(subscriber AddressStoreSubscriber) func() {
+ if subscriber == nil {
+ return func() {}
+ }
+
+ s.mu.Lock()
+ id := s.nextID
+ s.nextID++
+ s.subscribers[id] = subscriber
+ s.mu.Unlock()
+
+ return func() {
+ s.mu.Lock()
+ delete(s.subscribers, id)
+ s.mu.Unlock()
+ }
+}
+
+func (s *AddressStore) upsert(cluster string, instance *ServiceInstance) {
+ if instance == nil {
+ return
+ }
+
+ s.mu.Lock()
+ next := cloneServiceInstances(s.clusters[cluster])
+ key := serviceInstanceKey(instance)
+ for i := range next {
+ if serviceInstanceKey(next[i]) == key {
+ s.mu.Unlock()
+ return
+ }
+ }
+ next = append(next, cloneServiceInstance(instance))
+ snapshot := s.saveLocked(cluster, next)
+ subscribers := s.subscribersLocked()
+ s.mu.Unlock()
+ notifyAddressStoreSubscribers(subscribers, cluster, snapshot)
+}
+
+func (s *AddressStore) remove(cluster, addr string, port int) bool {
+ s.mu.Lock()
+ current := s.clusters[cluster]
+ next := make([]*ServiceInstance, 0, len(current))
+ removeKey := serviceInstanceKey(&ServiceInstance{Addr: addr, Port:
port})
+ removed := false
+ for _, instance := range current {
+ if serviceInstanceKey(instance) == removeKey {
+ removed = true
+ continue
+ }
+ next = append(next, cloneServiceInstance(instance))
+ }
+ if !removed {
+ s.mu.Unlock()
+ return false
+ }
+ snapshot := s.saveLocked(cluster, next)
+ subscribers := s.subscribersLocked()
+ s.mu.Unlock()
+ notifyAddressStoreSubscribers(subscribers, cluster, snapshot)
+ return true
+}
+
+func (s *AddressStore) saveAndNotify(cluster string, instances
[]*ServiceInstance) {
+ s.mu.Lock()
+ snapshot := s.saveLocked(cluster, instances)
+ subscribers := s.subscribersLocked()
+ s.mu.Unlock()
+ notifyAddressStoreSubscribers(subscribers, cluster, snapshot)
+}
+
+func (s *AddressStore) saveLocked(cluster string, instances
[]*ServiceInstance) []*ServiceInstance {
+ next := cloneServiceInstances(instances)
+ if len(next) == 0 {
+ delete(s.clusters, cluster)
+ } else {
+ s.clusters[cluster] = next
+ }
+ return next
+}
+
+func (s *AddressStore) subscribersLocked() []AddressStoreSubscriber {
+ subscribers := make([]AddressStoreSubscriber, 0, len(s.subscribers))
+ for _, subscriber := range s.subscribers {
+ subscribers = append(subscribers, subscriber)
+ }
+ return subscribers
+}
+
+func notifyAddressStoreSubscribers(subscribers []AddressStoreSubscriber,
cluster string, instances []*ServiceInstance) {
+ // Callbacks run without the store lock so subscribers can call
Snapshot.
+ for _, subscriber := range subscribers {
+ subscriber(cluster, cloneServiceInstances(instances))
+ }
+}
+
+func cloneServiceInstances(instances []*ServiceInstance) []*ServiceInstance {
+ if len(instances) == 0 {
+ return nil
+ }
+
+ result := make([]*ServiceInstance, 0, len(instances))
+ for _, instance := range instances {
+ if instance == nil {
+ continue
+ }
+ result = append(result, cloneServiceInstance(instance))
+ }
+ if len(result) == 0 {
+ return nil
+ }
+ return result
+}
+
+func cloneServiceInstance(instance *ServiceInstance) *ServiceInstance {
+ if instance == nil {
+ return nil
+ }
+ clone := *instance
+ return &clone
+}
+
+func serviceInstanceKey(instance *ServiceInstance) string {
+ if instance == nil {
+ return ""
+ }
+ return instance.Addr + ":" + strconv.Itoa(instance.Port)
+}
diff --git a/pkg/discovery/store_test.go b/pkg/discovery/store_test.go
new file mode 100644
index 00000000..5832d8ac
--- /dev/null
+++ b/pkg/discovery/store_test.go
@@ -0,0 +1,142 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package discovery
+
+import (
+ "fmt"
+ "sync"
+ "testing"
+)
+
+func TestAddressStoreUpdateAndSnapshotClone(t *testing.T) {
+ store := NewAddressStore()
+ input := []*ServiceInstance{
+ {Addr: "127.0.0.1", Port: 8091},
+ {Addr: "127.0.0.2", Port: 8092},
+ }
+
+ store.Update("default", input)
+ input[0].Addr = "10.0.0.1"
+
+ snapshot := store.Snapshot("default")
+ if len(snapshot) != 2 {
+ t.Fatalf("snapshot length = %d, want 2", len(snapshot))
+ }
+ if snapshot[0].Addr != "127.0.0.1" {
+ t.Fatalf("snapshot was changed by caller input: %v",
snapshot[0])
+ }
+
+ snapshot[0].Port = 9999
+ next := store.Snapshot("default")
+ if next[0].Port != 8091 {
+ t.Fatalf("store was changed by snapshot mutation: %v", next[0])
+ }
+}
+
+func TestAddressStoreSubscribeAndUnsubscribe(t *testing.T) {
+ store := NewAddressStore()
+ received := make(chan []*ServiceInstance, 1)
+
+ unsubscribe := store.Subscribe(func(cluster string, instances
[]*ServiceInstance) {
+ if cluster != "default" {
+ t.Fatalf("cluster = %s, want default", cluster)
+ }
+ instances[0].Addr = "mutated"
+ received <- instances
+ })
+
+ store.Update("default", []*ServiceInstance{{Addr: "127.0.0.1", Port:
8091}})
+ got := <-received
+ if got[0].Addr != "mutated" {
+ t.Fatalf("subscriber did not receive a mutable copy: %v",
got[0])
+ }
+ if snapshot := store.Snapshot("default"); snapshot[0].Addr !=
"127.0.0.1" {
+ t.Fatalf("subscriber mutated store snapshot: %v", snapshot[0])
+ }
+
+ unsubscribe()
+ store.Update("default", []*ServiceInstance{{Addr: "127.0.0.2", Port:
8092}})
+ select {
+ case got := <-received:
+ t.Fatalf("received update after unsubscribe: %v", got)
+ default:
+ }
+}
+
+func TestAddressStoreUpsertRemove(t *testing.T) {
+ store := NewAddressStore()
+ updates := 0
+ store.Subscribe(func(string, []*ServiceInstance) {
+ updates++
+ })
+
+ store.upsert("default", &ServiceInstance{Addr: "127.0.0.1", Port: 8091})
+ store.upsert("default", &ServiceInstance{Addr: "127.0.0.1", Port: 8091})
+ store.upsert("default", &ServiceInstance{Addr: "127.0.0.2", Port: 8092})
+
+ snapshot := store.Snapshot("default")
+ if len(snapshot) != 2 {
+ t.Fatalf("snapshot length = %d, want 2", len(snapshot))
+ }
+ if updates != 2 {
+ t.Fatalf("updates = %d, want 2", updates)
+ }
+
+ if !store.remove("default", "127.0.0.1", 8091) {
+ t.Fatal("expected remove to delete an instance")
+ }
+ snapshot = store.Snapshot("default")
+ if len(snapshot) != 1 {
+ t.Fatalf("snapshot length after remove = %d, want 1",
len(snapshot))
+ }
+ if snapshot[0].Addr != "127.0.0.2" || snapshot[0].Port != 8092 {
+ t.Fatalf("unexpected instance after remove: %v", snapshot[0])
+ }
+
+ if store.remove("default", "127.0.0.3", 8093) {
+ t.Fatal("remove should report false for a missing instance")
+ }
+ if updates != 3 {
+ t.Fatalf("updates after missing remove = %d, want 3", updates)
+ }
+}
+
+func TestAddressStoreConcurrentUpdateSnapshot(t *testing.T) {
+ store := NewAddressStore()
+ var wg sync.WaitGroup
+
+ for i := 0; i < 8; i++ {
+ wg.Add(1)
+ go func(worker int) {
+ defer wg.Done()
+ for j := 0; j < 100; j++ {
+ store.Update("default", []*ServiceInstance{{
+ Addr: fmt.Sprintf("127.0.0.%d", worker),
+ Port: 8091 + j,
+ }})
+ for _, instance := range
store.Snapshot("default") {
+ if instance == nil {
+ t.Error("snapshot contains nil
instance")
+ }
+ }
+ }
+ }(i)
+ }
+
+ wg.Wait()
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
