feat: add SubscribeOpt.Receipt

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
 }
 

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