Updated codec interface and code. Painful stuff
This commit is contained in:
Asim
@@ -1,23 +0,0 @@
|
||||
package bson
|
||||
|
||||
import (
|
||||
"labix.org/v2/mgo/bson"
|
||||
)
|
||||
|
||||
var (
|
||||
Codec = bsonCodec{}
|
||||
)
|
||||
|
||||
type bsonCodec struct {}
|
||||
|
||||
func (bsonCodec) Marshal(v interface{}) ([]byte, error) {
|
||||
return bson.Marshal(v)
|
||||
}
|
||||
|
||||
func (bsonCodec) Unmarshal(data []byte, v interface{}) error {
|
||||
return bson.Unmarshal(data, v)
|
||||
}
|
||||
|
||||
func (bsonCodec) String() string {
|
||||
return "bson"
|
||||
}
|
||||
@@ -1,12 +1,47 @@
|
||||
package codec
|
||||
|
||||
// Codec used to encode and decode request/responses
|
||||
import (
|
||||
"io"
|
||||
)
|
||||
|
||||
const (
|
||||
Error MessageType = iota
|
||||
Request
|
||||
Response
|
||||
Publication
|
||||
)
|
||||
|
||||
type MessageType int
|
||||
|
||||
// Takes in a connection/buffer and returns a new Codec
|
||||
type NewCodec func(io.ReadWriteCloser) Codec
|
||||
|
||||
// Codec encodes/decodes various types of
|
||||
// messages used within go-micro
|
||||
type Codec interface {
|
||||
// Marshal returns the wire format of v.
|
||||
Marshal(v interface{}) ([]byte, error)
|
||||
// Unmarshal parses the wire format into v.
|
||||
Unmarshal(data []byte, v interface{}) error
|
||||
// String returns the name of the Codec implementation. The returned
|
||||
// string will be used as part of content type in transmission.
|
||||
Encoder
|
||||
Decoder
|
||||
Close() error
|
||||
String() string
|
||||
}
|
||||
|
||||
type Encoder interface {
|
||||
Write(*Message, interface{}) error
|
||||
}
|
||||
|
||||
type Decoder interface {
|
||||
ReadHeader(*Message, MessageType) error
|
||||
ReadBody(interface{}) error
|
||||
}
|
||||
|
||||
// Message represents detailed information about
|
||||
// the communication, likely followed by the body.
|
||||
// In the case of an error, body may be nil.
|
||||
type Message struct {
|
||||
Id uint64
|
||||
Type MessageType
|
||||
Target string
|
||||
Method string
|
||||
Error string
|
||||
Headers map[string]string
|
||||
}
|
||||
|
||||
@@ -1,24 +0,0 @@
|
||||
package pb
|
||||
|
||||
import (
|
||||
"github.com/golang/protobuf/proto"
|
||||
)
|
||||
|
||||
var (
|
||||
Codec = protoCodec{}
|
||||
)
|
||||
|
||||
type protoCodec struct {}
|
||||
|
||||
func (protoCodec) Marshal(v interface{}) ([]byte, error) {
|
||||
return proto.Marshal(v.(proto.Message))
|
||||
}
|
||||
|
||||
func (protoCodec) Unmarshal(data []byte, v interface{}) error {
|
||||
return proto.Unmarshal(data, v.(proto.Message))
|
||||
}
|
||||
|
||||
func (protoCodec) String() string {
|
||||
return "proto"
|
||||
}
|
||||
|
||||
36
codec/proto/netstring.go
Normal file
36
codec/proto/netstring.go
Normal file
@@ -0,0 +1,36 @@
|
||||
package proto
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"io"
|
||||
)
|
||||
|
||||
// WriteNetString writes data to a big-endian netstring on a Writer.
|
||||
// Size is always a 32-bit unsigned int.
|
||||
func WriteNetString(w io.Writer, data []byte) (written int, err error) {
|
||||
size := make([]byte, 4)
|
||||
binary.BigEndian.PutUint32(size, uint32(len(data)))
|
||||
if written, err = w.Write(size); err != nil {
|
||||
return
|
||||
}
|
||||
return w.Write(data)
|
||||
}
|
||||
|
||||
// ReadNetString reads data from a big-endian netstring.
|
||||
func ReadNetString(r io.Reader) (data []byte, err error) {
|
||||
sizeBuf := make([]byte, 4)
|
||||
_, err = r.Read(sizeBuf)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
size := binary.BigEndian.Uint32(sizeBuf)
|
||||
if size == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
data = make([]byte, size)
|
||||
_, err = r.Read(data)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return
|
||||
}
|
||||
163
codec/proto/proto.go
Normal file
163
codec/proto/proto.go
Normal file
@@ -0,0 +1,163 @@
|
||||
package proto
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io"
|
||||
"sync"
|
||||
|
||||
"github.com/golang/protobuf/proto"
|
||||
"github.com/micro/go-micro/codec"
|
||||
rpc "github.com/youtube/vitess/go/rpcplus/pbrpc"
|
||||
)
|
||||
|
||||
type flusher interface {
|
||||
Flush() error
|
||||
}
|
||||
|
||||
type protoCodec struct {
|
||||
sync.Mutex
|
||||
rwc io.ReadWriteCloser
|
||||
mt codec.MessageType
|
||||
buf *bytes.Buffer
|
||||
}
|
||||
|
||||
func (c *protoCodec) Close() error {
|
||||
c.buf.Reset()
|
||||
return c.rwc.Close()
|
||||
}
|
||||
|
||||
func (c *protoCodec) String() string {
|
||||
return "proto"
|
||||
}
|
||||
|
||||
func (c *protoCodec) Write(m *codec.Message, b interface{}) error {
|
||||
switch m.Type {
|
||||
case codec.Request:
|
||||
c.Lock()
|
||||
defer c.Unlock()
|
||||
// This is protobuf, of course we copy it.
|
||||
pbr := &rpc.Request{ServiceMethod: &m.Method, Seq: &m.Id}
|
||||
data, err := proto.Marshal(pbr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = WriteNetString(c.rwc, data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Of course this is a protobuf! Trust me or detonate the program.
|
||||
data, err = proto.Marshal(b.(proto.Message))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = WriteNetString(c.rwc, data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if flusher, ok := c.rwc.(flusher); ok {
|
||||
err = flusher.Flush()
|
||||
}
|
||||
case codec.Response:
|
||||
c.Lock()
|
||||
defer c.Unlock()
|
||||
rtmp := &rpc.Response{ServiceMethod: &m.Method, Seq: &m.Id, Error: &m.Error}
|
||||
data, err := proto.Marshal(rtmp)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = WriteNetString(c.rwc, data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if pb, ok := b.(proto.Message); ok {
|
||||
data, err = proto.Marshal(pb)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
} else {
|
||||
data = nil
|
||||
}
|
||||
_, err = WriteNetString(c.rwc, data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if flusher, ok := c.rwc.(flusher); ok {
|
||||
err = flusher.Flush()
|
||||
}
|
||||
case codec.Publication:
|
||||
data, err := proto.Marshal(b.(proto.Message))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
c.rwc.Write(data)
|
||||
default:
|
||||
return fmt.Errorf("Unrecognised message type: %v", m.Type)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *protoCodec) ReadHeader(m *codec.Message, mt codec.MessageType) error {
|
||||
c.buf.Reset()
|
||||
c.mt = mt
|
||||
|
||||
switch mt {
|
||||
case codec.Request:
|
||||
data, err := ReadNetString(c.rwc)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
rtmp := new(rpc.Request)
|
||||
err = proto.Unmarshal(data, rtmp)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
m.Method = *rtmp.ServiceMethod
|
||||
m.Id = *rtmp.Seq
|
||||
case codec.Response:
|
||||
data, err := ReadNetString(c.rwc)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
rtmp := new(rpc.Response)
|
||||
err = proto.Unmarshal(data, rtmp)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
m.Method = *rtmp.ServiceMethod
|
||||
m.Id = *rtmp.Seq
|
||||
m.Error = *rtmp.Error
|
||||
case codec.Publication:
|
||||
io.Copy(c.buf, c.rwc)
|
||||
default:
|
||||
return fmt.Errorf("Unrecognised message type: %v", mt)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *protoCodec) ReadBody(b interface{}) error {
|
||||
var data []byte
|
||||
switch c.mt {
|
||||
case codec.Request, codec.Response:
|
||||
var err error
|
||||
data, err = ReadNetString(c.rwc)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
case codec.Publication:
|
||||
data = c.buf.Bytes()
|
||||
default:
|
||||
return fmt.Errorf("Unrecognised message type: %v", c.mt)
|
||||
}
|
||||
if b != nil {
|
||||
return proto.Unmarshal(data, b.(proto.Message))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func NewCodec(rwc io.ReadWriteCloser) codec.Codec {
|
||||
return &protoCodec{
|
||||
buf: bytes.NewBuffer(nil),
|
||||
rwc: rwc,
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user