63 lines
1.4 KiB
Go
63 lines
1.4 KiB
Go
package grpc_conn
|
|
|
|
import (
|
|
"fmt"
|
|
"git.hlsq.asia/mmorpg/service-common/log"
|
|
"google.golang.org/grpc"
|
|
"math/rand"
|
|
)
|
|
|
|
type GrpcConnectionMgr struct {
|
|
poolM map[string]*GrpcConnection
|
|
poolS []*GrpcConnection
|
|
}
|
|
|
|
func NewGrpcConnectionMgr() *GrpcConnectionMgr {
|
|
return &GrpcConnectionMgr{
|
|
poolM: make(map[string]*GrpcConnection),
|
|
poolS: make([]*GrpcConnection, 0),
|
|
}
|
|
}
|
|
|
|
func (p *GrpcConnectionMgr) Store(sid string, addr string) {
|
|
pool, err := NewGrpcConnection(sid, addr)
|
|
if err != nil {
|
|
log.Errorf("create grpc err: %v, sid: %v, addr: %v", err, sid, addr)
|
|
return
|
|
}
|
|
p.poolM[sid] = pool
|
|
p.poolS = append(p.poolS, pool)
|
|
}
|
|
|
|
func (p *GrpcConnectionMgr) Delete(sid string) int {
|
|
delete(p.poolM, sid)
|
|
for i, pool := range p.poolS {
|
|
if pool.sid == sid {
|
|
p.poolS = append(p.poolS[:i], p.poolS[i+1:]...)
|
|
break
|
|
}
|
|
}
|
|
return len(p.poolS)
|
|
}
|
|
|
|
func (p *GrpcConnectionMgr) Load(sid ...string) (*grpc.ClientConn, error) {
|
|
var pool *GrpcConnection
|
|
if len(sid) > 0 && sid[0] != "" {
|
|
pool = p.poolM[sid[0]]
|
|
} else {
|
|
pool = p.poolS[rand.Intn(len(p.poolS))]
|
|
}
|
|
if pool == nil {
|
|
return nil, fmt.Errorf("cannot find connection")
|
|
}
|
|
return pool.GetConnection(), nil
|
|
}
|
|
|
|
func (p *GrpcConnectionMgr) LoadAll() map[string]*grpc.ClientConn {
|
|
sidM := make(map[string]*grpc.ClientConn)
|
|
for sid, pool := range p.poolM {
|
|
sidM[sid] = pool.GetConnection()
|
|
}
|
|
return sidM
|
|
}
|