sig-pub/pkg/grpc/discovery/etcd_naming.go

package discovery

import (
	"context"
	"encoding/json"
	"fmt"
	"net"
	"sig-pub/pkg/utils/nets"
	"sig-pub/pkg/zlog"
	"time"

	"github.com/bytedance/sonic"
	"go.etcd.io/etcd/api/v3/mvccpb"
	clientv3 "go.etcd.io/etcd/client/v3"
	"go.etcd.io/etcd/client/v3/naming/endpoints"
	etcdResolver "go.etcd.io/etcd/client/v3/naming/resolver"
	"google.golang.org/grpc/resolver"
)

var (
	DefaultRegisterTTL int64 = 30
)

const (
	EtcdSchema = "etcd"
)

func EtcdDialUrl(serviceName string) string {
	return fmt.Sprintf("%s:///%s", EtcdSchema, serviceName)
}

type EtcdDiscovery struct {
	client *clientv3.Client
}

func NewEtcdDiscovery(client *clientv3.Client) *EtcdDiscovery {
	return &EtcdDiscovery{
		client: client,
	}
}

func (r *EtcdDiscovery) Registry(ctx context.Context, server Server) (err error) {
	em, err := endpoints.NewManager(r.client, server.Name)
	if err != nil {
		return
	}
	ip, port, err := RegisterIpPort(server.Addr)
	if err != nil {
		return
	}

	addr := fmt.Sprintf("%s:%d", ip, port)
	// 序列化 metadata 信息
	meta := "{}"
	if server.Attrs != nil {
		bytes, e := json.Marshal(server.Attrs)
		if e != nil {
			err = e
			return
		}
		meta = string(bytes)
	}

	ctxLease, cancelLease := context.WithTimeout(ctx, time.Second*5)
	defer cancelLease()
	lease, err := r.client.Grant(ctxLease, DefaultRegisterTTL) // 使用租约注册端点确保如果主机无法维持保活心跳, 从服务中删除
	if err != nil {
		return
	}
	endpointKey := fmt.Sprintf("%s/%s", server.Name, addr)
	err = em.AddEndpoint(ctx,
		endpointKey,
		endpoints.Endpoint{
			Addr:     addr,
			Metadata: meta,
		},
		clientv3.WithLease(lease.ID),
	)

	// keepalive lease
	keepCtx, keepCancel := context.WithCancel(context.Background())
	keepAliveCh, err := r.client.KeepAlive(keepCtx, lease.ID)
	if err != nil {
		keepCancel()
		zlog.Error("registry keepalive error: ", err)
		return
	}

	// ticker := time.NewTicker(time.Second * time.Duration(DefaultRegisterTTL))
	w := r.client.Watch(ctx, endpointKey)
	go func() {
		for {
			select {
			case <-ctx.Done():
				zlog.Info("registry keepalive done")
				ctx, c := context.WithTimeout(context.Background(), time.Second*3)
				_, _ = r.client.Revoke(ctx, lease.ID)
				c()
				keepCancel()
				//ticker.Stop()
				return
			case <-keepAliveCh:
				//zlog.Infof("keepalive: lease id %d, %+v", lease.ID, res)
			case res := <-w:
				if err := res.Err(); err != nil {
					zlog.Errorf("registry watch endpoint %s error: %v", endpointKey, err)
					continue
				}
				deleted := false
				for _, event := range res.Events {
					if event.Type == clientv3.EventTypeDelete {
						deleted = true
					}
				}
				// endpoint key被删除, 重放
				if deleted {
					zlog.Infof("registry endpoint %s deleted ", endpointKey)
					// 删除旧的 lease
					keepCancel()
					_, _ = r.client.Revoke(context.Background(), lease.ID)
					lease, err = r.client.Grant(ctx, DefaultRegisterTTL)
					if err != nil {
						zlog.Error("registry grant lease again error: ", err)
					}
					err = em.AddEndpoint(ctx,
						endpointKey,
						endpoints.Endpoint{
							Addr:     addr,
							Metadata: meta,
						},
						clientv3.WithLease(lease.ID),
					)
					if err != nil {
						zlog.Error("refresh endpoint error: ", err)
						continue
					}
					// keepalive lease
					keepCtx, keepCancel = context.WithCancel(context.Background())
					keepAliveCh, err = r.client.KeepAlive(keepCtx, lease.ID)
					if err != nil {
						zlog.Error("registry keepalive error: ", err)
					}
				}
				//case <-ticker.C: // 定时重放防止etcd中key被删除
			}
		}
	}()
	return
}

func (r *EtcdDiscovery) DialUrl(serviceName string) string {
	return fmt.Sprintf("%s:///%s", EtcdSchema, serviceName)
}

func (r *EtcdDiscovery) Resolver() (builder resolver.Builder, err error) {
	builder, err = etcdResolver.NewBuilder(r.client)
	return
}

func (r *EtcdDiscovery) ResolveAll(ctx context.Context, serviceName string) (servers []Server, err error) {
	res, err := r.client.Get(ctx, serviceName+"/", clientv3.WithPrefix())
	if err != nil {
		return
	}

	for _, kv := range res.Kvs {
		endpoint := endpoints.Endpoint{}
		if err = sonic.Unmarshal(kv.Value, &endpoint); err != nil {
			zlog.Errorf("resolve service %s error: ", string(kv.Value))
			return
		}
		server := Server{
			Name: serviceName,
			Addr: endpoint.Addr,
		}
		if endpoint.Metadata != nil {
			if strMeta, ok := endpoint.Metadata.(string); ok {
				err = json.Unmarshal([]byte(strMeta), &server.Attrs)
				if err != nil {
					return
				}
			}
		}

		servers = append(servers, server)
	}
	return
}

func (r *EtcdDiscovery) Watch(ctx context.Context, serviceName string) (ch chan []Server, err error) {
	w := r.client.Watch(ctx, serviceName+"/", clientv3.WithPrefix())
	ch = make(chan []Server, 1)
	go func() {
		for {
			select {
			case <-ctx.Done():
				close(ch)
				return
			case res := <-w:
				if err := res.Err(); err != nil {
					zlog.Errorf("watch service %s error: %v", serviceName, err)
					continue
				}

				for _, event := range res.Events {
					switch event.Type {
					case mvccpb.DELETE:
						fallthrough
					case mvccpb.PUT:
						servers, err := r.ResolveAll(context.Background(), serviceName)
						if err != nil {
							zlog.Errorf("watch event %v for service %s error: %v", event.Type, serviceName, err)
							continue
						}
						ch <- servers
					}
				}
			}
		}
	}()
	return
}

func RegisterIpPort(addr string) (ip string, port int, err error) {
	tcpAddr, err := net.ResolveTCPAddr("tcp", addr)
	if err != nil {
		return
	}
	port = tcpAddr.Port
	if tcpAddr.IP != nil {
		ip = tcpAddr.IP.String()
	} else {
		ip, err = nets.GetHostIpv4()
		if err != nil {
			return
		}
	}
	return
}