+++ /dev/null
-package websocket
-
-import (
- "container/list"
- "encoding/json"
- "fmt"
- "io"
- "net/http"
- "sync"
- "time"
-
- "github.com/gorilla/websocket"
- log "github.com/sirupsen/logrus"
-
- "github.com/vapor/errors"
-)
-
-// websocketSendBufferSize is the number of elements the send channel
-// can queue before blocking. Note that this only applies to requests
-// handled directly in the websocket client input handler or the async
-// handler since notifications have their own queuing mechanism
-// independent of the send channel buffer.
-const (
- logModule = "websocket"
- websocketSendBufferSize = 50
-)
-
-var (
- // ErrWSParse means a request parsing error
- ErrWSParse = errors.New("Websocket request parsing error")
- // ErrWSInternal means service handling errors
- ErrWSInternal = errors.New("Websocket Internal error")
- // ErrWSClientQuit means the websocket client is disconnected
- ErrWSClientQuit = errors.New("Websocket client quit")
-
- // timeZeroVal is simply the zero value for a time.Time and is used to avoid creating multiple instances.
- timeZeroVal time.Time
-)
-
-type semaphore chan struct{}
-
-func makeSemaphore(n int) semaphore {
- return make(chan struct{}, n)
-}
-
-func (s semaphore) acquire() { s <- struct{}{} }
-func (s semaphore) release() { <-s }
-
-// wsTopicHandler describes a callback function used to handle a specific topic.
-type wsTopicHandler func(*WSClient)
-
-// wsHandlers maps websocket topic strings to appropriate websocket handler
-// functions. This is set by init because help references wsHandlers and thus
-// causes a dependency loop.
-var wsHandlers = map[string]wsTopicHandler{
- "notify_raw_blocks": handleNotifyBlocks,
- "notify_new_transactions": handleNotifyNewTransactions,
- "stop_notify_raw_blocks": handleStopNotifyBlocks,
- "stop_notify_new_transactions": handleStopNotifyNewTransactions,
-}
-
-// responseMessage houses a message to send to a connected websocket client as
-// well as a channel to reply on when the message is sent.
-type responseMessage struct {
- msg []byte
- doneChan chan bool
-}
-
-// WSClient provides an abstraction for handling a websocket client. The
-// overall data flow is split into 3 main goroutines, a possible 4th goroutine
-// for long-running operations (only started if request is made), and a
-// websocket manager which is used to allow things such as broadcasting
-// requested notifications to all connected websocket clients. Inbound
-// messages are read via the inHandler goroutine and generally dispatched to
-// their own handler. However, certain potentially long-running operations such
-// as rescans, are sent to the asyncHander goroutine and are limited to one at a
-// time. There are two outbound message types - one for responding to client
-// requests and another for async notifications. Responses to client requests
-// use SendMessage which employs a buffered channel thereby limiting the number
-// of outstanding requests that can be made. Notifications are sent via
-// QueueNotification which implements a queue via notificationQueueHandler to
-// ensure sending notifications from other subsystems can't block. Ultimately,
-// all messages are sent via the outHandler.
-type WSClient struct {
- sync.Mutex
- conn *websocket.Conn
- // disconnected indicated whether or not the websocket client is disconnected.
- disconnected bool
- // addr is the remote address of the client.
- addr string
- serviceRequestSem semaphore
- ntfnChan chan []byte
- sendChan chan responseMessage
- quit chan struct{}
- wg sync.WaitGroup
- notificationMgr *WSNotificationManager
-}
-
-// NewWebsocketClient means to create a new object to the connected websocket client
-func NewWebsocketClient(w http.ResponseWriter, r *http.Request, notificationMgr *WSNotificationManager) (*WSClient, error) {
- // Limit max number of websocket clients.
- if notificationMgr.IsMaxConnect() {
- return nil, fmt.Errorf("numOfMaxWS: %d, disconnecting: %s", notificationMgr.MaxNumWebsockets, r.RemoteAddr)
- }
-
- // Attempt to upgrade the connection to a websocket connection using the default size for read/write buffers.
- conn, err := websocket.Upgrade(w, r, nil, 0, 0)
- if err != nil {
- return nil, err
- }
-
- conn.SetReadDeadline(timeZeroVal)
-
- client := &WSClient{
- conn: conn,
- addr: r.RemoteAddr,
- serviceRequestSem: makeSemaphore(notificationMgr.maxNumConcurrentReqs),
- ntfnChan: make(chan []byte, 1), // nonblocking sync
- sendChan: make(chan responseMessage, websocketSendBufferSize),
- quit: make(chan struct{}),
- notificationMgr: notificationMgr,
- }
- return client, nil
-}
-
-// inHandler handles all incoming messages for the websocket connection.
-func (c *WSClient) inHandler() {
-out:
- for {
- // Break out of the loop once the quit channel has been closed.
- // Use a non-blocking select here so we fall through otherwise.
- select {
- case <-c.quit:
- break out
- default:
- }
-
- _, msg, err := c.conn.ReadMessage()
- if err != nil {
- if err != io.EOF {
- log.WithFields(log.Fields{"module": logModule, "remoteAddress": c.addr, "error": err}).Error("Websocket receive error")
- }
- break out
- }
-
- var request WSRequest
- if err = json.Unmarshal(msg, &request); err != nil {
- respError := errors.Wrap(err, ErrWSParse)
- resp := NewWSResponse(NTRequestStatus.String(), nil, respError)
- reply, err := json.Marshal(resp)
- if err != nil {
- log.WithFields(log.Fields{"module": logModule, "error": err}).Error("Failed to marshal parse failure reply")
- continue
- }
-
- c.SendMessage(reply, nil)
- continue
- }
-
- c.serviceRequestSem.acquire()
- go func() {
- c.serviceRequest(request.Topic)
- c.serviceRequestSem.release()
- }()
- }
-
- // Ensure the connection is closed.
- c.Disconnect()
- c.wg.Done()
- log.WithFields(log.Fields{"module": logModule, "remoteAddress": c.addr}).Debug("Websocket client input handler done")
-}
-
-func (c *WSClient) serviceRequest(topic string) {
- var respErr error
-
- if wsHandler, ok := wsHandlers[topic]; ok {
- wsHandler(c)
- } else {
- err := fmt.Errorf("There is not this topic: %s", topic)
- respErr = errors.Wrap(err, ErrWSInternal)
- log.WithFields(log.Fields{"module": logModule, "topic": topic}).Debug("There is not this topic")
- }
-
- resp := NewWSResponse(NTRequestStatus.String(), nil, respErr)
- reply, err := json.Marshal(resp)
- if err != nil {
- log.WithFields(log.Fields{"module": logModule, "error": err}).Debug("Failed to marshal parse failure reply")
- return
- }
-
- c.SendMessage(reply, nil)
-}
-
-// notificationQueueHandler handles the queuing of outgoing notifications for the websocket client.
-func (c *WSClient) notificationQueueHandler() {
- ntfnSentChan := make(chan bool, 1) // nonblocking sync
-
- // pendingNtfns is used as a queue for notifications that are ready to
- // be sent once there are no outstanding notifications currently being
- // sent.
- pendingNtfns := list.New()
- waiting := false
-out:
- for {
- select {
- // This channel is notified when a message is being queued to
- // be sent across the network socket. It will either send the
- // message immediately if a send is not already in progress, or
- // queue the message to be sent once the other pending messages
- // are sent.
- case msg := <-c.ntfnChan:
- if !waiting {
- c.SendMessage(msg, ntfnSentChan)
- } else {
- pendingNtfns.PushBack(msg)
- }
- waiting = true
- // This channel is notified when a notification has been sent across the network socket.
- case <-ntfnSentChan:
- // No longer waiting if there are no more messages in the pending messages queue.
- next := pendingNtfns.Front()
- if next == nil {
- waiting = false
- continue
- }
-
- // Notify the outHandler about the next item to asynchronously send.
- msg := pendingNtfns.Remove(next).([]byte)
- c.SendMessage(msg, ntfnSentChan)
- case <-c.quit:
- break out
- }
- }
-
- // Drain any wait channels before exiting so nothing is left waiting around to send.
-cleanup:
- for {
- select {
- case <-c.ntfnChan:
- case <-ntfnSentChan:
- default:
- break cleanup
- }
- }
- c.wg.Done()
- log.WithFields(log.Fields{"module": logModule, "remoteAddress": c.addr}).Debug("Websocket client notification queue handler done")
-}
-
-// outHandler handles all outgoing messages for the websocket connection.
-func (c *WSClient) outHandler() {
-out:
- for {
- // Send any messages ready for send until the quit channel is closed.
- select {
- case r := <-c.sendChan:
- if err := c.conn.WriteMessage(websocket.TextMessage, r.msg); err != nil {
- log.WithFields(log.Fields{"module": logModule, "error": err}).Error("Failed to send message to wesocket client")
- c.Disconnect()
- break out
- }
- if r.doneChan != nil {
- r.doneChan <- true
- }
- case <-c.quit:
- break out
- }
- }
-
- // Drain any wait channels before exiting so nothing is left waiting around to send.
-cleanup:
- for {
- select {
- case r := <-c.sendChan:
- if r.doneChan != nil {
- r.doneChan <- false
- }
- default:
- break cleanup
- }
- }
- c.wg.Done()
- log.WithFields(log.Fields{"module": logModule, "remoteAddress": c.addr}).Debug("Websocket client output handler done")
-}
-
-// SendMessage sends the passed json to the websocket client. It is backed
-// by a buffered channel, so it will not block until the send channel is full.
-// Note however that QueueNotification must be used for sending async
-// notifications instead of the this function. This approach allows a limit to
-// the number of outstanding requests a client can make without preventing or
-// blocking on async notifications.
-func (c *WSClient) SendMessage(marshalledJSON []byte, doneChan chan bool) {
- // Don't send the message if disconnected.
- if c.Disconnected() {
- if doneChan != nil {
- doneChan <- false
- }
- return
- }
-
- c.sendChan <- responseMessage{msg: marshalledJSON, doneChan: doneChan}
-}
-
-// QueueNotification queues the passed notification to be sent to the websocket client.
-func (c *WSClient) QueueNotification(marshalledJSON []byte) error {
- // Don't queue the message if disconnected.
- if c.Disconnected() {
- return ErrWSClientQuit
- }
-
- c.ntfnChan <- marshalledJSON
- return nil
-}
-
-// Disconnected returns whether or not the websocket client is disconnected.
-func (c *WSClient) Disconnected() bool {
- c.Lock()
- defer c.Unlock()
-
- return c.disconnected
-}
-
-// Disconnect disconnects the websocket client.
-func (c *WSClient) Disconnect() {
- c.Lock()
- defer c.Unlock()
-
- // Nothing to do if already disconnected.
- if c.disconnected {
- return
- }
-
- log.WithFields(log.Fields{"module": logModule, "remoteAddress": c.addr}).Info("Disconnecting websocket client")
-
- close(c.quit)
- c.conn.Close()
- c.disconnected = true
-}
-
-// Start begins processing input and output messages.
-func (c *WSClient) Start() {
- log.WithFields(log.Fields{"module": logModule, "remoteAddress": c.addr}).Info("Starting websocket client")
-
- c.wg.Add(3)
- go c.inHandler()
- go c.notificationQueueHandler()
- go c.outHandler()
-}
-
-// WaitForShutdown blocks until the websocket client goroutines are stopped and the connection is closed.
-func (c *WSClient) WaitForShutdown() {
- c.wg.Wait()
-}
-
-// handleNotifyBlocks implements the notifyblocks topic extension for websocket connections.
-func handleNotifyBlocks(wsc *WSClient) {
- wsc.notificationMgr.RegisterBlockUpdates(wsc)
-}
-
-// handleStopNotifyBlocks implements the stopnotifyblocks topic extension for websocket connections.
-func handleStopNotifyBlocks(wsc *WSClient) {
- wsc.notificationMgr.UnregisterBlockUpdates(wsc)
-}
-
-// handleNotifyNewTransations implements the notifynewtransactions topic extension for websocket connections.
-func handleNotifyNewTransactions(wsc *WSClient) {
- wsc.notificationMgr.RegisterNewMempoolTxsUpdates(wsc)
-}
-
-// handleStopNotifyNewTransations implements the stopnotifynewtransactions topic extension for websocket connections.
-func handleStopNotifyNewTransactions(wsc *WSClient) {
- wsc.notificationMgr.UnregisterNewMempoolTxsUpdates(wsc)
-}
OSDN Git Service
