JCS-pub/common/pkgs/rpc/hub/pool.go

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)
				}
			}
		}()
	}
}