temp

Server/common/discover/server.go

@@ -0,0 +1,101 @@
+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)
+			}
+		}
+	}
+}