diff --git a/conn.go b/conn.go index a39a658..4ac6b5e 100644 --- a/conn.go +++ b/conn.go @@ -20,14 +20,19 @@ const DefaultHeartBeatError = 5 * time.Second // Default send timeout in Conn.Send function const DefaultMsgSendTimeout = 10 * time.Second +const defaultReceiptTimeout = 30 * time.Second + // Default receipt timeout in Conn.Send function -const DefaultRcvReceiptTimeout = 30 * time.Second +const DefaultRcvReceiptTimeout = defaultReceiptTimeout // Default receipt timeout in Conn.Disconnect function -const DefaultDisconnectReceiptTimeout = 30 * time.Second +const DefaultDisconnectReceiptTimeout = defaultReceiptTimeout // Default receipt timeout in Subscription.Unsubscribe function -const DefaultUnsubscribeReceiptTimeout = 30 * time.Second +const DefaultUnsubscribeReceiptTimeout = defaultReceiptTimeout + +// Default receipt timeout in Conn.Subscribe function, when SubscribeOpt.Receipt is used +const DefaultSubscribeReceiptTimeout = defaultReceiptTimeout // Reply-To header used for temporary queues/RPC with rabbit. const ReplyToHeader = "reply-to" @@ -47,6 +52,7 @@ type Conn struct { rcvReceiptTimeout time.Duration disconnectReceiptTimeout time.Duration unsubscribeReceiptTimeout time.Duration + subscribeReceiptTimeout time.Duration hbGracePeriodMultiplier float64 closed bool closeMutex *sync.Mutex @@ -60,6 +66,11 @@ type Conn struct { type writeRequest struct { Frame *frame.Frame // frame to send C chan *frame.Frame // response channel + + // ReceiptC, if non-nil, receives a resulting RECEIPT frame instead of C. + // Used by SUBSCRIBE requests, where C is the subscription's own message + // channel and must not be closed by the RECEIPT. + ReceiptC chan *frame.Frame } // Dial creates a network connection to a STOMP server and performs @@ -232,6 +243,7 @@ func ConnectWithContext(ctx context.Context, conn io.ReadWriteCloser, opts ...fu c.rcvReceiptTimeout = options.RcvReceiptTimeout c.disconnectReceiptTimeout = options.DisconnectReceiptTimeout c.unsubscribeReceiptTimeout = options.UnsubscribeReceiptTimeout + c.subscribeReceiptTimeout = options.SubscribeReceiptTimeout if options.ResponseHeadersCallback != nil { options.ResponseHeadersCallback(response.Header) @@ -408,8 +420,10 @@ func processLoop(c *Conn, writer *frame.Writer) { sendError(channels, errors.New("write channel closed")) return } - if req.C != nil { - if receipt, ok := req.Frame.Header.Contains(frame.Receipt); ok { + if receipt, ok := req.Frame.Header.Contains(frame.Receipt); ok { + if req.ReceiptC != nil { + channels[receipt] = req.ReceiptC + } else if req.C != nil { // remember the channel for this receipt channels[receipt] = req.C } @@ -711,9 +725,11 @@ func (c *Conn) sendFrame(f *frame.Frame) error { // will be received by this subscription. A subscription has a channel // on which the calling program can receive messages. func (c *Conn) Subscribe(destination string, ack AckMode, opts ...func(*frame.Frame) error) (*Subscription, error) { + // Not deferred: released before waiting on the receipt below, so a slow + // server doesn't stall other operations on this connection. c.closeMutex.Lock() - defer c.closeMutex.Unlock() if c.closed { + c.closeMutex.Unlock() c.conn.Close() return nil, ErrClosedUnexpectedly } @@ -730,6 +746,7 @@ func (c *Conn) Subscribe(destination string, ack AckMode, opts ...func(*frame.Fr } err := opt(subscribeFrame) if err != nil { + c.closeMutex.Unlock() return nil, err } } @@ -738,6 +755,13 @@ func (c *Conn) Subscribe(destination string, ack AckMode, opts ...func(*frame.Fr if replyToSet { subscribeFrame.Header.Set(frame.Id, replyTo) + + // Reply-to subscriptions are never sent to the server, so no RECEIPT + // can ever arrive for one. + if _, ok := subscribeFrame.Header.Contains(frame.Receipt); ok { + c.closeMutex.Unlock() + return nil, ErrReceiptNotSupportedForReplyTo + } } // If the option functions have not specified the "id" header entry, @@ -753,6 +777,12 @@ func (c *Conn) Subscribe(destination string, ack AckMode, opts ...func(*frame.Fr C: ch, } + var receiptC chan *frame.Frame + if _, ok := subscribeFrame.Header.Contains(frame.Receipt); ok { + receiptC = make(chan *frame.Frame, 1) + request.ReceiptC = receiptC + } + closeMutex := &sync.Mutex{} sub := &Subscription{ id: id, @@ -774,8 +804,21 @@ func (c *Conn) Subscribe(destination string, ack AckMode, opts ...func(*frame.Fr // TODO is this safe? There is no check if writeCh is actually open. err := sendDataToWriteChWithTimeout(c.writeCh, request, c.msgSendTimeout) if err != nil { + c.closeMutex.Unlock() + // Request never reached processLoop; readLoop would otherwise block on ch forever. + close(ch) return nil, err } + c.closeMutex.Unlock() + + if receiptC != nil { + if err := readReceiptWithTimeout(receiptC, c.subscribeReceiptTimeout, ErrSubscribeReceiptTimeout); err != nil { + // The caller gets no handle to sub, so tear it down here or it leaks. + sub.abandon() + return nil, err + } + } + return sub, nil } diff --git a/conn_options.go b/conn_options.go index ff82bcd..176b876 100644 --- a/conn_options.go +++ b/conn_options.go @@ -21,6 +21,7 @@ type connOptions struct { RcvReceiptTimeout time.Duration DisconnectReceiptTimeout time.Duration UnsubscribeReceiptTimeout time.Duration + SubscribeReceiptTimeout time.Duration HeartBeatGracePeriodMultiplier float64 Login, Passcode string AcceptVersions []string @@ -42,6 +43,7 @@ func newConnOptions(conn *Conn, opts []func(*Conn) error) (*connOptions, error) RcvReceiptTimeout: DefaultRcvReceiptTimeout, DisconnectReceiptTimeout: DefaultDisconnectReceiptTimeout, UnsubscribeReceiptTimeout: DefaultUnsubscribeReceiptTimeout, + SubscribeReceiptTimeout: DefaultSubscribeReceiptTimeout, Logger: log.StdLogger{}, } @@ -163,6 +165,12 @@ var ConnOpt struct { // avoid deadlocks. If this is not specified, the default is 30 seconds. UnsubscribeReceiptTimeout func(unsubscribeReceiptTimeout time.Duration) func(*Conn) error + // SubscribeReceiptTimeout is a connect option that allows the client to specify + // how long to wait for a receipt in the Conn.Subscribe function, when the + // SubscribeOpt.Receipt option is used. This helps avoid deadlocks. If this + // is not specified, the default is 30 seconds. + SubscribeReceiptTimeout func(subscribeReceiptTimeout time.Duration) func(*Conn) error + // HeartBeatGracePeriodMultiplier is used to calculate the effective read heart-beat timeout // the broker will enforce for each client’s connection. The multiplier is applied to // the read-timeout interval the client specifies in its CONNECT frame @@ -272,6 +280,13 @@ func init() { } } + ConnOpt.SubscribeReceiptTimeout = func(subscribeReceiptTimeout time.Duration) func(*Conn) error { + return func(c *Conn) error { + c.options.SubscribeReceiptTimeout = subscribeReceiptTimeout + return nil + } + } + ConnOpt.UnsubscribeReceiptTimeout = func(unsubscribeReceiptTimeout time.Duration) func(*Conn) error { return func(c *Conn) error { c.options.UnsubscribeReceiptTimeout = unsubscribeReceiptTimeout diff --git a/conn_test.go b/conn_test.go index dec0bea..bd9d559 100644 --- a/conn_test.go +++ b/conn_test.go @@ -342,6 +342,189 @@ func (s *StompSuite) Test_successful_disconnect_with_receipt_timeout(c *C) { c.Assert(client.closed, Equals, true) } +func (s *StompSuite) Test_subscribe_receipt_timeout(c *C) { + resetId() + fc1, fc2 := testutil.NewFakeConn(c) + stop := make(chan struct{}) + + go func() { + defer func() { + fc2.Close() + close(stop) + }() + + reader := frame.NewReader(fc2) + writer := frame.NewWriter(fc2) + + f1, err := reader.Read() + c.Assert(err, IsNil) + c.Assert(f1.Command, Equals, "CONNECT") + err = writer.Write(frame.New("CONNECTED")) + c.Assert(err, IsNil) + + // read the SUBSCRIBE frame, but never send a RECEIPT for it + f2, err := reader.Read() + c.Assert(err, IsNil) + c.Assert(f2.Command, Equals, "SUBSCRIBE") + id, ok := f2.Header.Contains(frame.Id) + c.Assert(ok, Equals, true) + _, ok = f2.Header.Contains(frame.Receipt) + c.Assert(ok, Equals, true) + + // having given up on the receipt, the client must not leave the + // subscription behind: the caller never gets a handle to it, so it + // would otherwise be impossible to unsubscribe + f3, err := reader.Read() + c.Assert(err, IsNil) + c.Assert(f3.Command, Equals, "UNSUBSCRIBE") + c.Assert(f3.Header.Get(frame.Id), Equals, id) + err = writer.Write(frame.New(frame.RECEIPT, frame.ReceiptId, f3.Header.Get(frame.Receipt))) + c.Assert(err, IsNil) + }() + + client, err := Connect(fc1, ConnOpt.SubscribeReceiptTimeout(1*time.Millisecond)) + c.Assert(err, IsNil) + c.Assert(client, NotNil) + + sub, err := client.Subscribe("/queue/test-1", AckAuto, SubscribeOpt.Receipt("")) + c.Assert(err, Equals, ErrSubscribeReceiptTimeout) + c.Assert(sub, IsNil) + + select { + case <-stop: + case <-time.After(5 * time.Second): + c.Fatal("timed out waiting for the client to unsubscribe") + } +} + +// Regression test for abandon() wedging the whole connection: if the broker +// actually did create the subscription (it only lost or delayed the RECEIPT) +// and starts delivering messages before Subscribe() gives up, readLoop must +// not block forever trying to push them onto the now-unreachable +// Subscription.C, since that would in turn block processLoop - and with it +// every other frame on the connection, including the abandoning UNSUBSCRIBE +// itself. +func (s *StompSuite) Test_subscribe_abandon_does_not_wedge_connection(c *C) { + resetId() + fc1, fc2 := testutil.NewFakeConn(c) + stop := make(chan struct{}) + + go func() { + defer func() { + fc2.Close() + close(stop) + }() + + reader := frame.NewReader(fc2) + writer := frame.NewWriter(fc2) + + f1, err := reader.Read() + c.Assert(err, IsNil) + c.Assert(f1.Command, Equals, "CONNECT") + err = writer.Write(frame.New("CONNECTED")) + c.Assert(err, IsNil) + + // read the SUBSCRIBE frame, but never send a RECEIPT for it + f2, err := reader.Read() + c.Assert(err, IsNil) + c.Assert(f2.Command, Equals, "SUBSCRIBE") + id, ok := f2.Header.Contains(frame.Id) + c.Assert(ok, Equals, true) + + // flood more messages than Subscription.C can buffer (16), so that + // without draining, readLoop would block delivering one of them + for i := 0; i < 25; i++ { + err = writer.Write(frame.New(frame.MESSAGE, + frame.Subscription, id, + frame.Destination, "/queue/test-1", + frame.MessageId, allocateId())) + c.Assert(err, IsNil) + } + + // the abandoning UNSUBSCRIBE must still reach the wire even though + // processLoop had messages queued up for the abandoned subscription + f3, err := reader.Read() + c.Assert(err, IsNil) + c.Assert(f3.Command, Equals, "UNSUBSCRIBE") + c.Assert(f3.Header.Get(frame.Id), Equals, id) + err = writer.Write(frame.New(frame.RECEIPT, frame.ReceiptId, f3.Header.Get(frame.Receipt))) + c.Assert(err, IsNil) + + // the connection must still be usable afterwards + f4, err := reader.Read() + c.Assert(err, IsNil) + c.Assert(f4.Command, Equals, "DISCONNECT") + err = writer.Write(frame.New(frame.RECEIPT, frame.ReceiptId, f4.Header.Get(frame.Receipt))) + c.Assert(err, IsNil) + }() + + client, err := Connect(fc1, ConnOpt.SubscribeReceiptTimeout(1*time.Millisecond)) + c.Assert(err, IsNil) + c.Assert(client, NotNil) + + sub, err := client.Subscribe("/queue/test-1", AckAuto, SubscribeOpt.Receipt("")) + c.Assert(err, Equals, ErrSubscribeReceiptTimeout) + c.Assert(sub, IsNil) + + err = client.Disconnect() + c.Assert(err, IsNil) + + select { + case <-stop: + case <-time.After(5 * time.Second): + c.Fatal("connection wedged: abandoned subscription's messages blocked processLoop") + } +} + +// A reply-to (temporary queue) subscription is never sent to the server, so it +// cannot be confirmed: SubscribeOpt.Receipt must fail fast with +// ErrReceiptNotSupportedForReplyTo rather than blocking every such call until +// the receipt timeout expires. +func (s *StompSuite) Test_subscribe_reply_to_rejects_receipt(c *C) { + resetId() + fc1, fc2 := testutil.NewFakeConn(c) + stop := make(chan struct{}) + + go func() { + defer func() { + fc2.Close() + close(stop) + }() + + reader := frame.NewReader(fc2) + writer := frame.NewWriter(fc2) + + f1, err := reader.Read() + c.Assert(err, IsNil) + c.Assert(f1.Command, Equals, "CONNECT") + err = writer.Write(frame.New("CONNECTED")) + c.Assert(err, IsNil) + + // no SUBSCRIBE frame is sent for a reply-to subscription, so the next + // frame the server sees is the DISCONNECT + f2, err := reader.Read() + c.Assert(err, IsNil) + c.Assert(f2.Command, Equals, "DISCONNECT") + err = writer.Write(frame.New(frame.RECEIPT, frame.ReceiptId, f2.Header.Get(frame.Receipt))) + c.Assert(err, IsNil) + }() + + client, err := Connect(fc1, ConnOpt.SubscribeReceiptTimeout(1*time.Millisecond)) + c.Assert(err, IsNil) + c.Assert(client, NotNil) + + sub, err := client.Subscribe("/temp-queue/reply", AckAuto, + SubscribeOpt.Header(ReplyToHeader, "/temp-queue/reply"), + SubscribeOpt.Receipt("")) + c.Assert(err, Equals, ErrReceiptNotSupportedForReplyTo) + c.Assert(sub, IsNil) + + err = client.Disconnect() + c.Assert(err, IsNil) + + <-stop +} + // Sets up a connection for testing func connectHelper(c *C, version Version) (*Conn, *fakeReaderWriter) { fc1, fc2 := testutil.NewFakeConn(c) @@ -385,6 +568,58 @@ func (s *StompSuite) Test_subscribe(c *C) { } } +func (s *StompSuite) Test_subscribe_with_receipt(c *C) { + conn, rw := connectHelper(c, V12) + stop := make(chan struct{}) + + go func() { + defer func() { + rw.Close() + close(stop) + }() + + f1, err := rw.Read() + c.Assert(err, IsNil) + c.Assert(f1.Command, Equals, "SUBSCRIBE") + id, ok := f1.Header.Contains(frame.Id) + c.Assert(ok, Equals, true) + receipt, ok := f1.Header.Contains(frame.Receipt) + c.Assert(ok, Equals, true) + + // confirm the subscription + err = rw.Write(frame.New(frame.RECEIPT, frame.ReceiptId, receipt)) + c.Assert(err, IsNil) + + // the subscription must still be able to receive messages afterwards, + // i.e. the confirmation RECEIPT must not have closed it + f2 := frame.New("MESSAGE", + frame.Subscription, id, + frame.MessageId, "message-1", + frame.Destination, "/queue/test-1") + f2.Body = []byte("hello") + err = rw.Write(f2) + c.Assert(err, IsNil) + + f3, err := rw.Read() + c.Assert(err, IsNil) + c.Assert(f3.Command, Equals, "DISCONNECT") + err = rw.Write(frame.New(frame.RECEIPT, frame.ReceiptId, f3.Header.Get(frame.Receipt))) + c.Assert(err, IsNil) + }() + + sub, err := conn.Subscribe("/queue/test-1", AckAuto, SubscribeOpt.Receipt("")) + c.Assert(err, IsNil) + c.Assert(sub, NotNil) + + msg := <-sub.C + c.Assert(msg.Body, DeepEquals, []byte("hello")) + + err = conn.Disconnect() + c.Assert(err, IsNil) + + <-stop +} + func subscribeHelper(c *C, ackMode AckMode, version Version, opts ...func(*frame.Frame) error) { conn, rw := connectHelper(c, version) stop := make(chan struct{}) diff --git a/errors.go b/errors.go index 913ca90..e81552e 100644 --- a/errors.go +++ b/errors.go @@ -6,21 +6,23 @@ import ( // Error values var ( - ErrInvalidCommand = newErrorMessage("invalid command") - ErrInvalidFrameFormat = newErrorMessage("invalid frame format") - ErrUnsupportedVersion = newErrorMessage("unsupported version") - ErrCompletedTransaction = newErrorMessage("transaction is completed") - ErrNackNotSupported = newErrorMessage("NACK not supported in STOMP 1.0") - ErrNotReceivedMessage = newErrorMessage("cannot ack/nack a message, not from server") - ErrCannotNackAutoSub = newErrorMessage("cannot send NACK for a subscription with ack:auto") - ErrCompletedSubscription = newErrorMessage("subscription is unsubscribed") - ErrClosedUnexpectedly = newErrorMessage("connection closed unexpectedly") - ErrAlreadyClosed = newErrorMessage("connection already closed") - ErrMsgSendTimeout = newErrorMessage("msg send timeout") - ErrMsgReceiptTimeout = newErrorMessage("msg receipt timeout") - ErrDisconnectReceiptTimeout = newErrorMessage("disconnect receipt timeout") - ErrUnsubscribeReceiptTimeout = newErrorMessage("unsubscribe receipt timeout") - ErrNilOption = newErrorMessage("nil option") + ErrInvalidCommand = newErrorMessage("invalid command") + ErrInvalidFrameFormat = newErrorMessage("invalid frame format") + ErrUnsupportedVersion = newErrorMessage("unsupported version") + ErrCompletedTransaction = newErrorMessage("transaction is completed") + ErrNackNotSupported = newErrorMessage("NACK not supported in STOMP 1.0") + ErrNotReceivedMessage = newErrorMessage("cannot ack/nack a message, not from server") + ErrCannotNackAutoSub = newErrorMessage("cannot send NACK for a subscription with ack:auto") + ErrCompletedSubscription = newErrorMessage("subscription is unsubscribed") + ErrClosedUnexpectedly = newErrorMessage("connection closed unexpectedly") + ErrAlreadyClosed = newErrorMessage("connection already closed") + ErrMsgSendTimeout = newErrorMessage("msg send timeout") + ErrMsgReceiptTimeout = newErrorMessage("msg receipt timeout") + ErrDisconnectReceiptTimeout = newErrorMessage("disconnect receipt timeout") + ErrUnsubscribeReceiptTimeout = newErrorMessage("unsubscribe receipt timeout") + ErrSubscribeReceiptTimeout = newErrorMessage("subscribe receipt timeout") + ErrReceiptNotSupportedForReplyTo = newErrorMessage("SubscribeOpt.Receipt is not supported for reply-to subscriptions") + ErrNilOption = newErrorMessage("nil option") ) // StompError implements the Error interface, and provides diff --git a/subscribe_options.go b/subscribe_options.go index e5a5b18..0368e05 100644 --- a/subscribe_options.go +++ b/subscribe_options.go @@ -16,6 +16,22 @@ var SubscribeOpt struct { // Header provides the opportunity to include custom header entries // in the SUBSCRIBE frame that the client sends to the server. Header func(key, value string) func(*frame.Frame) error + + // Receipt makes Conn.Subscribe wait for the server to confirm the + // subscription with a RECEIPT frame before returning, avoiding a race + // where a message published right after Subscribe() returns is lost + // because the server hasn't finished creating the subscription yet. + // + // receiptId is the value of the "receipt" header sent to the server. If + // left empty, a unique value is generated. + // + // If confirmation doesn't arrive within ConnOpt.SubscribeReceiptTimeout, + // Subscribe unsubscribes again and returns ErrSubscribeReceiptTimeout. + // + // Reply-to (temporary queue) subscriptions are never sent to the server + // and so cannot be confirmed: using Receipt with one makes Subscribe + // return ErrReceiptNotSupportedForReplyTo. + Receipt func(receiptId string) func(*frame.Frame) error } func init() { @@ -39,4 +55,17 @@ func init() { return nil } } + + SubscribeOpt.Receipt = func(receiptId string) func(*frame.Frame) error { + return func(f *frame.Frame) error { + if f.Command != frame.SUBSCRIBE { + return ErrInvalidCommand + } + if receiptId == "" { + receiptId = allocateId() + } + f.Header.Set(frame.Receipt, receiptId) + return nil + } + } } diff --git a/subscription.go b/subscription.go index 8ca54af..8f45ce4 100644 --- a/subscription.go +++ b/subscription.go @@ -115,6 +115,41 @@ func (s *Subscription) Unsubscribe(opts ...func(*frame.Frame) error) error { return err } +// abandon tears down a subscription that Conn.Subscribe created but never +// handed to the caller, because it failed while waiting for the confirming +// RECEIPT. Unlike Unsubscribe, it doesn't wait for the RECEIPT itself: there +// is nobody left to report the outcome to. +// +// If the broker never confirms the resulting UNSUBSCRIBE either, the drain +// goroutine started below (and C itself) is only reclaimed when the +// connection closes. +func (s *Subscription) abandon() { + // transition to the "closing" state + if !atomic.CompareAndSwapInt32(&s.state, subStateActive, subStateClosing) { + return + } + + // Nobody holds this subscription, so nobody will ever read C. Drain it + // until readLoop closes it, or messages already in flight would fill C, + // block readLoop on the frame channel and wedge processLoop with it. + go func() { + for range s.C { + } + }() + + f := frame.New(frame.UNSUBSCRIBE, frame.Id, s.id) + if s.replyToSet { + f.Header.Set(ReplyToHeader, s.id) + } + + if err := s.conn.sendFrame(f); err != nil { + // Nothing more we can do; the frame channel (and with it C) is closed + // when the connection is torn down. + s.conn.log.Infof("could not unsubscribe unconfirmed subscription %s: %s: %v", + s.id, s.destination, err) + } +} + func waitWithTimeout(cond *sync.Cond, timeout time.Duration) error { if timeout == 0 { cond.Wait()