eta_mini_crm_ht/utils/grpc_client_pool.go

package utils

import (
	"google.golang.org/grpc"
	"google.golang.org/grpc/connectivity"
	"sync"
)

type ClientPool interface {
	Get() *grpc.ClientConn
	Put(conn *grpc.ClientConn)
}

type clientPool struct {
	pool sync.Pool
}

func GetPoolInstance(target string, opts ...grpc.DialOption) (ClientPool, error) {
	return &clientPool{
		pool: sync.Pool{
			New: func() any {
				conn, err := grpc.NewClient(target, opts...)
				if err != nil {
					FileLog.Error("连接sso登录服务失败:%v", err)
				}
				return conn
			},
		},
	}, nil
}

func (c *clientPool) Get() *grpc.ClientConn {
	conn := c.pool.Get().(*grpc.ClientConn)

	// 当连接不可用时，关闭当前连接，并新建一个连接
	if conn.GetState() == connectivity.Shutdown || conn.GetState() == connectivity.TransientFailure {
		conn.Close()
		conn = c.pool.New().(*grpc.ClientConn)
	}
	return conn
}

func (c *clientPool) Put(conn *grpc.ClientConn) {
	// 当连接不可用时，关闭当前连接，并不再放回池中
	if conn.GetState() == connectivity.Shutdown || conn.GetState() == connectivity.TransientFailure {
		conn.Close()
		return
	}
	c.pool.Put(conn)
}