✨ feat(etcd): 添加 etcd 服务注册和 gRPC 客户端支持
- 新增 `etcd_registry.go` 文件，实现带自动重连的 etcd 服务注册功能 - 新增 `grpc_etcd.go` 文件，提供基于 etcd 的 gRPC 客户端，支持服务发现和重连机制 - 更新 `go.mod` 文件，添加 `github.com/rabbitmq/amqp091-go` 依赖 - 实现了服务注册、注销、续约及健康检查等功能，增强了连接管理能力
This commit is contained in:
200
etcd_registry.go
Normal file
200
etcd_registry.go
Normal file
@@ -0,0 +1,200 @@
|
||||
package reconnect
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
clientv3 "go.etcd.io/etcd/client/v3"
|
||||
)
|
||||
|
||||
// EtcdRegistryConfig etcd 服务注册配置
|
||||
type EtcdRegistryConfig struct {
|
||||
// Endpoints etcd 服务地址列表
|
||||
Endpoints []string
|
||||
// DialTimeout 连接超时时间
|
||||
DialTimeout time.Duration
|
||||
// ReconnectConfig 重连配置
|
||||
ReconnectConfig Config
|
||||
}
|
||||
|
||||
// EtcdRegistry etcd 服务注册（带自动重连）
|
||||
type EtcdRegistry struct {
|
||||
config EtcdRegistryConfig
|
||||
etcdClient *EtcdClient
|
||||
cli *clientv3.Client
|
||||
key string
|
||||
val string
|
||||
serviceName string
|
||||
ttl int64
|
||||
leaseID clientv3.LeaseID
|
||||
keepAliveCh <-chan *clientv3.LeaseKeepAliveResponse
|
||||
mu sync.RWMutex
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
}
|
||||
|
||||
// NewEtcdRegistry 创建带重连功能的 etcd 服务注册器
|
||||
func NewEtcdRegistry(cfg EtcdRegistryConfig) (*EtcdRegistry, error) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
// 创建带重连的 etcd 客户端
|
||||
etcdCfg := EtcdClientConfig{
|
||||
Endpoints: cfg.Endpoints,
|
||||
DialTimeout: cfg.DialTimeout,
|
||||
ReconnectConfig: cfg.ReconnectConfig,
|
||||
}
|
||||
|
||||
etcdClient, err := NewEtcdClient(etcdCfg)
|
||||
if err != nil {
|
||||
cancel()
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &EtcdRegistry{
|
||||
config: cfg,
|
||||
etcdClient: etcdClient,
|
||||
cli: etcdClient.GetClient(),
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Register 注册服务
|
||||
func (r *EtcdRegistry) Register(serviceName, addr string, ttl int64) error {
|
||||
r.mu.Lock()
|
||||
r.serviceName = serviceName
|
||||
r.ttl = ttl
|
||||
r.key = fmt.Sprintf("/%s/%s", serviceName, addr)
|
||||
r.val = addr
|
||||
r.mu.Unlock()
|
||||
|
||||
return r.registerWithKV(ttl)
|
||||
}
|
||||
|
||||
func (r *EtcdRegistry) registerWithKV(ttl int64) error {
|
||||
r.mu.RLock()
|
||||
cli := r.cli
|
||||
key := r.key
|
||||
val := r.val
|
||||
r.mu.RUnlock()
|
||||
|
||||
if cli == nil {
|
||||
// 尝试重新获取客户端
|
||||
if r.etcdClient != nil {
|
||||
r.mu.Lock()
|
||||
r.cli = r.etcdClient.GetClient()
|
||||
cli = r.cli
|
||||
r.mu.Unlock()
|
||||
}
|
||||
if cli == nil {
|
||||
return fmt.Errorf("etcd client is nil")
|
||||
}
|
||||
}
|
||||
|
||||
resp, err := cli.Grant(context.Background(), ttl)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
r.mu.Lock()
|
||||
r.leaseID = resp.ID
|
||||
r.mu.Unlock()
|
||||
|
||||
_, err = cli.Put(context.Background(), key, val, clientv3.WithLease(resp.ID))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
keepAliveCh, err := cli.KeepAlive(context.Background(), resp.ID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
r.mu.Lock()
|
||||
r.keepAliveCh = keepAliveCh
|
||||
r.mu.Unlock()
|
||||
|
||||
go r.watcher(ttl)
|
||||
log.Printf("[EtcdRegistry] 服务注册成功: %s", key)
|
||||
return nil
|
||||
}
|
||||
|
||||
// watcher 监听续约
|
||||
func (r *EtcdRegistry) watcher(ttl int64) {
|
||||
for {
|
||||
select {
|
||||
case <-r.ctx.Done():
|
||||
log.Println("[EtcdRegistry] context done, watcher exiting.")
|
||||
return
|
||||
case ka, ok := <-r.keepAliveCh:
|
||||
if !ok {
|
||||
log.Println("[EtcdRegistry] keep-alive channel closed, attempting to re-register...")
|
||||
|
||||
// 更新 cli 引用
|
||||
if r.etcdClient != nil {
|
||||
r.mu.Lock()
|
||||
r.cli = r.etcdClient.GetClient()
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
// 使用指数退避重试
|
||||
delay := time.Second
|
||||
maxDelay := 30 * time.Second
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-r.ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
if err := r.registerWithKV(ttl); err != nil {
|
||||
log.Printf("[EtcdRegistry] failed to re-register service: %v, retrying in %v...", err, delay)
|
||||
time.Sleep(delay)
|
||||
delay *= 2
|
||||
if delay > maxDelay {
|
||||
delay = maxDelay
|
||||
}
|
||||
} else {
|
||||
return // 退出当前 goroutine，让新的 watcher 启动
|
||||
}
|
||||
}
|
||||
}
|
||||
// 续约成功
|
||||
_ = ka
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// UnRegister 注销服务
|
||||
func (r *EtcdRegistry) UnRegister() {
|
||||
r.cancel() // 停止 watcher
|
||||
|
||||
r.mu.RLock()
|
||||
cli := r.cli
|
||||
leaseID := r.leaseID
|
||||
key := r.key
|
||||
r.mu.RUnlock()
|
||||
|
||||
if cli != nil {
|
||||
cli.Revoke(context.Background(), leaseID)
|
||||
cli.Delete(context.Background(), key)
|
||||
log.Printf("[EtcdRegistry] 服务注销成功: %s", key)
|
||||
}
|
||||
|
||||
if r.etcdClient != nil {
|
||||
r.etcdClient.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// State 获取连接状态
|
||||
func (r *EtcdRegistry) State() ConnectionState {
|
||||
if r.etcdClient != nil {
|
||||
return r.etcdClient.State()
|
||||
}
|
||||
return StateDisconnected
|
||||
}
|
||||
|
||||
2
go.mod
2
go.mod
@@ -4,6 +4,7 @@ go 1.24.0
|
||||
|
||||
require (
|
||||
github.com/go-redis/redis/v8 v8.11.5
|
||||
github.com/rabbitmq/amqp091-go v1.10.0
|
||||
go.etcd.io/etcd/client/v3 v3.6.5
|
||||
google.golang.org/grpc v1.75.1
|
||||
gorm.io/gorm v1.31.1
|
||||
@@ -19,7 +20,6 @@ require (
|
||||
github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.3 // indirect
|
||||
github.com/jinzhu/inflection v1.0.0 // indirect
|
||||
github.com/jinzhu/now v1.1.5 // indirect
|
||||
github.com/rabbitmq/amqp091-go v1.10.0 // indirect
|
||||
go.etcd.io/etcd/api/v3 v3.6.5 // indirect
|
||||
go.etcd.io/etcd/client/pkg/v3 v3.6.5 // indirect
|
||||
go.uber.org/multierr v1.11.0 // indirect
|
||||
|
||||
274
grpc_etcd.go
Normal file
274
grpc_etcd.go
Normal file
@@ -0,0 +1,274 @@
|
||||
package reconnect
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
clientv3 "go.etcd.io/etcd/client/v3"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/connectivity"
|
||||
"google.golang.org/grpc/credentials/insecure"
|
||||
"google.golang.org/grpc/resolver"
|
||||
)
|
||||
|
||||
// GRPCEtcdClientConfig 使用 etcd 服务发现的 gRPC 客户端配置
|
||||
type GRPCEtcdClientConfig struct {
|
||||
// EtcdEndpoints etcd 服务地址列表
|
||||
EtcdEndpoints []string
|
||||
// EtcdDialTimeout etcd 连接超时时间
|
||||
EtcdDialTimeout time.Duration
|
||||
// ServiceName 服务名称（在 etcd 中注册的 key 前缀）
|
||||
ServiceName string
|
||||
// DialOptions 额外的 gRPC 拨号选项
|
||||
DialOptions []grpc.DialOption
|
||||
// ReconnectConfig 重连配置
|
||||
ReconnectConfig Config
|
||||
}
|
||||
|
||||
// GRPCEtcdClient 带 etcd 服务发现和重连功能的 gRPC 客户端
|
||||
type GRPCEtcdClient struct {
|
||||
config GRPCEtcdClientConfig
|
||||
etcdClient *clientv3.Client
|
||||
conn *grpc.ClientConn
|
||||
manager *ConnectionManager
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// grpcEtcdConnector gRPC + etcd 连接器
|
||||
type grpcEtcdConnector struct {
|
||||
client *GRPCEtcdClient
|
||||
}
|
||||
|
||||
func (g *grpcEtcdConnector) Connect(ctx context.Context) error {
|
||||
g.client.mu.Lock()
|
||||
defer g.client.mu.Unlock()
|
||||
|
||||
cfg := g.client.config
|
||||
|
||||
// 创建 etcd 客户端
|
||||
etcdTimeout := cfg.EtcdDialTimeout
|
||||
if etcdTimeout == 0 {
|
||||
etcdTimeout = 5 * time.Second
|
||||
}
|
||||
|
||||
cli, err := clientv3.New(clientv3.Config{
|
||||
Endpoints: cfg.EtcdEndpoints,
|
||||
DialTimeout: etcdTimeout,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 注册 etcd resolver
|
||||
etcdResolver := &etcdResolverBuilder{cli: cli}
|
||||
resolver.Register(etcdResolver)
|
||||
|
||||
target := "etcd:///" + cfg.ServiceName + "/"
|
||||
|
||||
// 默认 DialOptions
|
||||
opts := []grpc.DialOption{
|
||||
grpc.WithResolvers(etcdResolver),
|
||||
grpc.WithTransportCredentials(insecure.NewCredentials()),
|
||||
grpc.WithDefaultServiceConfig(`{"loadBalancingPolicy":"round_robin"}`),
|
||||
}
|
||||
// 追加用户自定义选项
|
||||
opts = append(opts, cfg.DialOptions...)
|
||||
|
||||
conn, err := grpc.NewClient(target, opts...)
|
||||
if err != nil {
|
||||
cli.Close()
|
||||
return err
|
||||
}
|
||||
|
||||
// 等待连接就绪
|
||||
connCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
conn.Connect()
|
||||
for {
|
||||
state := conn.GetState()
|
||||
if state == connectivity.Ready {
|
||||
break
|
||||
}
|
||||
if state == connectivity.TransientFailure || state == connectivity.Shutdown {
|
||||
conn.Close()
|
||||
cli.Close()
|
||||
return context.DeadlineExceeded
|
||||
}
|
||||
if !conn.WaitForStateChange(connCtx, state) {
|
||||
conn.Close()
|
||||
cli.Close()
|
||||
return connCtx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
g.client.etcdClient = cli
|
||||
g.client.conn = conn
|
||||
return nil
|
||||
}
|
||||
|
||||
func (g *grpcEtcdConnector) Close() error {
|
||||
g.client.mu.Lock()
|
||||
defer g.client.mu.Unlock()
|
||||
|
||||
if g.client.conn != nil {
|
||||
g.client.conn.Close()
|
||||
g.client.conn = nil
|
||||
}
|
||||
if g.client.etcdClient != nil {
|
||||
g.client.etcdClient.Close()
|
||||
g.client.etcdClient = nil
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// grpcEtcdHealthChecker gRPC + etcd 健康检查器
|
||||
type grpcEtcdHealthChecker struct {
|
||||
client *GRPCEtcdClient
|
||||
}
|
||||
|
||||
func (g *grpcEtcdHealthChecker) HealthCheck(ctx context.Context) error {
|
||||
g.client.mu.RLock()
|
||||
conn := g.client.conn
|
||||
g.client.mu.RUnlock()
|
||||
|
||||
if conn == nil {
|
||||
return context.Canceled
|
||||
}
|
||||
|
||||
state := conn.GetState()
|
||||
if state == connectivity.Ready || state == connectivity.Idle {
|
||||
return nil
|
||||
}
|
||||
|
||||
conn.Connect()
|
||||
if !conn.WaitForStateChange(ctx, state) {
|
||||
return ctx.Err()
|
||||
}
|
||||
|
||||
newState := conn.GetState()
|
||||
if newState != connectivity.Ready && newState != connectivity.Idle {
|
||||
return context.DeadlineExceeded
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// NewGRPCClientWithEtcd 创建带 etcd 服务发现和重连功能的 gRPC 客户端
|
||||
func NewGRPCClientWithEtcd(cfg GRPCEtcdClientConfig) (*GRPCEtcdClient, error) {
|
||||
client := &GRPCEtcdClient{
|
||||
config: cfg,
|
||||
}
|
||||
|
||||
connector := &grpcEtcdConnector{client: client}
|
||||
checker := &grpcEtcdHealthChecker{client: client}
|
||||
|
||||
client.manager = NewConnectionManager(connector, checker, cfg.ReconnectConfig)
|
||||
|
||||
// 首次连接
|
||||
ctx := context.Background()
|
||||
if err := client.manager.ConnectWithRetry(ctx); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return client, nil
|
||||
}
|
||||
|
||||
// GetConn 获取 gRPC 连接
|
||||
func (c *GRPCEtcdClient) GetConn() *grpc.ClientConn {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
return c.conn
|
||||
}
|
||||
|
||||
// GetEtcdClient 获取 etcd 客户端
|
||||
func (c *GRPCEtcdClient) GetEtcdClient() *clientv3.Client {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
return c.etcdClient
|
||||
}
|
||||
|
||||
// State 获取连接状态
|
||||
func (c *GRPCEtcdClient) State() ConnectionState {
|
||||
return c.manager.State()
|
||||
}
|
||||
|
||||
// Close 关闭客户端
|
||||
func (c *GRPCEtcdClient) Close() error {
|
||||
return c.manager.Close()
|
||||
}
|
||||
|
||||
// TriggerReconnect 手动触发重连
|
||||
func (c *GRPCEtcdClient) TriggerReconnect() {
|
||||
c.manager.TriggerReconnect()
|
||||
}
|
||||
|
||||
// ================== etcd resolver 实现 ==================
|
||||
|
||||
// etcdResolverBuilder 实现 resolver.Builder
|
||||
type etcdResolverBuilder struct {
|
||||
cli *clientv3.Client
|
||||
}
|
||||
|
||||
func (b *etcdResolverBuilder) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
|
||||
r := &etcdResolver{
|
||||
cli: b.cli,
|
||||
cc: cc,
|
||||
target: target,
|
||||
}
|
||||
go r.watch(target.Endpoint())
|
||||
return r, nil
|
||||
}
|
||||
|
||||
func (b *etcdResolverBuilder) Scheme() string {
|
||||
return "etcd"
|
||||
}
|
||||
|
||||
// etcdResolver 实现 resolver.Resolver
|
||||
type etcdResolver struct {
|
||||
cli *clientv3.Client
|
||||
cc resolver.ClientConn
|
||||
target resolver.Target
|
||||
}
|
||||
|
||||
func (r *etcdResolver) watch(keyPrefix string) {
|
||||
// 首次获取服务列表
|
||||
resp, err := r.cli.Get(context.Background(), "/"+keyPrefix, clientv3.WithPrefix())
|
||||
if err != nil {
|
||||
log.Printf("[etcdResolver] 获取服务列表失败: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
addrs := make([]resolver.Address, 0, len(resp.Kvs))
|
||||
for _, kv := range resp.Kvs {
|
||||
addrs = append(addrs, resolver.Address{Addr: string(kv.Value)})
|
||||
}
|
||||
r.cc.UpdateState(resolver.State{Addresses: addrs})
|
||||
|
||||
// 监听变化
|
||||
watchCh := r.cli.Watch(context.Background(), "/"+keyPrefix, clientv3.WithPrefix())
|
||||
for wresp := range watchCh {
|
||||
for _, ev := range wresp.Events {
|
||||
switch ev.Type {
|
||||
case clientv3.EventTypePut:
|
||||
addrs = append(addrs, resolver.Address{Addr: string(ev.Kv.Value)})
|
||||
case clientv3.EventTypeDelete:
|
||||
newAddrs := make([]resolver.Address, 0)
|
||||
for _, addr := range addrs {
|
||||
if addr.Addr != string(ev.Kv.Value) {
|
||||
newAddrs = append(newAddrs, addr)
|
||||
}
|
||||
}
|
||||
addrs = newAddrs
|
||||
}
|
||||
r.cc.UpdateState(resolver.State{Addresses: addrs})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (r *etcdResolver) ResolveNow(opts resolver.ResolveNowOptions) {}
|
||||
|
||||
func (r *etcdResolver) Close() {}
|
||||
|
||||
Reference in New Issue
Block a user