userProxyManager.go 1.27 KB
Newer Older
hujiebin's avatar
hujiebin committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58
package manager

import (
	"context"
	"google.golang.org/grpc"
	"log"
	"sync"
	"time"
)

var UserProxyMgr = UserProxyManager{}

type UserProxyManager struct {
	mutex   sync.RWMutex
	clients map[string]*grpc.ClientConn
}

func init() {
	UserProxyMgr.clients = make(map[string]*grpc.ClientConn, 0)
}

// 创建Client连接,保存并返回
func (m *UserProxyManager) MakeClient(addr string) *grpc.ClientConn {
	conn := m.GetClient(addr)
	if conn == nil {
		ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
		defer cancel()

		c, err := grpc.DialContext(ctx, addr, grpc.WithInsecure(), grpc.WithBlock())
		if err == nil {
			if c == nil {
				log.Printf("Failed in connecting to %s:\n", addr)
			} else {
				conn = c
				m.SetClient(addr, conn)

				log.Printf("Connection to %s established.\n", addr)
			}
		} else {
			log.Printf("Failed to connect %s: %v\n", addr, err)
		}
	} else {
		//log.Printf("Connection to %s already exists, state = %s\n", addr, conn.GetState().String())
	}
	return conn
}

func (m *UserProxyManager) GetClient(addr string) *grpc.ClientConn {
	m.mutex.RLock()
	defer m.mutex.RUnlock()
	return m.clients[addr]
}

func (m *UserProxyManager) SetClient(addr string, c *grpc.ClientConn) {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	m.clients[addr] = c
}