126 lines
3.0 KiB
Go
126 lines
3.0 KiB
Go
package rpc2
|
|
|
|
import (
|
|
"bufio"
|
|
"encoding/gob"
|
|
"io"
|
|
"sync"
|
|
)
|
|
|
|
// A Codec implements reading and writing of RPC requests and responses.
|
|
// The client calls ReadHeader to read a message header.
|
|
// The implementation must populate either Request or Response argument.
|
|
// Depending on which argument is populated, ReadRequestBody or
|
|
// ReadResponseBody is called right after ReadHeader.
|
|
// ReadRequestBody and ReadResponseBody may be called with a nil
|
|
// argument to force the body to be read and then discarded.
|
|
type Codec interface {
|
|
// ReadHeader must read a message and populate either the request
|
|
// or the response by inspecting the incoming message.
|
|
ReadHeader(*Request, *Response) error
|
|
|
|
// ReadRequestBody into args argument of handler function.
|
|
ReadRequestBody(interface{}) error
|
|
|
|
// ReadResponseBody into reply argument of handler function.
|
|
ReadResponseBody(interface{}) error
|
|
|
|
// WriteRequest must be safe for concurrent use by multiple goroutines.
|
|
WriteRequest(*Request, interface{}) error
|
|
|
|
// WriteResponse must be safe for concurrent use by multiple goroutines.
|
|
WriteResponse(*Response, interface{}) error
|
|
|
|
// Close is called when client/server finished with the connection.
|
|
Close() error
|
|
}
|
|
|
|
// Request is a header written before every RPC call.
|
|
type Request struct {
|
|
Seq uint64 // sequence number chosen by client
|
|
Method string
|
|
}
|
|
|
|
// Response is a header written before every RPC return.
|
|
type Response struct {
|
|
Seq uint64 // echoes that of the request
|
|
Error string // error, if any.
|
|
}
|
|
|
|
type gobCodec struct {
|
|
rwc io.ReadWriteCloser
|
|
dec *gob.Decoder
|
|
enc *gob.Encoder
|
|
encBuf *bufio.Writer
|
|
mutex sync.Mutex
|
|
}
|
|
|
|
type message struct {
|
|
Seq uint64
|
|
Method string
|
|
Error string
|
|
}
|
|
|
|
// NewGobCodec returns a new rpc2.Codec using gob encoding/decoding on conn.
|
|
func NewGobCodec(conn io.ReadWriteCloser) Codec {
|
|
buf := bufio.NewWriter(conn)
|
|
return &gobCodec{
|
|
rwc: conn,
|
|
dec: gob.NewDecoder(conn),
|
|
enc: gob.NewEncoder(buf),
|
|
encBuf: buf,
|
|
}
|
|
}
|
|
|
|
func (c *gobCodec) ReadHeader(req *Request, resp *Response) error {
|
|
var msg message
|
|
if err := c.dec.Decode(&msg); err != nil {
|
|
return err
|
|
}
|
|
|
|
if msg.Method != "" {
|
|
req.Seq = msg.Seq
|
|
req.Method = msg.Method
|
|
} else {
|
|
resp.Seq = msg.Seq
|
|
resp.Error = msg.Error
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (c *gobCodec) ReadRequestBody(body interface{}) error {
|
|
return c.dec.Decode(body)
|
|
}
|
|
|
|
func (c *gobCodec) ReadResponseBody(body interface{}) error {
|
|
return c.dec.Decode(body)
|
|
}
|
|
|
|
func (c *gobCodec) WriteRequest(r *Request, body interface{}) (err error) {
|
|
c.mutex.Lock()
|
|
defer c.mutex.Unlock()
|
|
if err = c.enc.Encode(r); err != nil {
|
|
return
|
|
}
|
|
if err = c.enc.Encode(body); err != nil {
|
|
return
|
|
}
|
|
return c.encBuf.Flush()
|
|
}
|
|
|
|
func (c *gobCodec) WriteResponse(r *Response, body interface{}) (err error) {
|
|
c.mutex.Lock()
|
|
defer c.mutex.Unlock()
|
|
if err = c.enc.Encode(r); err != nil {
|
|
return
|
|
}
|
|
if err = c.enc.Encode(body); err != nil {
|
|
return
|
|
}
|
|
return c.encBuf.Flush()
|
|
}
|
|
|
|
func (c *gobCodec) Close() error {
|
|
return c.rwc.Close()
|
|
}
|