2021-11-21 15:25:11 +00:00

365 lines
9.7 KiB
Go

// Package rpc2 provides bi-directional RPC client and server similar to net/rpc.
package rpc2
import (
"context"
"errors"
"io"
"log"
"reflect"
"sync"
)
// Client represents an RPC Client.
// There may be multiple outstanding Calls associated
// with a single Client, and a Client may be used by
// multiple goroutines simultaneously.
type Client struct {
mutex sync.Mutex // protects pending, seq, request
sending sync.Mutex
request Request // temp area used in send()
seq uint64
pending map[uint64]*Call
closing bool
shutdown bool
server bool
codec Codec
handlers map[string]*handler
disconnect chan struct{}
State *State // additional information to associate with client
blocking bool // whether to block request handling
}
// NewClient returns a new Client to handle requests to the
// set of services at the other end of the connection.
// It adds a buffer to the write side of the connection so
// the header and payload are sent as a unit.
func NewClient(conn io.ReadWriteCloser) *Client {
return NewClientWithCodec(NewGobCodec(conn))
}
// NewClientWithCodec is like NewClient but uses the specified
// codec to encode requests and decode responses.
func NewClientWithCodec(codec Codec) *Client {
return &Client{
codec: codec,
pending: make(map[uint64]*Call),
handlers: make(map[string]*handler),
disconnect: make(chan struct{}),
seq: 1, // 0 means notification.
}
}
// SetBlocking puts the client in blocking mode.
// In blocking mode, received requests are processes synchronously.
// If you have methods that may take a long time, other subsequent requests may time out.
func (c *Client) SetBlocking(blocking bool) {
c.blocking = blocking
}
// Run the client's read loop.
// You must run this method before calling any methods on the server.
func (c *Client) Run() {
c.readLoop()
}
// DisconnectNotify returns a channel that is closed
// when the client connection has gone away.
func (c *Client) DisconnectNotify() chan struct{} {
return c.disconnect
}
// Handle registers the handler function for the given method. If a handler already exists for method, Handle panics.
func (c *Client) Handle(method string, handlerFunc interface{}) {
addHandler(c.handlers, method, handlerFunc)
}
// readLoop reads messages from codec.
// It reads a reqeust or a response to the previous request.
// If the message is request, calls the handler function.
// If the message is response, sends the reply to the associated call.
func (c *Client) readLoop() {
var err error
var req Request
var resp Response
for err == nil {
req = Request{}
resp = Response{}
if err = c.codec.ReadHeader(&req, &resp); err != nil {
break
}
if req.Method != "" {
// request comes to server
if err = c.readRequest(&req); err != nil {
debugln("rpc2: error reading request:", err.Error())
}
} else {
// response comes to client
if err = c.readResponse(&resp); err != nil {
debugln("rpc2: error reading response:", err.Error())
}
}
}
// Terminate pending calls.
c.sending.Lock()
c.mutex.Lock()
c.shutdown = true
closing := c.closing
if err == io.EOF {
if closing {
err = ErrShutdown
} else {
err = io.ErrUnexpectedEOF
}
}
for _, call := range c.pending {
call.Error = err
call.done()
}
c.mutex.Unlock()
c.sending.Unlock()
if err != io.EOF && !closing && !c.server {
debugln("rpc2: client protocol error:", err)
}
close(c.disconnect)
if !closing {
c.codec.Close()
}
}
func (c *Client) handleRequest(req Request, method *handler, argv reflect.Value) {
// Invoke the method, providing a new value for the reply.
replyv := reflect.New(method.replyType.Elem())
returnValues := method.fn.Call([]reflect.Value{reflect.ValueOf(c), argv, replyv})
// Do not send response if request is a notification.
if req.Seq == 0 {
return
}
// The return value for the method is an error.
errInter := returnValues[0].Interface()
errmsg := ""
if errInter != nil {
errmsg = errInter.(error).Error()
}
resp := &Response{
Seq: req.Seq,
Error: errmsg,
}
if err := c.codec.WriteResponse(resp, replyv.Interface()); err != nil {
debugln("rpc2: error writing response:", err.Error())
}
}
func (c *Client) readRequest(req *Request) error {
method, ok := c.handlers[req.Method]
if !ok {
resp := &Response{
Seq: req.Seq,
Error: "rpc2: can't find method " + req.Method,
}
return c.codec.WriteResponse(resp, resp)
}
// Decode the argument value.
var argv reflect.Value
argIsValue := false // if true, need to indirect before calling.
if method.argType.Kind() == reflect.Ptr {
argv = reflect.New(method.argType.Elem())
} else {
argv = reflect.New(method.argType)
argIsValue = true
}
// argv guaranteed to be a pointer now.
if err := c.codec.ReadRequestBody(argv.Interface()); err != nil {
return err
}
if argIsValue {
argv = argv.Elem()
}
if c.blocking {
c.handleRequest(*req, method, argv)
} else {
go c.handleRequest(*req, method, argv)
}
return nil
}
func (c *Client) readResponse(resp *Response) error {
seq := resp.Seq
c.mutex.Lock()
call := c.pending[seq]
delete(c.pending, seq)
c.mutex.Unlock()
var err error
switch {
case call == nil:
// We've got no pending call. That usually means that
// WriteRequest partially failed, and call was already
// removed; response is a server telling us about an
// error reading request body. We should still attempt
// to read error body, but there's no one to give it to.
err = c.codec.ReadResponseBody(nil)
if err != nil {
err = errors.New("reading error body: " + err.Error())
}
case resp.Error != "":
// We've got an error response. Give this to the request;
// any subsequent requests will get the ReadResponseBody
// error if there is one.
call.Error = ServerError(resp.Error)
err = c.codec.ReadResponseBody(nil)
if err != nil {
err = errors.New("reading error body: " + err.Error())
}
call.done()
default:
err = c.codec.ReadResponseBody(call.Reply)
if err != nil {
call.Error = errors.New("reading body " + err.Error())
}
call.done()
}
return err
}
// Close waits for active calls to finish and closes the codec.
func (c *Client) Close() error {
c.mutex.Lock()
if c.shutdown || c.closing {
c.mutex.Unlock()
return ErrShutdown
}
c.closing = true
c.mutex.Unlock()
return c.codec.Close()
}
// Go invokes the function asynchronously. It returns the Call structure representing
// the invocation. The done channel will signal when the call is complete by returning
// the same Call object. If done is nil, Go will allocate a new channel.
// If non-nil, done must be buffered or Go will deliberately crash.
func (c *Client) Go(method string, args interface{}, reply interface{}, done chan *Call) *Call {
call := new(Call)
call.Method = method
call.Args = args
call.Reply = reply
if done == nil {
done = make(chan *Call, 10) // buffered.
} else {
// If caller passes done != nil, it must arrange that
// done has enough buffer for the number of simultaneous
// RPCs that will be using that channel. If the channel
// is totally unbuffered, it's best not to run at all.
if cap(done) == 0 {
log.Panic("rpc2: done channel is unbuffered")
}
}
call.Done = done
c.send(call)
return call
}
// CallWithContext invokes the named function, waits for it to complete, and
// returns its error status, or an error from Context timeout.
func (c *Client) CallWithContext(ctx context.Context, method string, args interface{}, reply interface{}) error {
call := c.Go(method, args, reply, make(chan *Call, 1))
select {
case <-call.Done:
return call.Error
case <-ctx.Done():
return ctx.Err()
}
return nil
}
// Call invokes the named function, waits for it to complete, and returns its error status.
func (c *Client) Call(method string, args interface{}, reply interface{}) error {
return c.CallWithContext(context.Background(), method, args, reply)
}
func (call *Call) done() {
select {
case call.Done <- call:
// ok
default:
// We don't want to block here. It is the caller's responsibility to make
// sure the channel has enough buffer space. See comment in Go().
debugln("rpc2: discarding Call reply due to insufficient Done chan capacity")
}
}
// ServerError represents an error that has been returned from
// the remote side of the RPC connection.
type ServerError string
func (e ServerError) Error() string {
return string(e)
}
// ErrShutdown is returned when the connection is closing or closed.
var ErrShutdown = errors.New("connection is shut down")
// Call represents an active RPC.
type Call struct {
Method string // The name of the service and method to call.
Args interface{} // The argument to the function (*struct).
Reply interface{} // The reply from the function (*struct).
Error error // After completion, the error status.
Done chan *Call // Strobes when call is complete.
}
func (c *Client) send(call *Call) {
c.sending.Lock()
defer c.sending.Unlock()
// Register this call.
c.mutex.Lock()
if c.shutdown || c.closing {
call.Error = ErrShutdown
c.mutex.Unlock()
call.done()
return
}
seq := c.seq
c.seq++
c.pending[seq] = call
c.mutex.Unlock()
// Encode and send the request.
c.request.Seq = seq
c.request.Method = call.Method
err := c.codec.WriteRequest(&c.request, call.Args)
if err != nil {
c.mutex.Lock()
call = c.pending[seq]
delete(c.pending, seq)
c.mutex.Unlock()
if call != nil {
call.Error = err
call.done()
}
}
}
// Notify sends a request to the receiver but does not wait for a return value.
func (c *Client) Notify(method string, args interface{}) error {
c.sending.Lock()
defer c.sending.Unlock()
if c.shutdown || c.closing {
return ErrShutdown
}
c.request.Seq = 0
c.request.Method = method
return c.codec.WriteRequest(&c.request, args)
}