overflow_probes/central/client/client.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, &noti); 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, &noti); 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
}