102 lines
2.4 KiB
Go
102 lines
2.4 KiB
Go
package discover
|
||
|
||
import (
|
||
"common/db/etcd"
|
||
"common/discover/common"
|
||
"common/discover/grpc_client"
|
||
"common/log"
|
||
"context"
|
||
"fmt"
|
||
clientv3 "go.etcd.io/etcd/client/v3"
|
||
"google.golang.org/grpc"
|
||
"sync"
|
||
)
|
||
|
||
// 大量读少量写的情况下，读写锁比同步Map更高效
|
||
var (
|
||
serverMU = sync.RWMutex{}
|
||
conn = make(map[string]*grpc_client.GrpcConnectionMgr)
|
||
serverLeaseM = make(map[string]clientv3.LeaseID)
|
||
)
|
||
|
||
func init() {
|
||
RegisterListener(common.ListenerTypeNewServer, onServerStart)
|
||
RegisterListener(common.ListenerTypeCloseServer, onServerStop)
|
||
}
|
||
|
||
// FindServer 根据SID或随机查找服务
|
||
func FindServer(target string, sid ...string) (*grpc.ClientConn, error) {
|
||
serverMU.RLock()
|
||
defer serverMU.RUnlock()
|
||
if v, ok := conn[target]; ok {
|
||
return v.Load(sid...)
|
||
}
|
||
return nil, fmt.Errorf("cannot find server")
|
||
}
|
||
|
||
func FindServerAll(target string) map[string]*grpc.ClientConn {
|
||
serverMU.RLock()
|
||
defer serverMU.RUnlock()
|
||
if v, ok := conn[target]; ok {
|
||
return v.LoadAll()
|
||
}
|
||
return make(map[string]*grpc.ClientConn)
|
||
}
|
||
|
||
// RegisterGrpcServer 注册服务提供者
|
||
func RegisterGrpcServer(target, sid, addr string, ttl int64) error {
|
||
serverMU.Lock()
|
||
defer serverMU.Unlock()
|
||
leaseID, err := common.NewLeaseAndKeepAlive(ttl)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
_, err = etcd.Client().Put(context.Background(), target+"/"+sid, addr, clientv3.WithLease(leaseID))
|
||
if err != nil {
|
||
return err
|
||
}
|
||
serverLeaseM[sid] = leaseID
|
||
return nil
|
||
}
|
||
|
||
// UnRegisterGrpcServer 解注册服务提供者
|
||
func UnRegisterGrpcServer(sid string) {
|
||
serverMU.Lock()
|
||
defer serverMU.Unlock()
|
||
if leaseID, ok := serverLeaseM[sid]; ok {
|
||
_, err := etcd.Client().Revoke(context.Background(), leaseID)
|
||
if err != nil {
|
||
log.Errorf("server.go UnRegisterGrpcServer err: %v", err)
|
||
}
|
||
delete(serverLeaseM, sid)
|
||
}
|
||
}
|
||
|
||
// 某个服务启动了
|
||
func onServerStart(data any) {
|
||
if provider, ok := data.(*common.ServiceProvider); ok {
|
||
serverMU.Lock()
|
||
defer serverMU.Unlock()
|
||
if v, ok := conn[provider.Target]; ok {
|
||
v.Store(provider.SID, provider.Addr)
|
||
} else {
|
||
mgr := grpc_client.NewGrpcConnectionMgr()
|
||
mgr.Store(provider.SID, provider.Addr)
|
||
conn[provider.Target] = mgr
|
||
}
|
||
}
|
||
}
|
||
|
||
// 某个服务关闭了
|
||
func onServerStop(data any) {
|
||
if provider, ok := data.(*common.ServiceProvider); ok {
|
||
serverMU.Lock()
|
||
defer serverMU.Unlock()
|
||
if v, ok := conn[provider.Target]; ok {
|
||
if v.Delete(provider.SID) == 0 {
|
||
delete(conn, provider.Target)
|
||
}
|
||
}
|
||
}
|
||
}
|