package hubrpc import ( "fmt" "sync" "time" "gitlink.org.cn/cloudream/common/consts/errorcode" "gitlink.org.cn/cloudream/jcs-pub/common/pkgs/rpc" grpc "google.golang.org/grpc" "google.golang.org/grpc/credentials/insecure" ) type PoolConfig struct{} type Pool struct { grpcCons map[grpcAddr]*grpcCon lock sync.Mutex } type grpcAddr struct { IP string Port int } type grpcCon struct { grpcCon *grpc.ClientConn refCount int stopClosing chan any } func NewPool(cfg PoolConfig) *Pool { return &Pool{ grpcCons: make(map[grpcAddr]*grpcCon), } } func (p *Pool) Get(ip string, port int) *Client { p.lock.Lock() defer p.lock.Unlock() ga := grpcAddr{IP: ip, Port: port} con := p.grpcCons[ga] if con == nil { gcon, err := grpc.NewClient(fmt.Sprintf("%v:%v", ip, port), grpc.WithTransportCredentials(insecure.NewCredentials())) if err != nil { return &Client{ addr: ga, con: nil, pool: p, fusedErr: rpc.Failed(errorcode.OperationFailed, err.Error()), } } con = &grpcCon{ grpcCon: gcon, refCount: 0, stopClosing: nil, } p.grpcCons[ga] = con } con.refCount++ return &Client{ addr: ga, con: con.grpcCon, cli: NewHubClient(con.grpcCon), pool: p, } } func (p *Pool) release(addr grpcAddr) { p.lock.Lock() defer p.lock.Unlock() grpcCon := p.grpcCons[addr] if grpcCon == nil { return } grpcCon.refCount-- grpcCon.refCount = max(grpcCon.refCount, 0) if grpcCon.refCount == 0 { stopClosing := make(chan any) grpcCon.stopClosing = stopClosing go func() { select { case <-stopClosing: return case <-time.After(time.Minute): p.lock.Lock() defer p.lock.Unlock() if grpcCon.refCount == 0 { grpcCon.grpcCon.Close() delete(p.grpcCons, addr) } } }() } }