feat: registry.

This commit is contained in:
Bobo
2025-06-02 11:53:18 +08:00
parent b0e91998e1
commit 59c731904a
83 changed files with 7154 additions and 1784 deletions

31
registry/zookeeper/go.mod Normal file
View 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
View 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=

View 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
}
}

View 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
}
}
}

View 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 := &registry.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 := &registry.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: &registry.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: &registry.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 := &registry.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 := &registry.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)
}
})
}
}

View 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
}

View 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
}