Add pluggable codec support

This commit is contained in:
Asim 2015-11-25 19:50:05 +00:00
parent 04e07f4b39
commit 12a1e1eeda
8 changed files with 248 additions and 64 deletions

75
client/codec.go Normal file
View File

@ -0,0 +1,75 @@
package client
import (
"io"
"net/rpc"
"github.com/youtube/vitess/go/rpcplus"
"github.com/youtube/vitess/go/rpcplus/jsonrpc"
"github.com/youtube/vitess/go/rpcplus/pbrpc"
)
var (
defaultContentType = "application/octet-stream"
defaultCodecs = map[string]codecFunc{
"application/json": jsonrpc.NewClientCodec,
"application/json-rpc": jsonrpc.NewClientCodec,
"application/protobuf": pbrpc.NewClientCodec,
"application/proto-rpc": pbrpc.NewClientCodec,
"application/octet-stream": pbrpc.NewClientCodec,
}
)
type CodecFunc func(io.ReadWriteCloser) rpc.ClientCodec
// only for internal use
type codecFunc func(io.ReadWriteCloser) rpcplus.ClientCodec
// wraps an net/rpc ClientCodec to provide an rpcplus.ClientCodec
// temporary until we strip out use of rpcplus
type rpcCodecWrap struct {
r rpc.ClientCodec
}
func (cw *rpcCodecWrap) WriteRequest(r *rpcplus.Request, b interface{}) error {
rc := &rpc.Request{
ServiceMethod: r.ServiceMethod,
Seq: r.Seq,
}
err := cw.r.WriteRequest(rc, b)
r.ServiceMethod = rc.ServiceMethod
r.Seq = rc.Seq
return err
}
func (cw *rpcCodecWrap) ReadResponseHeader(r *rpcplus.Response) error {
rc := &rpc.Response{
ServiceMethod: r.ServiceMethod,
Seq: r.Seq,
Error: r.Error,
}
err := cw.r.ReadResponseHeader(rc)
r.ServiceMethod = rc.ServiceMethod
r.Seq = rc.Seq
r.Error = r.Error
return err
}
func (cw *rpcCodecWrap) ReadResponseBody(b interface{}) error {
return cw.r.ReadResponseBody(b)
}
func (cw *rpcCodecWrap) Close() error {
return cw.r.Close()
}
// wraps a CodecFunc to provide an internal codecFunc
// temporary until we strip rpcplus out
func codecWrap(cf CodecFunc) codecFunc {
return func(rwc io.ReadWriteCloser) rpcplus.ClientCodec {
return &rpcCodecWrap{
r: cf(rwc),
}
}
}

View File

