/
client.go
42 lines (36 loc) · 1.23 KB
/
client.go
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
package rpc
import (
"fmt"
"time"
"github.com/smallnest/rpcx/client"
"github.com/smallnest/rpcx/protocol"
"github.com/smallnest/rpcx/share"
"github.com/bububa/ratelimitd/conf"
"github.com/bububa/ratelimitd/pkg/logger"
"github.com/bububa/ratelimitd/pkg/rpc/codec"
)
func NewProtobufClientPool(cfg conf.ClientConfig, poolSize int) (*client.XClientPool, error) {
share.Codecs[protocol.SerializeType(4)] = &codec.Msgpack{}
share.Codecs[protocol.SerializeType(5)] = &codec.Protobuf{}
return newClientPool(cfg, poolSize, protocol.SerializeType(5))
}
func newClientPool(cfg conf.ClientConfig, poolSize int, serializeType protocol.SerializeType) (*client.XClientPool, error) {
option := client.DefaultOption
option.SerializeType = serializeType
option.Heartbeat = true
option.HeartbeatInterval = time.Second
discover, err := client.NewPeer2PeerDiscovery(fmt.Sprintf("tcp@%s:%d", cfg.Host, cfg.Port), "")
if err != nil {
logger.Error().Err(err).Send()
return nil, err
}
name := cfg.Name
if name == "" {
name = "ratelimitd"
}
return client.NewXClientPool(poolSize, name, client.Failtry, client.RandomSelect, discover, option), nil
}
func GetClientFromPool(pool *client.XClientPool) client.XClient {
clt := pool.Get()
return clt
}