skiffer-git
6 months ago
parent
73824cc5f2
commit
f16b3c50bd
@ -1,289 +0,0 @@
|
|||||||
package etcd
|
|
||||||
|
|
||||||
import (
|
|
||||||
"context"
|
|
||||||
"fmt"
|
|
||||||
"github.com/pkg/errors"
|
|
||||||
clientv3 "go.etcd.io/etcd/client/v3"
|
|
||||||
"go.etcd.io/etcd/client/v3/naming/endpoints"
|
|
||||||
"go.etcd.io/etcd/client/v3/naming/resolver"
|
|
||||||
"google.golang.org/grpc"
|
|
||||||
gresolver "google.golang.org/grpc/resolver"
|
|
||||||
"strings"
|
|
||||||
"sync"
|
|
||||||
"time"
|
|
||||||
)
|
|
||||||
|
|
||||||
// ZkOption defines a function type for modifying clientv3.Config
|
|
||||||
type ZkOption func(*clientv3.Config)
|
|
||||||
|
|
||||||
// SvcDiscoveryRegistryImpl implementation
|
|
||||||
type SvcDiscoveryRegistryImpl struct {
|
|
||||||
client *clientv3.Client
|
|
||||||
resolver gresolver.Builder
|
|
||||||
dialOptions []grpc.DialOption
|
|
||||||
serviceKey string
|
|
||||||
endpointMgr endpoints.Manager
|
|
||||||
leaseID clientv3.LeaseID
|
|
||||||
rpcRegisterTarget string
|
|
||||||
|
|
||||||
rootDirectory string
|
|
||||||
|
|
||||||
mu sync.RWMutex
|
|
||||||
connMap map[string][]*grpc.ClientConn
|
|
||||||
}
|
|
||||||
|
|
||||||
// NewSvcDiscoveryRegistry creates a new service discovery registry implementation
|
|
||||||
func NewSvcDiscoveryRegistry(rootDirectory string, endpoints []string, options ...ZkOption) (*SvcDiscoveryRegistryImpl, error) {
|
|
||||||
cfg := clientv3.Config{
|
|
||||||
Endpoints: endpoints,
|
|
||||||
DialTimeout: 5 * time.Second,
|
|
||||||
// Increase keep-alive queue capacity and message size
|
|
||||||
PermitWithoutStream: true,
|
|
||||||
MaxCallSendMsgSize: 10 * 1024 * 1024, // 10 MB
|
|
||||||
}
|
|
||||||
|
|
||||||
// Apply provided options to the config
|
|
||||||
for _, opt := range options {
|
|
||||||
opt(&cfg)
|
|
||||||
}
|
|
||||||
|
|
||||||
client, err := clientv3.New(cfg)
|
|
||||||
if err != nil {
|
|
||||||
return nil, err
|
|
||||||
}
|
|
||||||
r, err := resolver.NewBuilder(client)
|
|
||||||
if err != nil {
|
|
||||||
return nil, err
|
|
||||||
}
|
|
||||||
|
|
||||||
s := &SvcDiscoveryRegistryImpl{
|
|
||||||
client: client,
|
|
||||||
resolver: r,
|
|
||||||
rootDirectory: rootDirectory,
|
|
||||||
connMap: make(map[string][]*grpc.ClientConn),
|
|
||||||
}
|
|
||||||
go s.watchServiceChanges()
|
|
||||||
return s, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
// WithDialTimeout sets a custom dial timeout for the etcd client
|
|
||||||
func WithDialTimeout(timeout time.Duration) ZkOption {
|
|
||||||
return func(cfg *clientv3.Config) {
|
|
||||||
cfg.DialTimeout = timeout
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// WithMaxCallSendMsgSize sets a custom max call send message size for the etcd client
|
|
||||||
func WithMaxCallSendMsgSize(size int) ZkOption {
|
|
||||||
return func(cfg *clientv3.Config) {
|
|
||||||
cfg.MaxCallSendMsgSize = size
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// WithUsernameAndPassword sets a username and password for the etcd client
|
|
||||||
func WithUsernameAndPassword(username, password string) ZkOption {
|
|
||||||
return func(cfg *clientv3.Config) {
|
|
||||||
cfg.Username = username
|
|
||||||
cfg.Password = password
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// GetUserIdHashGatewayHost returns the gateway host for a given user ID hash
|
|
||||||
func (r *SvcDiscoveryRegistryImpl) GetUserIdHashGatewayHost(ctx context.Context, userId string) (string, error) {
|
|
||||||
return "", nil
|
|
||||||
}
|
|
||||||
|
|
||||||
// GetConns returns gRPC client connections for a given service name
|
|
||||||
func (r *SvcDiscoveryRegistryImpl) GetConns(ctx context.Context, serviceName string, opts ...grpc.DialOption) ([]*grpc.ClientConn, error) {
|
|
||||||
fullServiceKey := fmt.Sprintf("%s/%s", r.rootDirectory, serviceName)
|
|
||||||
r.mu.RLock()
|
|
||||||
defer r.mu.RUnlock()
|
|
||||||
return r.connMap[fullServiceKey], nil
|
|
||||||
}
|
|
||||||
|
|
||||||
// GetConn returns a single gRPC client connection for a given service name
|
|
||||||
func (r *SvcDiscoveryRegistryImpl) GetConn(ctx context.Context, serviceName string, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
|
|
||||||
target := fmt.Sprintf("etcd:///%s/%s", r.rootDirectory, serviceName)
|
|
||||||
return grpc.DialContext(ctx, target, append(append(r.dialOptions, opts...), grpc.WithResolvers(r.resolver))...)
|
|
||||||
}
|
|
||||||
|
|
||||||
// GetSelfConnTarget returns the connection target for the current service
|
|
||||||
func (r *SvcDiscoveryRegistryImpl) GetSelfConnTarget() string {
|
|
||||||
return r.rpcRegisterTarget
|
|
||||||
}
|
|
||||||
|
|
||||||
// AddOption appends gRPC dial options to the existing options
|
|
||||||
func (r *SvcDiscoveryRegistryImpl) AddOption(opts ...grpc.DialOption) {
|
|
||||||
r.dialOptions = append(r.dialOptions, opts...)
|
|
||||||
}
|
|
||||||
|
|
||||||
// CloseConn closes a given gRPC client connection
|
|
||||||
func (r *SvcDiscoveryRegistryImpl) CloseConn(conn *grpc.ClientConn) {
|
|
||||||
conn.Close()
|
|
||||||
}
|
|
||||||
|
|
||||||
// Register registers a new service endpoint with etcd
|
|
||||||
func (r *SvcDiscoveryRegistryImpl) Register(serviceName, host string, port int, opts ...grpc.DialOption) error {
|
|
||||||
r.serviceKey = fmt.Sprintf("%s/%s/%s:%d", r.rootDirectory, serviceName, host, port)
|
|
||||||
em, err := endpoints.NewManager(r.client, r.rootDirectory+"/"+serviceName)
|
|
||||||
if err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
r.endpointMgr = em
|
|
||||||
|
|
||||||
leaseResp, err := r.client.Grant(context.Background(), 30) //
|
|
||||||
if err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
r.leaseID = leaseResp.ID
|
|
||||||
|
|
||||||
r.rpcRegisterTarget = fmt.Sprintf("%s:%d", host, port)
|
|
||||||
endpoint := endpoints.Endpoint{Addr: r.rpcRegisterTarget}
|
|
||||||
|
|
||||||
err = em.AddEndpoint(context.TODO(), r.serviceKey, endpoint, clientv3.WithLease(leaseResp.ID))
|
|
||||||
if err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
|
|
||||||
go r.keepAliveLease(r.leaseID)
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
|
|
||||||
// keepAliveLease maintains the lease alive by sending keep-alive requests
|
|
||||||
func (r *SvcDiscoveryRegistryImpl) keepAliveLease(leaseID clientv3.LeaseID) {
|
|
||||||
ch, err := r.client.KeepAlive(context.Background(), leaseID)
|
|
||||||
if err != nil {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
for ka := range ch {
|
|
||||||
if ka != nil {
|
|
||||||
} else {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// watchServiceChanges watches for changes in the service directory
|
|
||||||
func (r *SvcDiscoveryRegistryImpl) watchServiceChanges() {
|
|
||||||
watchChan := r.client.Watch(context.Background(), r.rootDirectory, clientv3.WithPrefix())
|
|
||||||
for watchResp := range watchChan {
|
|
||||||
updatedPrefixes := make(map[string]struct{}) // Create a set to track updated prefixes
|
|
||||||
|
|
||||||
for _, event := range watchResp.Events {
|
|
||||||
prefix, _ := r.splitEndpoint(string(event.Kv.Key))
|
|
||||||
if _, alreadyUpdated := updatedPrefixes[prefix]; !alreadyUpdated {
|
|
||||||
updatedPrefixes[prefix] = struct{}{} // Mark this prefix as updated
|
|
||||||
r.refreshConnMap(prefix)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// refreshConnMap fetches the latest endpoints and updates the local map
|
|
||||||
func (r *SvcDiscoveryRegistryImpl) refreshConnMap(prefix string) {
|
|
||||||
r.mu.Lock()
|
|
||||||
defer r.mu.Unlock()
|
|
||||||
|
|
||||||
fullPrefix := fmt.Sprintf("%s/", prefix)
|
|
||||||
resp, err := r.client.Get(context.Background(), fullPrefix, clientv3.WithPrefix())
|
|
||||||
if err != nil {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
r.connMap[prefix] = []*grpc.ClientConn{} // Update the connMap with new connections
|
|
||||||
for _, kv := range resp.Kvs {
|
|
||||||
_, addr := r.splitEndpoint(string(kv.Key))
|
|
||||||
conn, err := grpc.DialContext(context.Background(), addr, append(r.dialOptions, grpc.WithResolvers(r.resolver))...)
|
|
||||||
if err != nil {
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
r.connMap[prefix] = append(r.connMap[prefix], conn)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// splitEndpoint splits the endpoint string into prefix and address
|
|
||||||
func (r *SvcDiscoveryRegistryImpl) splitEndpoint(input string) (string, string) {
|
|
||||||
lastSlashIndex := strings.LastIndex(input, "/")
|
|
||||||
if lastSlashIndex != -1 {
|
|
||||||
part1 := input[:lastSlashIndex]
|
|
||||||
part2 := input[lastSlashIndex+1:]
|
|
||||||
return part1, part2
|
|
||||||
}
|
|
||||||
return input, ""
|
|
||||||
}
|
|
||||||
|
|
||||||
// UnRegister removes the service endpoint from etcd
|
|
||||||
func (r *SvcDiscoveryRegistryImpl) UnRegister() error {
|
|
||||||
if r.endpointMgr == nil {
|
|
||||||
return fmt.Errorf("endpoint manager is not initialized")
|
|
||||||
}
|
|
||||||
err := r.endpointMgr.DeleteEndpoint(context.TODO(), r.serviceKey)
|
|
||||||
if err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
|
|
||||||
// Close closes the etcd client connection
|
|
||||||
func (r *SvcDiscoveryRegistryImpl) Close() {
|
|
||||||
if r.client != nil {
|
|
||||||
_ = r.client.Close()
|
|
||||||
}
|
|
||||||
|
|
||||||
r.mu.Lock()
|
|
||||||
defer r.mu.Unlock()
|
|
||||||
}
|
|
||||||
|
|
||||||
// Check verifies if etcd is running by checking the existence of the root node and optionally creates it with a lease
|
|
||||||
func Check(ctx context.Context, etcdServers []string, etcdRoot string, createIfNotExist bool, options ...ZkOption) error {
|
|
||||||
cfg := clientv3.Config{
|
|
||||||
Endpoints: etcdServers,
|
|
||||||
}
|
|
||||||
for _, opt := range options {
|
|
||||||
opt(&cfg)
|
|
||||||
}
|
|
||||||
client, err := clientv3.New(cfg)
|
|
||||||
if err != nil {
|
|
||||||
return errors.Wrap(err, "failed to connect to etcd")
|
|
||||||
}
|
|
||||||
defer client.Close()
|
|
||||||
|
|
||||||
var opCtx context.Context
|
|
||||||
var cancel context.CancelFunc
|
|
||||||
if cfg.DialTimeout != 0 {
|
|
||||||
opCtx, cancel = context.WithTimeout(ctx, cfg.DialTimeout)
|
|
||||||
} else {
|
|
||||||
opCtx, cancel = context.WithTimeout(ctx, 10*time.Second)
|
|
||||||
}
|
|
||||||
defer cancel()
|
|
||||||
|
|
||||||
resp, err := client.Get(opCtx, etcdRoot)
|
|
||||||
if err != nil {
|
|
||||||
return errors.Wrap(err, "failed to get the root node from etcd")
|
|
||||||
}
|
|
||||||
|
|
||||||
if len(resp.Kvs) == 0 {
|
|
||||||
if createIfNotExist {
|
|
||||||
var leaseTTL int64 = 10
|
|
||||||
var leaseResp *clientv3.LeaseGrantResponse
|
|
||||||
if leaseTTL > 0 {
|
|
||||||
leaseResp, err = client.Grant(opCtx, leaseTTL)
|
|
||||||
if err != nil {
|
|
||||||
return errors.Wrap(err, "failed to create lease in etcd")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
putOpts := []clientv3.OpOption{}
|
|
||||||
if leaseResp != nil {
|
|
||||||
putOpts = append(putOpts, clientv3.WithLease(leaseResp.ID))
|
|
||||||
}
|
|
||||||
|
|
||||||
_, err := client.Put(opCtx, etcdRoot, "", putOpts...)
|
|
||||||
if err != nil {
|
|
||||||
return errors.Wrap(err, "failed to create the root node in etcd")
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
return fmt.Errorf("root node %s does not exist in etcd", etcdRoot)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
Loading…
Reference in new issue