@ -7,9 +7,11 @@ import (
)
type options struct {
broker broker.Broker
registry registry.Registry
transport transport.Transport
contentType string
codecs map[string]CodecFunc
broker broker.Broker
registry registry.Registry
transport transport.Transport
}
func Broker(b broker.Broker) Option {
@ -18,6 +20,18 @@ func Broker(b broker.Broker) Option {
}
}
func Codec(contentType string, cf CodecFunc) Option {
return func(o *options) {
o.codecs[contentType] = cf
}
}
func ContentType(ct string) Option {
return func(o *options) {
o.contentType = ct
}
}
func Registry(r registry.Registry) Option {
return func(o *options) {
o.registry = r

View File

@ -28,12 +28,18 @@ type rpcClient struct {
}
func newRpcClient(opt ...Option) Client {
var opts options
opts := options{
codecs: make(map[string]CodecFunc),
}
for _, o := range opt {
o(&opts)
}
if len(opts.contentType) == 0 {
opts.contentType = defaultContentType
}
if opts.transport == nil {
opts.transport = transport.DefaultTransport
}
@ -55,6 +61,16 @@ func (t *headerRoundTripper) RoundTrip(r *http.Request) (*http.Response, error)
return t.r.RoundTrip(r)
}
func (r *rpcClient) codecFunc(contentType string) (codecFunc, error) {
if cf, ok := r.opts.codecs[contentType]; ok {
return codecWrap(cf), nil
}
if cf, ok := defaultCodecs[contentType]; ok {
return cf, nil
}
return nil, fmt.Errorf("Unsupported Content-Type: %s", contentType)
}
func (r *rpcClient) call(ctx context.Context, address string, request Request, response interface{}) error {
msg := &transport.Message{
Header: make(map[string]string),
@ -69,13 +85,18 @@ func (r *rpcClient) call(ctx context.Context, address string, request Request, r
msg.Header["Content-Type"] = request.ContentType()
cf, err := r.codecFunc(request.ContentType())
if err != nil {
return errors.InternalServerError("go.micro.client", err.Error())
}
c, err := r.opts.transport.Dial(address)
if err != nil {
return errors.InternalServerError("go.micro.client", fmt.Sprintf("Error sending request: %v", err))
}
defer c.Close()
client := rpc.NewClientWithCodec(newRpcPlusCodec(msg, c))
client := rpc.NewClientWithCodec(newRpcPlusCodec(msg, c, cf))
err = client.Call(ctx, request.Method(), request.Request(), response)
if err != nil {
return err
@ -97,12 +118,17 @@ func (r *rpcClient) stream(ctx context.Context, address string, request Request,
msg.Header["Content-Type"] = request.ContentType()
cf, err := r.codecFunc(request.ContentType())
if err != nil {
return nil, errors.InternalServerError("go.micro.client", err.Error())
}
c, err := r.opts.transport.Dial(address, transport.WithStream())
if err != nil {
return nil, errors.InternalServerError("go.micro.client", fmt.Sprintf("Error sending request: %v", err))
}
client := rpc.NewClientWithCodec(newRpcPlusCodec(msg, c))
client := rpc.NewClientWithCodec(newRpcPlusCodec(msg, c, cf))
call := client.StreamGo(request.Method(), request.Request(), responseChan)
return &rpcStream{
@ -195,14 +221,14 @@ func (r *rpcClient) Publish(ctx context.Context, p Publication) error {
}
func (r *rpcClient) NewPublication(topic string, message interface{}) Publication {
return r.NewProtoPublication(topic, message)
return newRpcPublication(topic, message, r.opts.contentType)
}
func (r *rpcClient) NewProtoPublication(topic string, message interface{}) Publication {
return newRpcPublication(topic, message, "application/octet-stream")
}
func (r *rpcClient) NewRequest(service, method string, request interface{}) Request {
return r.NewProtoRequest(service, method, request)
return newRpcRequest(service, method, request, r.opts.contentType)
}
func (r *rpcClient) NewProtoRequest(service, method string, request interface{}) Request {

View File

@ -2,12 +2,9 @@ package client
import (
"bytes"
"fmt"
"github.com/micro/go-micro/transport"
rpc "github.com/youtube/vitess/go/rpcplus"
js "github.com/youtube/vitess/go/rpcplus/jsonrpc"
pb "github.com/youtube/vitess/go/rpcplus/pbrpc"
)
type rpcPlusCodec struct {
@ -37,50 +34,33 @@ func (rwc *readWriteCloser) Close() error {
return nil
}
func newRpcPlusCodec(req *transport.Message, client transport.Client) *rpcPlusCodec {
func newRpcPlusCodec(req *transport.Message, client transport.Client, cf codecFunc) *rpcPlusCodec {
rwc := &readWriteCloser{
wbuf: bytes.NewBuffer(nil),
rbuf: bytes.NewBuffer(nil),
}
r := &rpcPlusCodec{
req: req,
buf: rwc,
client: client,
buf: &readWriteCloser{
wbuf: bytes.NewBuffer(nil),
rbuf: bytes.NewBuffer(nil),
},
codec: cf(rwc),
req: req,
}
switch req.Header["Content-Type"] {
case "application/octet-stream":
r.codec = pb.NewClientCodec(r.buf)
case "application/json":
r.codec = js.NewClientCodec(r.buf)
}
return r
}
func (c *rpcPlusCodec) WriteRequest(req *rpc.Request, body interface{}) error {
if c.codec == nil {
return fmt.Errorf("unsupported request type: %s", c.req.Header["Content-Type"])
}
if err := c.codec.WriteRequest(req, body); err != nil {
return err
}
c.req.Body = c.buf.wbuf.Bytes()
return c.client.Send(c.req)
}
func (c *rpcPlusCodec) ReadResponseHeader(r *rpc.Response) error {
var m transport.Message
if err := c.client.Recv(&m); err != nil {
return err
}
if c.codec == nil {
return fmt.Errorf("%s", string(m.Body))
}
c.buf.rbuf.Reset()
c.buf.rbuf.Write(m.Body)
return c.codec.ReadResponseHeader(r)

73
server/codec.go Normal file
View File

@ -0,0 +1,73 @@
package server
import (
"io"
"net/rpc"
"github.com/youtube/vitess/go/rpcplus"
"github.com/youtube/vitess/go/rpcplus/jsonrpc"
"github.com/youtube/vitess/go/rpcplus/pbrpc"
)
type CodecFunc func(io.ReadWriteCloser) rpc.ServerCodec
// for internal use only
type codecFunc func(io.ReadWriteCloser) rpcplus.ServerCodec
var (
defaultCodecs = map[string]codecFunc{
"application/json": jsonrpc.NewServerCodec,
"application/json-rpc": jsonrpc.NewServerCodec,
"application/protobuf": pbrpc.NewServerCodec,
"application/proto-rpc": pbrpc.NewServerCodec,
"application/octet-stream": pbrpc.NewServerCodec,
}
)
// wraps an net/rpc ServerCodec to provide an rpcplus.ServerCodec
// temporary until we strip out use of rpcplus
type rpcCodecWrap struct {
r rpc.ServerCodec
}
func (cw *rpcCodecWrap) ReadRequestHeader(r *rpcplus.Request) error {
rc := &rpc.Request{
ServiceMethod: r.ServiceMethod,
Seq: r.Seq,
}
err := cw.r.ReadRequestHeader(rc)
r.ServiceMethod = rc.ServiceMethod
r.Seq = rc.Seq
return err
}
func (cw *rpcCodecWrap) ReadRequestBody(b interface{}) error {
return cw.r.ReadRequestBody(b)
}
func (cw *rpcCodecWrap) WriteResponse(r *rpcplus.Response, b interface{}, l bool) error {
rc := &rpc.Response{
ServiceMethod: r.ServiceMethod,
Seq: r.Seq,
Error: r.Error,
}
err := cw.r.WriteResponse(rc, b)
r.ServiceMethod = rc.ServiceMethod
r.Seq = rc.Seq
r.Error = r.Error
return err
}
func (cw *rpcCodecWrap) Close() error {
return cw.r.Close()
}
// wraps a CodecFunc to provide an internal codecFunc
// temporary until we strip rpcplus out
func codecWrap(cf CodecFunc) codecFunc {
return func(rwc io.ReadWriteCloser) rpcplus.ServerCodec {
return &rpcCodecWrap{
r: cf(rwc),
}
}
}

View File

@ -7,6 +7,7 @@ import (
)
type options struct {
codecs map[string]CodecFunc
broker broker.Broker
registry registry.Registry
transport transport.Transport
@ -19,7 +20,9 @@ type options struct {
}
func newOptions(opt ...Option) options {
var opts options
opts := options{
codecs: make(map[string]CodecFunc),
}
for _, o := range opt {
o(&opts)
@ -116,6 +119,12 @@ func Broker(b broker.Broker) Option {
}
}
func Codec(contentType string, cf CodecFunc) Option {
return func(o *options) {
o.codecs[contentType] = cf
}
}
func Registry(r registry.Registry) Option {
return func(o *options) {
o.registry = r

View File

@ -2,11 +2,9 @@ package server
import (
"bytes"
"fmt"
"github.com/micro/go-micro/transport"
rpc "github.com/youtube/vitess/go/rpcplus"
js "github.com/youtube/vitess/go/rpcplus/jsonrpc"
pb "github.com/youtube/vitess/go/rpcplus/pbrpc"
)
type rpcPlusCodec struct {
@ -36,30 +34,21 @@ func (rwc *readWriteCloser) Close() error {
return nil
}
func newRpcPlusCodec(req *transport.Message, socket transport.Socket) rpc.ServerCodec {
func newRpcPlusCodec(req *transport.Message, socket transport.Socket, cf codecFunc) rpc.ServerCodec {
rwc := &readWriteCloser{
rbuf: bytes.NewBuffer(req.Body),
wbuf: bytes.NewBuffer(nil),
}
r := &rpcPlusCodec{
socket: socket,
buf: rwc,
codec: cf(rwc),
req: req,
buf: &readWriteCloser{
rbuf: bytes.NewBuffer(req.Body),
wbuf: bytes.NewBuffer(nil),
},
socket: socket,
}
switch req.Header["Content-Type"] {
case "application/octet-stream":
r.codec = pb.NewServerCodec(r.buf)
case "application/json":
r.codec = js.NewServerCodec(r.buf)
}
return r
}
func (c *rpcPlusCodec) ReadRequestHeader(r *rpc.Request) error {
if c.codec == nil {
return fmt.Errorf("unsupported content type %s", c.req.Header["Content-Type"])
}
return c.codec.ReadRequestHeader(r)
}
@ -68,19 +57,14 @@ func (c *rpcPlusCodec) ReadRequestBody(r interface{}) error {
}
func (c *rpcPlusCodec) WriteResponse(r *rpc.Response, body interface{}, last bool) error {
if c.codec == nil {
return fmt.Errorf("unsupported request type: %s", c.req.Header["Content-Type"])
}
c.buf.wbuf.Reset()
if err := c.codec.WriteResponse(r, body, last); err != nil {
return err
}
return c.socket.Send(&transport.Message{
Header: map[string]string{"Content-Type": c.req.Header["Content-Type"]},
Body: c.buf.wbuf.Bytes(),
})
}
func (c *rpcPlusCodec) Close() error {

View File

@ -43,7 +43,20 @@ func (s *rpcServer) accept(sock transport.Socket) {
return
}
codec := newRpcPlusCodec(&msg, sock)
cf, err := s.codecFunc(msg.Header["Content-Type"])
// TODO: needs better error handling
if err != nil {
sock.Send(&transport.Message{
Header: map[string]string{
"Content-Type": "text/plain",
},
Body: []byte(err.Error()),
})
sock.Close()
return
}
codec := newRpcPlusCodec(&msg, sock, cf)
// strip our headers
hdr := make(map[string]string)
@ -55,11 +68,21 @@ func (s *rpcServer) accept(sock transport.Socket) {
ctx := c.WithMetadata(context.Background(), hdr)
// TODO: needs better error handling
if err := s.rpc.ServeRequestWithContext(ctx, codec); err != nil {
log.Errorf("Unexpected error servinc request, closing socket: %v", err)
log.Errorf("Unexpected error serving request, closing socket: %v", err)
sock.Close()
}
}
func (s *rpcServer) codecFunc(contentType string) (codecFunc, error) {
if cf, ok := s.opts.codecs[contentType]; ok {
return codecWrap(cf), nil
}
if cf, ok := defaultCodecs[contentType]; ok {
return cf, nil
}
return nil, fmt.Errorf("Unsupported Content-Type: %s", contentType)
}
func (s *rpcServer) Config() options {
s.RLock()
opts := s.opts