vyvanse/vendor/github.com/drone/mq/stomp/client.go

package stomp

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"runtime/debug"
	"strconv"
	"sync"
	"time"

	"github.com/drone/mq/logger"
	"github.com/drone/mq/stomp/dialer"
)

// Client defines a client connection to a STOMP server.
type Client struct {
	mu sync.Mutex

	peer Peer
	subs map[string]Handler
	wait map[string]chan struct{}
	done chan error

	seq int64

	skipVerify      bool
	readBufferSize  int
	writeBufferSize int
	timeout         time.Duration
}

// New returns a new STOMP client using the given connection.
func New(peer Peer) *Client {
	return &Client{
		peer: peer,
		subs: make(map[string]Handler),
		wait: make(map[string]chan struct{}),
		done: make(chan error, 1),
	}
}

// Dial creates a client connection to the given target.
func Dial(target string) (*Client, error) {
	conn, err := dialer.Dial(target)
	if err != nil {
		return nil, err
	}
	return New(Conn(conn)), nil
}

// Send sends the data to the given destination.
func (c *Client) Send(dest string, data []byte, opts ...MessageOption) error {
	m := NewMessage()
	m.Method = MethodSend
	m.Dest = []byte(dest)
	m.Body = data
	m.Apply(opts...)
	return c.sendMessage(m)
}

// SendJSON sends the JSON encoding of v to the given destination.
func (c *Client) SendJSON(dest string, v interface{}, opts ...MessageOption) error {
	data, err := json.Marshal(v)
	if err != nil {
		return err
	}
	opts = append(opts,
		WithHeader("content-type", "application/json"),
	)
	return c.Send(dest, data, opts...)
}

// Subscribe subscribes to the given destination.
func (c *Client) Subscribe(dest string, handler Handler, opts ...MessageOption) (id []byte, err error) {
	id = c.incr()

	m := NewMessage()
	m.Method = MethodSubscribe
	m.ID = id
	m.Dest = []byte(dest)
	m.Apply(opts...)

	c.mu.Lock()
	c.subs[string(id)] = handler
	c.mu.Unlock()

	err = c.sendMessage(m)
	if err != nil {
		c.mu.Lock()
		delete(c.subs, string(id))
		c.mu.Unlock()
		return
	}
	return
}

// Unsubscribe unsubscribes to the destination.
func (c *Client) Unsubscribe(id []byte, opts ...MessageOption) error {
	c.mu.Lock()
	delete(c.subs, string(id))
	c.mu.Unlock()

	m := NewMessage()
	m.Method = MethodUnsubscribe
	m.ID = id
	m.Apply(opts...)

	return c.sendMessage(m)
}

// Ack acknowledges the messages with the given id.
func (c *Client) Ack(id []byte, opts ...MessageOption) error {
	m := NewMessage()
	m.Method = MethodAck
	m.ID = id
	m.Apply(opts...)

	return c.sendMessage(m)
}

// Nack negative-acknowledges the messages with the given id.
func (c *Client) Nack(id []byte, opts ...MessageOption) error {
	m := NewMessage()
	m.Method = MethodNack
	m.ID = id
	m.Apply(opts...)

	return c.peer.Send(m)
}

// Connect opens the connection and establishes the session.
func (c *Client) Connect(opts ...MessageOption) error {
	m := NewMessage()
	m.Proto = STOMP
	m.Method = MethodStomp
	m.Apply(opts...)
	if err := c.sendMessage(m); err != nil {
		return err
	}

	m, ok := <-c.peer.Receive()
	if !ok {
		return io.EOF
	}
	defer m.Release()

	if !bytes.Equal(m.Method, MethodConnected) {
		return fmt.Errorf("stomp: inbound message: unexpected method, want connected")
	}
	go c.listen()
	return nil
}

// Disconnect terminates the session and closes the connection.
func (c *Client) Disconnect() error {
	m := NewMessage()
	m.Method = MethodDisconnect
	c.sendMessage(m)
	return c.peer.Close()
}

// Done returns a channel
func (c *Client) Done() <-chan error {
	return c.done
}

func (c *Client) incr() []byte {
	c.mu.Lock()
	i := c.seq
	c.seq++
	c.mu.Unlock()
	return strconv.AppendInt(nil, i, 10)
}

func (c *Client) listen() {
	defer func() {
		if r := recover(); r != nil {
			logger.Warningf("stomp client: recover panic: %s", r)
			err, ok := r.(error)
			if !ok {
				logger.Warningf("%v: %s", r, debug.Stack())
				c.done <- fmt.Errorf("%v", r)
			} else {
				logger.Warningf("%s", err)
				c.done <- err
			}
		}
	}()

	for {
		m, ok := <-c.peer.Receive()
		if !ok {
			c.done <- io.EOF
			return
		}

		switch {
		case bytes.Equal(m.Method, MethodMessage):
			c.handleMessage(m)
		case bytes.Equal(m.Method, MethodRecipet):
			c.handleReceipt(m)
		default:
			logger.Noticef("stomp client: unknown message type: %s",
				string(m.Method),
			)
		}
	}
}

func (c *Client) handleReceipt(m *Message) {
	c.mu.Lock()
	receiptc, ok := c.wait[string(m.Receipt)]
	c.mu.Unlock()
	if !ok {
		logger.Noticef("stomp client: unknown read receipt: %s",
			string(m.Receipt),
		)
		return
	}
	receiptc <- struct{}{}
}

func (c *Client) handleMessage(m *Message) {
	c.mu.Lock()
	handler, ok := c.subs[string(m.Subs)]
	c.mu.Unlock()
	if !ok {
		logger.Noticef("stomp client: subscription not found: %s",
			string(m.Subs),
		)
		return
	}
	handler.Handle(m)
}

func (c *Client) sendMessage(m *Message) error {
	if len(m.Receipt) == 0 {
		return c.peer.Send(m)
	}

	receiptc := make(chan struct{}, 1)
	c.wait[string(m.Receipt)] = receiptc

	defer func() {
		delete(c.wait, string(m.Receipt))
	}()

	err := c.peer.Send(m)
	if err != nil {
		return err
	}

	select {
	case <-receiptc:
		return nil
	}
}