315 lines
7.3 KiB
Go
315 lines
7.3 KiB
Go
package client
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"sync"
|
|
|
|
"git.loafle.net/commons_go/logging"
|
|
"git.loafle.net/overflow/overflow_probes/central/client/protocol"
|
|
"git.loafle.net/overflow/overflow_probes/config"
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
const (
|
|
ProtocolName = "RPC/1.0"
|
|
)
|
|
|
|
type (
|
|
OnNotifyFunc func(method string, params []string)
|
|
OnCloseFunc func(code int, text string)
|
|
)
|
|
|
|
type ServerError string
|
|
|
|
func (e ServerError) Error() string {
|
|
return string(e)
|
|
}
|
|
|
|
var ErrShutdown = errors.New("connection is shut down")
|
|
|
|
type Call struct {
|
|
Method string // The name of the service and method to call.
|
|
Args []string // The argument to the function (*struct).
|
|
Result interface{} // The reply from the function (*struct).
|
|
Error error // After completion, the error status.
|
|
Done chan *Call // Strobes when call is complete.
|
|
}
|
|
|
|
func (c *Call) done() {
|
|
select {
|
|
case c.Done <- c:
|
|
// 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().
|
|
logging.Logger.Debug("Client: discarding Call reply due to insufficient Done chan capacity")
|
|
}
|
|
}
|
|
|
|
type Client interface {
|
|
Dial(url string, header http.Header) (*http.Response, error)
|
|
Call(method string, args []string, result interface{}) error
|
|
Notify(method string, args []string) error
|
|
OnNotify(cb OnNotifyFunc)
|
|
OnClose(cb OnCloseFunc)
|
|
Shutdown(ctx context.Context) error
|
|
}
|
|
|
|
type client struct {
|
|
conn *websocket.Conn
|
|
sendMutex sync.Mutex
|
|
request protocol.Request
|
|
notification protocol.Notification
|
|
mutex sync.Mutex
|
|
requestID uint64
|
|
pending map[uint64]*Call
|
|
closing bool // user has called Close
|
|
shutdown bool // server has told us to stop
|
|
onNotifyHandler OnNotifyFunc
|
|
onCloseHandlers []OnCloseFunc
|
|
}
|
|
|
|
func New() Client {
|
|
c := &client{
|
|
requestID: 0,
|
|
pending: make(map[uint64]*Call),
|
|
onCloseHandlers: make([]OnCloseFunc, 1),
|
|
}
|
|
|
|
return c
|
|
}
|
|
|
|
func (c *client) Dial(url string, header http.Header) (*http.Response, error) {
|
|
var err error
|
|
var res *http.Response
|
|
|
|
dialer := websocket.Dialer{
|
|
ReadBufferSize: config.CFG.Central.ReadBufferSize,
|
|
WriteBufferSize: config.CFG.Central.WriteBufferSize,
|
|
}
|
|
|
|
if c.conn, res, err = dialer.Dial(url, header); nil != err {
|
|
return nil, err
|
|
}
|
|
|
|
c.conn.SetCloseHandler(c.connCloseHandler)
|
|
|
|
go c.input()
|
|
|
|
return res, nil
|
|
}
|
|
|
|
func (c *client) Call(method string, args []string, result interface{}) error {
|
|
call := <-c.goCall(method, args, result, make(chan *Call, 1)).Done
|
|
return call.Error
|
|
}
|
|
|
|
func (c *client) Notify(method string, args []string) error {
|
|
c.sendMutex.Lock()
|
|
defer c.sendMutex.Unlock()
|
|
|
|
c.notification.Protocol = ProtocolName
|
|
c.notification.Method = method
|
|
c.notification.Params = args
|
|
|
|
if err := c.conn.WriteJSON(c.notification); nil != err {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (c *client) OnNotify(cb OnNotifyFunc) {
|
|
c.onNotifyHandler = cb
|
|
}
|
|
|
|
func (c *client) OnClose(cb OnCloseFunc) {
|
|
c.onCloseHandlers = append(c.onCloseHandlers, cb)
|
|
}
|
|
|
|
func (c *client) Shutdown(ctx context.Context) error {
|
|
c.mutex.Lock()
|
|
if c.closing {
|
|
c.mutex.Unlock()
|
|
return ErrShutdown
|
|
}
|
|
c.closing = true
|
|
c.mutex.Unlock()
|
|
|
|
return c.conn.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) goCall(method string, args []string, result interface{}, done chan *Call) *Call {
|
|
call := new(Call)
|
|
call.Method = method
|
|
call.Args = args
|
|
call.Result = result
|
|
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 {
|
|
logging.Logger.Panic("Client: done channel is unbuffered")
|
|
}
|
|
}
|
|
call.Done = done
|
|
c.sendCall(call)
|
|
return call
|
|
}
|
|
|
|
func (c *client) sendCall(call *Call) {
|
|
c.sendMutex.Lock()
|
|
defer c.sendMutex.Unlock()
|
|
|
|
// Register this call.
|
|
c.mutex.Lock()
|
|
if c.shutdown || c.closing {
|
|
call.Error = ErrShutdown
|
|
c.mutex.Unlock()
|
|
call.done()
|
|
return
|
|
}
|
|
c.requestID++
|
|
id := c.requestID
|
|
c.pending[id] = call
|
|
c.mutex.Unlock()
|
|
|
|
// Encode and send the request.
|
|
c.request.Protocol = ProtocolName
|
|
c.request.Method = call.Method
|
|
c.request.Params = call.Args
|
|
c.request.ID = id
|
|
if err := c.conn.WriteJSON(c.request); nil != err {
|
|
c.mutex.Lock()
|
|
call = c.pending[id]
|
|
delete(c.pending, id)
|
|
c.mutex.Unlock()
|
|
if call != nil {
|
|
call.Error = err
|
|
call.done()
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *client) input() {
|
|
var err error
|
|
var res protocol.Response
|
|
var noti protocol.Notification
|
|
var messageType int
|
|
var buff []byte
|
|
|
|
for err == nil {
|
|
if messageType, buff, err = c.conn.ReadMessage(); nil != err {
|
|
logging.Logger.Error(fmt.Sprintf("Client: Reader error[%v]", err))
|
|
continue
|
|
}
|
|
logging.Logger.Debug(fmt.Sprintf("Client: messageType:%d", messageType))
|
|
|
|
if err = json.Unmarshal(buff, ¬i); nil != err {
|
|
if err = json.Unmarshal(buff, &res); nil != err {
|
|
logging.Logger.Error(fmt.Sprintf("Client: Decode error[%v]", err))
|
|
continue
|
|
} else {
|
|
if err = c.onResponse(res); nil != err {
|
|
logging.Logger.Error(fmt.Sprintf("Client: Response error[%v]", err))
|
|
continue
|
|
}
|
|
}
|
|
} else {
|
|
c.onNotification(noti)
|
|
}
|
|
|
|
// if err = json.Unmarshal(buff, &res); nil != err {
|
|
// if err = json.Unmarshal(buff, ¬i); nil != err {
|
|
// logging.Logger.Error(fmt.Sprintf("Client: Decode error[%v]", err))
|
|
// continue
|
|
// } else {
|
|
// c.onNotification(noti)
|
|
// }
|
|
// } else {
|
|
// err = c.onResponse(res)
|
|
// }
|
|
}
|
|
// Terminate pending calls.
|
|
c.sendMutex.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.sendMutex.Unlock()
|
|
|
|
if err != io.EOF && !closing {
|
|
logging.Logger.Debug(fmt.Sprintf("Client: client protocol error:%v", err))
|
|
}
|
|
}
|
|
|
|
func (c *client) onResponse(res protocol.Response) error {
|
|
var err error
|
|
id := res.ID
|
|
c.mutex.Lock()
|
|
call := c.pending[id]
|
|
delete(c.pending, id)
|
|
c.mutex.Unlock()
|
|
|
|
switch {
|
|
case call == nil:
|
|
|
|
case res.Error != nil:
|
|
// We've got an error response. Give this to the request;
|
|
// any subsequent requests will get the ReadResponseBody
|
|
// error if there is one.
|
|
if protocol.ProtocolErrorCodeInternal == res.Error.Code {
|
|
if nil != res.Error.Message {
|
|
call.Error = ServerError(*res.Error.Message)
|
|
}
|
|
}
|
|
|
|
call.done()
|
|
default:
|
|
if err = json.Unmarshal(*res.Result, call.Result); nil != err {
|
|
call.Error = errors.New("reading body " + err.Error())
|
|
}
|
|
call.done()
|
|
}
|
|
|
|
return err
|
|
}
|
|
|
|
func (c *client) onNotification(noti protocol.Notification) {
|
|
if nil == c.onNotifyHandler {
|
|
return
|
|
}
|
|
|
|
c.onNotifyHandler(noti.Method, noti.Params)
|
|
}
|
|
|
|
func (c *client) connCloseHandler(code int, text string) error {
|
|
for _, h := range c.onCloseHandlers {
|
|
h(code, text)
|
|
}
|
|
|
|
return nil
|
|
}
|