micro/client/rpc_client.go

225 lines
5.3 KiB
Go
Raw Normal View History

2015-01-14 02:31:27 +03:00
package client
import (
"encoding/json"
2015-01-14 02:31:27 +03:00
"fmt"
"math/rand"
"net/http"
"sync"
2015-01-14 02:31:27 +03:00
"time"
"github.com/myodc/go-micro/broker"
2015-05-23 13:53:40 +03:00
c "github.com/myodc/go-micro/context"
2015-05-05 21:05:06 +03:00
"github.com/myodc/go-micro/errors"
"github.com/myodc/go-micro/registry"
2015-05-21 00:57:19 +03:00
"github.com/myodc/go-micro/transport"
2015-05-23 13:53:40 +03:00
2015-01-14 02:31:27 +03:00
rpc "github.com/youtube/vitess/go/rpcplus"
2015-05-23 13:53:40 +03:00
"github.com/golang/protobuf/proto"
2015-05-23 13:53:40 +03:00
"golang.org/x/net/context"
2015-01-14 02:31:27 +03:00
)
type headerRoundTripper struct {
r http.RoundTripper
}
2015-05-23 19:40:53 +03:00
type rpcClient struct {
once sync.Once
2015-05-21 21:24:57 +03:00
opts options
2015-05-21 00:57:19 +03:00
}
2015-01-14 02:31:27 +03:00
func init() {
rand.Seed(time.Now().UnixNano())
}
2015-05-23 19:40:53 +03:00
func newRpcClient(opt ...Option) Client {
var opts options
for _, o := range opt {
o(&opts)
}
if opts.transport == nil {
opts.transport = transport.DefaultTransport
}
if opts.broker == nil {
opts.broker = broker.DefaultBroker
}
var once sync.Once
2015-05-23 19:40:53 +03:00
return &rpcClient{
once: once,
2015-05-23 19:40:53 +03:00
opts: opts,
}
}
2015-01-14 02:31:27 +03:00
func (t *headerRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
r.Header.Set("X-Client-Version", "1.0")
return t.r.RoundTrip(r)
}
2015-05-23 19:40:53 +03:00
func (r *rpcClient) call(ctx context.Context, address string, request Request, response interface{}) error {
2015-05-21 00:57:19 +03:00
msg := &transport.Message{
Header: make(map[string]string),
}
2015-01-14 02:31:27 +03:00
2015-05-27 00:39:48 +03:00
md, ok := c.GetMetadata(ctx)
2015-05-23 13:53:40 +03:00
if ok {
for k, v := range md {
msg.Header[k] = v
2015-05-21 00:57:19 +03:00
}
2015-01-14 02:31:27 +03:00
}
2015-05-21 00:57:19 +03:00
msg.Header["Content-Type"] = request.ContentType()
2015-05-21 23:08:19 +03:00
c, err := r.opts.transport.Dial(address)
2015-01-14 02:31:27 +03:00
if err != nil {
return errors.InternalServerError("go.micro.client", fmt.Sprintf("Error sending request: %v", err))
}
2015-10-22 12:33:33 +03:00
defer c.Close()
2015-01-14 02:31:27 +03:00
client := rpc.NewClientWithCodec(newRpcPlusCodec(msg, c))
2015-10-22 17:14:56 +03:00
err = client.Call(ctx, request.Method(), request.Request(), response)
if err != nil {
return err
}
return client.Close()
}
2015-01-14 02:31:27 +03:00
func (r *rpcClient) stream(ctx context.Context, address string, request Request, responseChan interface{}) (Streamer, error) {
msg := &transport.Message{
Header: make(map[string]string),
2015-01-14 02:31:27 +03:00
}
md, ok := c.GetMetadata(ctx)
if ok {
for k, v := range md {
msg.Header[k] = v
}
2015-01-14 02:31:27 +03:00
}
msg.Header["Content-Type"] = request.ContentType()
2015-01-14 02:31:27 +03:00
c, err := r.opts.transport.Dial(address, transport.WithStream())
2015-01-14 02:31:27 +03:00
if err != nil {
return nil, errors.InternalServerError("go.micro.client", fmt.Sprintf("Error sending request: %v", err))
2015-01-14 02:31:27 +03:00
}
client := rpc.NewClientWithCodec(newRpcPlusCodec(msg, c))
call := client.StreamGo(request.Method(), request.Request(), responseChan)
return &rpcStream{
request: request,
call: call,
client: client,
}, nil
2015-01-14 02:31:27 +03:00
}
2015-05-23 19:40:53 +03:00
func (r *rpcClient) CallRemote(ctx context.Context, address string, request Request, response interface{}) error {
2015-05-23 13:53:40 +03:00
return r.call(ctx, address, request, response)
2015-01-14 02:31:27 +03:00
}
// TODO: Call(..., opts *Options) error {
2015-05-23 19:40:53 +03:00
func (r *rpcClient) Call(ctx context.Context, request Request, response interface{}) error {
2015-01-14 02:31:27 +03:00
service, err := registry.GetService(request.Service())
if err != nil {
return errors.InternalServerError("go.micro.client", err.Error())
}
if len(service.Nodes) == 0 {
2015-01-14 02:31:27 +03:00
return errors.NotFound("go.micro.client", "Service not found")
}
n := rand.Int() % len(service.Nodes)
node := service.Nodes[n]
2015-05-21 21:24:57 +03:00
address := node.Address
if node.Port > 0 {
address = fmt.Sprintf("%s:%d", address, node.Port)
2015-05-21 21:24:57 +03:00
}
2015-05-23 13:53:40 +03:00
return r.call(ctx, address, request, response)
2015-01-14 02:31:27 +03:00
}
func (r *rpcClient) StreamRemote(ctx context.Context, address string, request Request, responseChan interface{}) (Streamer, error) {
return r.stream(ctx, address, request, responseChan)
}
func (r *rpcClient) Stream(ctx context.Context, request Request, responseChan interface{}) (Streamer, error) {
service, err := registry.GetService(request.Service())
if err != nil {
return nil, errors.InternalServerError("go.micro.client", err.Error())
}
if len(service.Nodes) == 0 {
return nil, errors.NotFound("go.micro.client", "Service not found")
}
n := rand.Int() % len(service.Nodes)
node := service.Nodes[n]
address := node.Address
if node.Port > 0 {
address = fmt.Sprintf("%s:%d", address, node.Port)
}
return r.stream(ctx, address, request, responseChan)
}
func (r *rpcClient) Publish(ctx context.Context, p Publication) error {
md, ok := c.GetMetadata(ctx)
if !ok {
md = make(map[string]string)
}
md["Content-Type"] = p.ContentType()
// encode message body
var body []byte
switch p.ContentType() {
case "application/octet-stream":
b, err := proto.Marshal(p.Message().(proto.Message))
if err != nil {
return err
}
body = b
case "application/json":
b, err := json.Marshal(p.Message())
if err != nil {
return err
}
body = b
}
r.once.Do(func() {
r.opts.broker.Connect()
})
return r.opts.broker.Publish(p.Topic(), &broker.Message{
Header: md,
Body: body,
})
}
func (r *rpcClient) NewPublication(topic string, message interface{}) Publication {
return r.NewProtoPublication(topic, message)
}
func (r *rpcClient) NewProtoPublication(topic string, message interface{}) Publication {
return newRpcPublication(topic, message, "application/octet-stream")
}
2015-05-23 19:40:53 +03:00
func (r *rpcClient) NewRequest(service, method string, request interface{}) Request {
2015-01-14 02:31:27 +03:00
return r.NewProtoRequest(service, method, request)
}
2015-05-23 19:40:53 +03:00
func (r *rpcClient) NewProtoRequest(service, method string, request interface{}) Request {
2015-01-14 02:31:27 +03:00
return newRpcRequest(service, method, request, "application/octet-stream")
}
2015-05-23 19:40:53 +03:00
func (r *rpcClient) NewJsonRequest(service, method string, request interface{}) Request {
2015-01-14 02:31:27 +03:00
return newRpcRequest(service, method, request, "application/json")
}