260 lines
4.9 KiB
Go
260 lines
4.9 KiB
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
|
||
|
}
|
||
|
}
|