feat: registry.
This commit is contained in:
31
registry/zookeeper/go.mod
Normal file
31
registry/zookeeper/go.mod
Normal file
@@ -0,0 +1,31 @@
|
||||
module github.com/tx7do/kratos-bootstrap/registry/zookeeper
|
||||
|
||||
go 1.24.0
|
||||
|
||||
toolchain go1.24.3
|
||||
|
||||
replace (
|
||||
github.com/armon/go-metrics => github.com/hashicorp/go-metrics v0.4.1
|
||||
|
||||
github.com/tx7do/kratos-bootstrap/api => ../../api
|
||||
github.com/tx7do/kratos-bootstrap/registry => ../registry
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/go-kratos/kratos/contrib/registry/zookeeper/v2 v2.0.0-20250527152916-d6f5f00cf562
|
||||
github.com/go-kratos/kratos/v2 v2.8.4
|
||||
github.com/go-zookeeper/zk v1.0.4
|
||||
github.com/stretchr/testify v1.10.0
|
||||
github.com/tx7do/kratos-bootstrap/api v0.0.20
|
||||
golang.org/x/sync v0.14.0
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
|
||||
github.com/google/go-cmp v0.7.0 // indirect
|
||||
github.com/kr/text v0.2.0 // indirect
|
||||
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
|
||||
github.com/rogpeppe/go-internal v1.13.1 // indirect
|
||||
google.golang.org/protobuf v1.36.6 // indirect
|
||||
gopkg.in/yaml.v3 v3.0.1 // indirect
|
||||
)
|
||||
30
registry/zookeeper/go.sum
Normal file
30
registry/zookeeper/go.sum
Normal file
@@ -0,0 +1,30 @@
|
||||
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
|
||||
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
|
||||
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/go-kratos/kratos/contrib/registry/zookeeper/v2 v2.0.0-20250527152916-d6f5f00cf562 h1:SBDGsiqkqT6BOidznJKweIcru7eVigX8fkNWp1mZjWk=
|
||||
github.com/go-kratos/kratos/contrib/registry/zookeeper/v2 v2.0.0-20250527152916-d6f5f00cf562/go.mod h1:/XjgaYgSaaoIe4YcUishDuqLpQvOSujQTk8D0cJx/Zo=
|
||||
github.com/go-kratos/kratos/v2 v2.8.4 h1:eIJLE9Qq9WSoKx+Buy2uPyrahtF/lPh+Xf4MTpxhmjs=
|
||||
github.com/go-kratos/kratos/v2 v2.8.4/go.mod h1:mq62W2101a5uYyRxe+7IdWubu7gZCGYqSNKwGFiiRcw=
|
||||
github.com/go-zookeeper/zk v1.0.4 h1:DPzxraQx7OrPyXq2phlGlNSIyWEsAox0RJmjTseMV6I=
|
||||
github.com/go-zookeeper/zk v1.0.4/go.mod h1:nOB03cncLtlp4t+UAkGSV+9beXP/akpekBwL+UX1Qcw=
|
||||
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
|
||||
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
|
||||
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
|
||||
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
|
||||
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
|
||||
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
|
||||
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
|
||||
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII=
|
||||
github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o=
|
||||
github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
|
||||
github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
|
||||
golang.org/x/sync v0.14.0 h1:woo0S4Yywslg6hp4eUFjTVOyKt0RookbpAHG4c1HmhQ=
|
||||
golang.org/x/sync v0.14.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=
|
||||
google.golang.org/protobuf v1.36.6 h1:z1NpPI8ku2WgiWnf+t9wTPsn6eP1L7ksHUlkfLvd9xY=
|
||||
google.golang.org/protobuf v1.36.6/go.mod h1:jduwjTPXsFjZGTmRluh+L6NjiWu7pchiJ2/5YcXBHnY=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
|
||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
|
||||
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
23
registry/zookeeper/options.go
Normal file
23
registry/zookeeper/options.go
Normal file
@@ -0,0 +1,23 @@
|
||||
package zookeeper
|
||||
|
||||
// Option is etcd registry option.
|
||||
type Option func(o *options)
|
||||
|
||||
type options struct {
|
||||
namespace string
|
||||
user string
|
||||
password string
|
||||
}
|
||||
|
||||
// WithRootPath with registry root path.
|
||||
func WithRootPath(path string) Option {
|
||||
return func(o *options) { o.namespace = path }
|
||||
}
|
||||
|
||||
// WithDigestACL with registry password.
|
||||
func WithDigestACL(user string, password string) Option {
|
||||
return func(o *options) {
|
||||
o.user = user
|
||||
o.password = password
|
||||
}
|
||||
}
|
||||
159
registry/zookeeper/register.go
Normal file
159
registry/zookeeper/register.go
Normal file
@@ -0,0 +1,159 @@
|
||||
package zookeeper
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"path"
|
||||
"time"
|
||||
|
||||
"github.com/go-zookeeper/zk"
|
||||
"golang.org/x/sync/singleflight"
|
||||
|
||||
"github.com/go-kratos/kratos/v2/registry"
|
||||
)
|
||||
|
||||
var (
|
||||
_ registry.Registrar = (*Registry)(nil)
|
||||
_ registry.Discovery = (*Registry)(nil)
|
||||
)
|
||||
|
||||
// Registry is consul registry
|
||||
type Registry struct {
|
||||
opts *options
|
||||
conn *zk.Conn
|
||||
|
||||
group singleflight.Group
|
||||
}
|
||||
|
||||
func New(conn *zk.Conn, opts ...Option) *Registry {
|
||||
opt := &options{
|
||||
namespace: "/microservices",
|
||||
}
|
||||
for _, o := range opts {
|
||||
o(opt)
|
||||
}
|
||||
return &Registry{
|
||||
opts: opt,
|
||||
conn: conn,
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Registry) Register(_ context.Context, service *registry.ServiceInstance) error {
|
||||
var (
|
||||
data []byte
|
||||
err error
|
||||
)
|
||||
if err = r.ensureName(r.opts.namespace, []byte(""), 0); err != nil {
|
||||
return err
|
||||
}
|
||||
serviceNamePath := path.Join(r.opts.namespace, service.Name)
|
||||
if err = r.ensureName(serviceNamePath, []byte(""), 0); err != nil {
|
||||
return err
|
||||
}
|
||||
if data, err = marshal(service); err != nil {
|
||||
return err
|
||||
}
|
||||
servicePath := path.Join(serviceNamePath, service.ID)
|
||||
if err = r.ensureName(servicePath, data, zk.FlagEphemeral); err != nil {
|
||||
return err
|
||||
}
|
||||
go r.reRegister(servicePath, data)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Deregister registry service to zookeeper.
|
||||
func (r *Registry) Deregister(ctx context.Context, service *registry.ServiceInstance) error {
|
||||
ch := make(chan error, 1)
|
||||
servicePath := path.Join(r.opts.namespace, service.Name, service.ID)
|
||||
go func() {
|
||||
err := r.conn.Delete(servicePath, -1)
|
||||
ch <- err
|
||||
}()
|
||||
var err error
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
err = ctx.Err()
|
||||
case err = <-ch:
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// GetService get services from zookeeper
|
||||
func (r *Registry) GetService(_ context.Context, serviceName string) ([]*registry.ServiceInstance, error) {
|
||||
instances, err, _ := r.group.Do(serviceName, func() (interface{}, error) {
|
||||
serviceNamePath := path.Join(r.opts.namespace, serviceName)
|
||||
servicesID, _, err := r.conn.Children(serviceNamePath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items := make([]*registry.ServiceInstance, 0, len(servicesID))
|
||||
for _, service := range servicesID {
|
||||
servicePath := path.Join(serviceNamePath, service)
|
||||
serviceInstanceByte, _, err := r.conn.Get(servicePath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
item, err := unmarshal(serviceInstanceByte)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return instances.([]*registry.ServiceInstance), nil
|
||||
}
|
||||
|
||||
func (r *Registry) Watch(ctx context.Context, serviceName string) (registry.Watcher, error) {
|
||||
prefix := path.Join(r.opts.namespace, serviceName)
|
||||
return newWatcher(ctx, prefix, serviceName, r.conn)
|
||||
}
|
||||
|
||||
// ensureName ensure node exists, if not exist, create and set data
|
||||
func (r *Registry) ensureName(path string, data []byte, flags int32) error {
|
||||
exists, stat, err := r.conn.Exists(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// ephemeral nodes handling after restart
|
||||
// fixes a race condition if the server crashes without using CreateProtectedEphemeralSequential()
|
||||
if flags&zk.FlagEphemeral == zk.FlagEphemeral {
|
||||
err = r.conn.Delete(path, stat.Version)
|
||||
if err != nil && !errors.Is(err, zk.ErrNoNode) {
|
||||
return err
|
||||
}
|
||||
exists = false
|
||||
}
|
||||
if !exists {
|
||||
if len(r.opts.user) > 0 && len(r.opts.password) > 0 {
|
||||
_, err = r.conn.Create(path, data, flags, zk.DigestACL(zk.PermAll, r.opts.user, r.opts.password))
|
||||
} else {
|
||||
_, err = r.conn.Create(path, data, flags, zk.WorldACL(zk.PermAll))
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// reRegister re-register data node info when bad connection recovered
|
||||
func (r *Registry) reRegister(path string, data []byte) {
|
||||
sessionID := r.conn.SessionID()
|
||||
ticker := time.NewTicker(time.Second)
|
||||
defer ticker.Stop()
|
||||
for range ticker.C {
|
||||
cur := r.conn.SessionID()
|
||||
// sessionID changed
|
||||
if cur > 0 && sessionID != cur {
|
||||
// re-ensureName
|
||||
if err := r.ensureName(path, data, zk.FlagEphemeral); err != nil {
|
||||
return
|
||||
}
|
||||
sessionID = cur
|
||||
}
|
||||
}
|
||||
}
|
||||
420
registry/zookeeper/register_test.go
Normal file
420
registry/zookeeper/register_test.go
Normal file
@@ -0,0 +1,420 @@
|
||||
package zookeeper
|
||||
|
||||
import (
|
||||
"context"
|
||||
"reflect"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/go-zookeeper/zk"
|
||||
|
||||
"github.com/go-kratos/kratos/v2/registry"
|
||||
)
|
||||
|
||||
func TestRegistry_GetService(t *testing.T) {
|
||||
conn, _, err := zk.Connect([]string{"127.0.0.1:2181"}, time.Second*15)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
return
|
||||
}
|
||||
r := New(conn)
|
||||
|
||||
svrHello := ®istry.ServiceInstance{
|
||||
ID: "1",
|
||||
Name: "hello",
|
||||
Version: "v1.0.0",
|
||||
Endpoints: []string{"127.0.0.1:8080"},
|
||||
}
|
||||
|
||||
type fields struct {
|
||||
registry *Registry
|
||||
}
|
||||
type args struct {
|
||||
ctx context.Context
|
||||
serviceName string
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
fields fields
|
||||
args args
|
||||
want []*registry.ServiceInstance
|
||||
wantErr bool
|
||||
preFunc func(t *testing.T)
|
||||
deferFunc func(t *testing.T)
|
||||
}{
|
||||
{
|
||||
name: "normal",
|
||||
preFunc: func(t *testing.T) {
|
||||
err = r.Register(context.Background(), svrHello)
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
},
|
||||
deferFunc: func(t *testing.T) {
|
||||
err = r.Deregister(context.Background(), svrHello)
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
},
|
||||
fields: fields{
|
||||
registry: r,
|
||||
},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
serviceName: svrHello.Name,
|
||||
},
|
||||
want: []*registry.ServiceInstance{svrHello},
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "can't get any",
|
||||
fields: fields{
|
||||
registry: r,
|
||||
},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
serviceName: "helloxxx",
|
||||
},
|
||||
want: nil,
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "conn close",
|
||||
preFunc: func(t *testing.T) {
|
||||
conn.Close()
|
||||
},
|
||||
fields: fields{
|
||||
registry: r,
|
||||
},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
serviceName: "hello",
|
||||
},
|
||||
want: nil,
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if tt.preFunc != nil {
|
||||
tt.preFunc(t)
|
||||
}
|
||||
if tt.deferFunc != nil {
|
||||
defer tt.deferFunc(t)
|
||||
}
|
||||
r := tt.fields.registry
|
||||
got, err := r.GetService(tt.args.ctx, tt.args.serviceName)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("GetService() error = %v, wantErr %v", err, tt.wantErr)
|
||||
t.Errorf("GetService() got = %v", got)
|
||||
return
|
||||
}
|
||||
if !reflect.DeepEqual(got, tt.want) {
|
||||
t.Errorf("GetService() got = %v, want %v", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_Register(t *testing.T) {
|
||||
conn, _, err := zk.Connect([]string{"127.0.0.1:2181"}, time.Second*15)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
return
|
||||
}
|
||||
r := New(conn)
|
||||
|
||||
svrHello := ®istry.ServiceInstance{
|
||||
ID: "1",
|
||||
Name: "hello",
|
||||
Version: "v1.0.0",
|
||||
Endpoints: []string{"127.0.0.1:8080"},
|
||||
}
|
||||
|
||||
type fields struct {
|
||||
registry *Registry
|
||||
}
|
||||
type args struct {
|
||||
ctx context.Context
|
||||
service *registry.ServiceInstance
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
fields fields
|
||||
args args
|
||||
wantErr bool
|
||||
preFunc func(t *testing.T)
|
||||
deferFunc func(t *testing.T)
|
||||
}{
|
||||
{
|
||||
name: "normal",
|
||||
fields: fields{
|
||||
registry: r,
|
||||
},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
service: svrHello,
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "invalid path",
|
||||
fields: fields{
|
||||
registry: New(conn, WithRootPath("invalid")),
|
||||
},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
service: ®istry.ServiceInstance{
|
||||
ID: "1",
|
||||
Name: "hello1",
|
||||
Version: "v1.0.0",
|
||||
Endpoints: []string{"127.0.0.1:8080"},
|
||||
},
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "auth",
|
||||
preFunc: func(t *testing.T) {
|
||||
err = conn.AddAuth("digest", []byte("test:test"))
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
},
|
||||
fields: fields{
|
||||
registry: New(conn, WithRootPath("/tt1"), WithDigestACL("test", "test")),
|
||||
},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
service: ®istry.ServiceInstance{
|
||||
ID: "1",
|
||||
Name: "hello2",
|
||||
Version: "v1.0.0",
|
||||
Endpoints: []string{"127.0.0.1:8080"},
|
||||
},
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if tt.preFunc != nil {
|
||||
tt.preFunc(t)
|
||||
}
|
||||
if tt.deferFunc != nil {
|
||||
defer tt.deferFunc(t)
|
||||
}
|
||||
r := tt.fields.registry
|
||||
if err := r.Register(tt.args.ctx, tt.args.service); (err != nil) != tt.wantErr {
|
||||
t.Errorf("Register() error = %v, wantErr %v", err, tt.wantErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_Deregister(t *testing.T) {
|
||||
conn, _, err := zk.Connect([]string{"127.0.0.1:2181"}, time.Second*15)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
return
|
||||
}
|
||||
r := New(conn)
|
||||
|
||||
svrHello := ®istry.ServiceInstance{
|
||||
ID: "1",
|
||||
Name: "hello",
|
||||
Version: "v1.0.0",
|
||||
Endpoints: []string{"127.0.0.1:8080"},
|
||||
}
|
||||
|
||||
cancelCtx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
type fields struct {
|
||||
registry *Registry
|
||||
}
|
||||
type args struct {
|
||||
ctx context.Context
|
||||
service *registry.ServiceInstance
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
fields fields
|
||||
args args
|
||||
wantErr bool
|
||||
preFunc func(t *testing.T)
|
||||
deferFunc func(t *testing.T)
|
||||
}{
|
||||
{
|
||||
name: "normal",
|
||||
fields: fields{
|
||||
registry: r,
|
||||
},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
service: svrHello,
|
||||
},
|
||||
wantErr: false,
|
||||
preFunc: func(t *testing.T) {
|
||||
err = r.Register(context.Background(), svrHello)
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "with ctx cancel",
|
||||
fields: fields{
|
||||
registry: r,
|
||||
},
|
||||
args: args{
|
||||
ctx: cancelCtx,
|
||||
service: svrHello,
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
r := tt.fields.registry
|
||||
if err := r.Deregister(tt.args.ctx, tt.args.service); (err != nil) != tt.wantErr {
|
||||
t.Errorf("Deregister() error = %v, wantErr %v", err, tt.wantErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_Watch(t *testing.T) {
|
||||
conn, _, err := zk.Connect([]string{"127.0.0.1:2181"}, time.Second*15)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
return
|
||||
}
|
||||
closeConn, _, err := zk.Connect([]string{"127.0.0.1:2181"}, time.Second*15)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
return
|
||||
}
|
||||
r := New(conn)
|
||||
|
||||
svrHello := ®istry.ServiceInstance{
|
||||
ID: "1",
|
||||
Name: "hello",
|
||||
Version: "v1.0.0",
|
||||
Endpoints: []string{"127.0.0.1:8080"},
|
||||
}
|
||||
|
||||
cancelCtx, cancel := context.WithCancel(context.Background())
|
||||
type fields struct {
|
||||
registry *Registry
|
||||
}
|
||||
type args struct {
|
||||
ctx context.Context
|
||||
serviceName string
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
fields fields
|
||||
args args
|
||||
wantErr bool
|
||||
want []*registry.ServiceInstance
|
||||
|
||||
preFunc func(t *testing.T)
|
||||
deferFunc func(t *testing.T)
|
||||
processFunc func(t *testing.T, w registry.Watcher)
|
||||
}{
|
||||
{
|
||||
name: "normal",
|
||||
fields: fields{
|
||||
registry: r,
|
||||
},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
serviceName: svrHello.Name,
|
||||
},
|
||||
wantErr: false,
|
||||
want: []*registry.ServiceInstance{svrHello},
|
||||
deferFunc: func(t *testing.T) {
|
||||
err = r.Deregister(context.Background(), svrHello)
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
},
|
||||
processFunc: func(t *testing.T, w registry.Watcher) {
|
||||
err = r.Register(context.Background(), svrHello)
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "ctx cancel",
|
||||
fields: fields{
|
||||
registry: r,
|
||||
},
|
||||
args: args{
|
||||
ctx: cancelCtx,
|
||||
serviceName: svrHello.Name,
|
||||
},
|
||||
wantErr: true,
|
||||
want: nil,
|
||||
processFunc: func(t *testing.T, w registry.Watcher) {
|
||||
cancel()
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "disconnect",
|
||||
fields: fields{
|
||||
registry: New(closeConn),
|
||||
},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
serviceName: svrHello.Name,
|
||||
},
|
||||
wantErr: true,
|
||||
want: nil,
|
||||
processFunc: func(t *testing.T, w registry.Watcher) {
|
||||
closeConn.Close()
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if tt.preFunc != nil {
|
||||
tt.preFunc(t)
|
||||
}
|
||||
if tt.deferFunc != nil {
|
||||
defer tt.deferFunc(t)
|
||||
}
|
||||
r := tt.fields.registry
|
||||
watcher, err := r.Watch(tt.args.ctx, tt.args.serviceName)
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
err = watcher.Stop()
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
}()
|
||||
_, err = watcher.Next()
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
return
|
||||
}
|
||||
|
||||
if tt.processFunc != nil {
|
||||
tt.processFunc(t, watcher)
|
||||
}
|
||||
|
||||
want, err := watcher.Next()
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("Watch() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if !reflect.DeepEqual(want, tt.want) {
|
||||
t.Errorf("Watch() watcher = %v, want %v", watcher, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
16
registry/zookeeper/service.go
Normal file
16
registry/zookeeper/service.go
Normal file
@@ -0,0 +1,16 @@
|
||||
package zookeeper
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
|
||||
"github.com/go-kratos/kratos/v2/registry"
|
||||
)
|
||||
|
||||
func marshal(si *registry.ServiceInstance) ([]byte, error) {
|
||||
return json.Marshal(si)
|
||||
}
|
||||
|
||||
func unmarshal(data []byte) (si *registry.ServiceInstance, err error) {
|
||||
err = json.Unmarshal(data, &si)
|
||||
return
|
||||
}
|
||||
111
registry/zookeeper/watcher.go
Normal file
111
registry/zookeeper/watcher.go
Normal file
@@ -0,0 +1,111 @@
|
||||
package zookeeper
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"path"
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/go-zookeeper/zk"
|
||||
|
||||
"github.com/go-kratos/kratos/v2/registry"
|
||||
)
|
||||
|
||||
var _ registry.Watcher = (*watcher)(nil)
|
||||
|
||||
var ErrWatcherStopped = errors.New("watcher stopped")
|
||||
|
||||
type watcher struct {
|
||||
ctx context.Context
|
||||
event chan zk.Event
|
||||
conn *zk.Conn
|
||||
cancel context.CancelFunc
|
||||
|
||||
first uint32
|
||||
// 前缀
|
||||
prefix string
|
||||
// watch 的服务名
|
||||
serviceName string
|
||||
}
|
||||
|
||||
func newWatcher(ctx context.Context, prefix, serviceName string, conn *zk.Conn) (*watcher, error) {
|
||||
w := &watcher{conn: conn, event: make(chan zk.Event, 1), prefix: prefix, serviceName: serviceName}
|
||||
w.ctx, w.cancel = context.WithCancel(ctx)
|
||||
go w.watch(w.ctx)
|
||||
return w, nil
|
||||
}
|
||||
|
||||
func (w *watcher) watch(ctx context.Context) {
|
||||
for {
|
||||
// 每次 watch 只有一次有效期 所以循环 watch
|
||||
_, _, ch, err := w.conn.ChildrenW(w.prefix)
|
||||
if err != nil {
|
||||
// If the target service node has not been created
|
||||
if errors.Is(err, zk.ErrNoNode) {
|
||||
// Add watcher for the node exists
|
||||
_, _, ch, err = w.conn.ExistsW(w.prefix)
|
||||
}
|
||||
if err != nil {
|
||||
w.event <- zk.Event{Err: err}
|
||||
return
|
||||
}
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case ev := <-ch:
|
||||
w.event <- ev
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (w *watcher) Next() ([]*registry.ServiceInstance, error) {
|
||||
// todo 如果多处调用 next 可能会导致多实例信息不同步
|
||||
if atomic.CompareAndSwapUint32(&w.first, 0, 1) {
|
||||
return w.getServices()
|
||||
}
|
||||
select {
|
||||
case <-w.ctx.Done():
|
||||
return nil, w.ctx.Err()
|
||||
case e := <-w.event:
|
||||
if e.State == zk.StateDisconnected {
|
||||
return nil, ErrWatcherStopped
|
||||
}
|
||||
if e.Err != nil {
|
||||
return nil, e.Err
|
||||
}
|
||||
return w.getServices()
|
||||
}
|
||||
}
|
||||
|
||||
func (w *watcher) Stop() error {
|
||||
w.cancel()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *watcher) getServices() ([]*registry.ServiceInstance, error) {
|
||||
servicesID, _, err := w.conn.Children(w.prefix)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items := make([]*registry.ServiceInstance, 0, len(servicesID))
|
||||
for _, id := range servicesID {
|
||||
servicePath := path.Join(w.prefix, id)
|
||||
b, _, err := w.conn.Get(servicePath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
item, err := unmarshal(b)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 与 watch 的服务名不同 则跳过
|
||||
if item.Name != w.serviceName {
|
||||
continue
|
||||
}
|
||||
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
Reference in New Issue
Block a user