tacitus/websocket/client.go

package websocket

import (
	"github.com/kcotugno/tacitus/gdax"

	"github.com/gorilla/websocket"

	"sync"
	"time"
)

const ENDPOINT = "wss://ws-feed.gdax.com"

type Client struct {
	msgs   chan gdax.Message
	err    chan error
	writer chan gdax.Request
	conn   *websocket.Conn

	conMux  sync.Mutex
	connect bool
	close   chan bool
}

func NewClient() *Client {
	c := Client{}
	c.msgs = make(chan gdax.Message, 1024)
	c.err = make(chan error)
	c.writer = make(chan gdax.Request)
	c.close = make(chan bool)

	return &c
}

func (c *Client) Open() error {
	var err error

	c.conn, _, err = websocket.DefaultDialer.Dial(ENDPOINT, nil)
	if err != nil {
		return err
	}
	c.setConnected(true)

	go func() {
		for c.connected() {
			var msg gdax.Message

			err := c.conn.ReadJSON(&msg)
			if err != nil {
				if websocket.IsUnexpectedCloseError(err,
					websocket.CloseNormalClosure) {
					c.sendClose(false)
					c.err <- err
					break
				}
			}
			c.msgs <- msg
		}
	}()

	go func() {
		for c.connected() {
			select {
			case msg := <-c.writer:
				err := c.conn.WriteJSON(msg)
				if err != nil {
					c.sendClose(false)
					c.err <- err
				}
			case send := <-c.close:
				if send {
					c.conn.WriteControl(websocket.CloseMessage,
						websocket.FormatCloseMessage(websocket.CloseNormalClosure,
							"done"),
						time.Now().Add(30*time.Second))
				}
			}
		}
	}()

	return nil
}

func (c *Client) Close() {
	c.sendClose(true)
}

func (c *Client) sendClose(send bool) {
	if c.connected() {
		c.setConnected(false)
		c.close <- send
	}
}

func (c *Client) Send(msg gdax.Request) {
	if c.connected() {
		c.writer <- msg
	}
}

func (c *Client) MessageStream() <-chan gdax.Message {
	return c.msgs
}

func (c *Client) Error() <-chan error {
	return c.err
}

func (c *Client) connected() bool {
	c.conMux.Lock()
	defer c.conMux.Unlock()

	return c.connect
}

func (c *Client) setConnected(connected bool) {
	c.conMux.Lock()
	defer c.conMux.Unlock()

	c.connect = connected
}