PR feedback: reject pipeline pushes once stopped

A push racing the connection-close drain could strand its JSA reply
registration in the account-scoped map. push() now rejects entries on a
stopped pipeline and re-drains after a raced enqueue.

Signed-off-by: Lev Brouk <levbrouk@gmail.com>

Author: levb

server/mqtt.go

@@ -246,6 +246,9 @@ var (
 	errMQTTInvalidRetainedMessage     = errors.New("invalid retained message")
 	errMQTTSessionCollision           = errors.New("stored session does not match client ID")
 	errMQTTInvalidPublishLength       = errors.New("invalid publish message, variable header exceeds remaining length")
+	errMQTTAckPipelineStopped         = errors.New("QoS1 PUBACK pipeline has shut down while admitting a message, " +
+		"abandoning the wait for its JetStream ack; failing the connection, " +
+		"the client will re-send unacknowledged PUBLISH packets on reconnect")
 )
 
 type srvMQTT struct {
@@ -4693,33 +4696,58 @@ func (s *Server) mqttPipelinePush(c *client, jsa *mqttJSA, ack *mqttPipelinedAck
 		c.mqtt.acks = pipe
 	}
 
-	// A just-stopped pipeline may admit an unconsumed entry; the
-	// connection-close handler (mqttHandleClosedClient) drains the queue
-	// via pipe.shutdown().
+	return pipe.push(ack)
+}
+
+// Admits an entry into the pipeline. Rejects it once the pipeline is
+// stopped: an entry admitted after the connection-close drain
+// (mqttHandleClosedClient -> pipe.shutdown) would strand its reply
+// registration in the account-scoped jsa.replies. readLoop only.
+func (pipe *mqttAckPipeline) push(ack *mqttPipelinedAck) error {
+	jsa := pipe.jsa
 	select {
-	case pipe.q <- ack:
-		return nil
+	case <-pipe.quitCh:
+		jsa.replies.Delete(ack.reply)
+		return errMQTTAckPipelineStopped
 	default:
 	}
 
-	// Window full: wait for an available slot in the pipeline.
-	t := time.NewTimer(jsa.timeout)
-	defer t.Stop()
+	enqueued := false
 	select {
 	case pipe.q <- ack:
-		return nil
-	case <-pipe.quitCh:
-		jsa.replies.Delete(ack.reply)
-		return errors.New("QoS1 PUBACK pipeline has shut down while admitting a message, " +
-			"abandoning the wait for its JetStream ack; failing the connection, " +
-			"the client will re-send unacknowledged PUBLISH packets on reconnect")
-	case <-t.C:
+		enqueued = true
+	default:
+		// Window full: wait for an available slot in the pipeline.
+		t := time.NewTimer(jsa.timeout)
+		defer t.Stop()
+		select {
+		case pipe.q <- ack:
+			enqueued = true
+		case <-pipe.quitCh:
+		case <-t.C:
+			jsa.replies.Delete(ack.reply)
+			return fmt.Errorf("QoS1 in-flight window is full (%d messages) and JetStream has not acknowledged "+
+				"the oldest message within %v; failing the connection, "+
+				"the client will re-send unacknowledged PUBLISH packets on reconnect",
+				mqttMaxAcksInFlight, jsa.timeout)
+		}
+	}
+	if !enqueued {
+		// Stopped while waiting for a slot, never admitted.
 		jsa.replies.Delete(ack.reply)
-		return fmt.Errorf("QoS1 in-flight window is full (%d messages) and JetStream has not acknowledged "+
-			"the oldest message within %v; failing the connection, "+
-			"the client will re-send unacknowledged PUBLISH packets on reconnect",
-			mqttMaxAcksInFlight, jsa.timeout)
+		return errMQTTAckPipelineStopped
+	}
+
+	// A stop may have raced the enqueue: the close-time drain could run
+	// before the entry landed and miss it. Drain again; both drains only
+	// dequeue and delete, so overlapping is harmless.
+	select {
+	case <-pipe.quitCh:
+		pipe.shutdown()
+		return errMQTTAckPipelineStopped
+	default:
 	}
+	return nil
 }
 
 func (s *Server) mqttInitiateMsgDelivery(c *client, pp *mqttPublish) error {

server/mqtt_test.go

@@ -10402,11 +10402,12 @@ func TestMQTTQoS1PubAckPipelineConnClose(t *testing.T) {
 	testMQTTCheckPubMsgNoAck(t, mcs, msr, "foo", mqttPubQos1, []byte("msg2"))
 }
 
-// A pipeline stopped while the readLoop is still admitting entries must
-// not leak JSA reply registrations: shutdown() runs after the last push.
+// A pipeline shut down while the readLoop is still admitting entries must
+// not leak JSA reply registrations, even when the close-time drain runs
+// concurrently with (or before) a racing push.
 func TestMQTTQoS1PubAckPipelineShutdownRace(t *testing.T) {
 	for i := 0; i < 100; i++ {
-		jsa := &mqttJSA{}
+		jsa := &mqttJSA{timeout: time.Second}
 		pipe := &mqttAckPipeline{
 			jsa:    jsa,
 			q:      make(chan *mqttPipelinedAck, 4),
@@ -10415,28 +10416,25 @@ func TestMQTTQoS1PubAckPipelineShutdownRace(t *testing.T) {
 
 		done := make(chan struct{})
 		go func() {
-			// The readLoop side: register, admit, clean up on stop.
+			// The readLoop side: register, admit until rejected.
 			defer close(done)
-			defer pipe.shutdown()
 			for n := 0; ; n++ {
 				ack := &mqttPipelinedAck{pi: uint16(n%0xFFFF + 1), reply: fmt.Sprintf("reply.%d", n), done: make(chan error, 1)}
 				jsa.replies.Store(ack.reply, func(any) {})
-				select {
-				case pipe.q <- ack:
-				case <-pipe.quitCh:
-					// As mqttPipelinePush does when stopped.
-					jsa.replies.Delete(ack.reply)
+				if err := pipe.push(ack); err != nil {
 					return
 				}
 			}
 		}()
 
-		// The consumer side: take a few entries, then stop mid-stream.
+		// The consumer side: take a few entries, then stop and drain
+		// concurrently with the pushes, as the connection-close handler
+		// does.
 		for j := 0; j < i%4; j++ {
 			ack := <-pipe.q
 			jsa.replies.Delete(ack.reply)
 		}
-		pipe.stop()
+		pipe.shutdown()
 		<-done
 
 		leaked := 0
@@ -10445,4 +10443,25 @@ func TestMQTTQoS1PubAckPipelineShutdownRace(t *testing.T) {
 			t.Fatalf("iteration %d: %v reply registration(s) leaked after shutdown", i, leaked)
 		}
 	}
+
+	// A push after the pipeline has stopped must be rejected outright,
+	// with its registration cleaned up and nothing left in the queue.
+	jsa := &mqttJSA{timeout: time.Second}
+	pipe := &mqttAckPipeline{
+		jsa:    jsa,
+		q:      make(chan *mqttPipelinedAck, 4),
+		quitCh: make(chan struct{}),
+	}
+	pipe.shutdown()
+	ack := &mqttPipelinedAck{pi: 1, reply: "reply.stopped", done: make(chan error, 1)}
+	jsa.replies.Store(ack.reply, func(any) {})
+	if err := pipe.push(ack); err != errMQTTAckPipelineStopped {
+		t.Fatalf("Expected errMQTTAckPipelineStopped, got %v", err)
+	}
+	if _, ok := jsa.replies.Load(ack.reply); ok {
+		t.Fatal("reply registration not cleaned up on rejected push")
+	}
+	if n := len(pipe.q); n != 0 {
+		t.Fatalf("%v entry(ies) admitted into a stopped pipeline", n)
+	}
 }