Repository: nats-io/nats.go
Pull request #1691 (2024-08-02 08:50:17)

[ADDED] ConsumeContext.Closed() method for waiting for consume to be closed/drained
Signed-off-by: Piotr Piotrowski <piotr@synadia.com>

AlexCuse: Thank you so much for this

Patch:
diff --git a/README.md b/README.md
index 64307b2..75e19fd 100644
--- a/README.md
+++ b/README.md
@@ -381,6 +381,8 @@ msg, err := nc.RequestWithContext(ctx, "foo", []byte("bar"))
 sub, err := nc.SubscribeSync("foo")
 msg, err := sub.NextMsgWithContext(ctx)
 
+```
+
 ## Backwards compatibility
 
 In the development of nats.go, we are committed to maintaining backward compatibility and ensuring a stable and reliable  experience for all users. In general, we follow the standard go compatibility guidelines.
diff --git a/jetstream/errors.go b/jetstream/errors.go
index fb36434..8d2fec6 100644
--- a/jetstream/errors.go
+++ b/jetstream/errors.go
@@ -270,6 +270,9 @@ var (
 	// of an ordered consumer which was not yet created.
 	ErrOrderedConsumerNotCreated JetStreamError = &jsError{message: "consumer instance not yet created"}
 
+	// ErrJetStreamPublisherClosed is returned for each unfinished ack future when JetStream.Cleanup is called.
+	ErrJetStreamPublisherClosed JetStreamError = &jsError{message: "jetstream context closed"}
+
 	// KeyValue Errors
 
 	// ErrKeyExists is returned when attempting to create a key that already
diff --git a/jetstream/jetstream.go b/jetstream/jetstream.go
index 554400b..e401cb9 100644
--- a/jetstream/jetstream.go
+++ b/jetstream/jetstream.go
@@ -101,6 +101,19 @@ type (
 		// outstanding asynchronously published messages are acknowledged by the
 		// server.
 		PublishAsyncComplete() <-chan struct{}
+
+		// CleanupPublisher will cleanup the publishing side of JetStreamContext.
+		//
+		// This will unsubscribe from the internal reply subject if needed.
+		// All pending async publishes will fail with ErrJetStreamContextClosed.
+		//
+		// If an error handler was provided, it will be called for each pending async
+		// publish and PublishAsyncComplete will be closed.
+		//
+		// After completing JetStreamContext is still usable - internal subscription
+		// will be recreated on next publish, but the acks from previous publishes will
+		// be lost.
+		CleanupPublisher()
 	}
 
 	// StreamManager provides CRUD API for managing streams. It is available as
@@ -1032,6 +1045,39 @@ func wrapContextWithoutDeadline(ctx context.Context) (context.Context, context.C
 	return context.WithTimeout(ctx, defaultAPITimeout)
 }
 
+// CleanupPublisher will cleanup the publishing side of JetStreamContext.
+//
+// This will unsubscribe from the internal reply subject if needed.
+// All pending async publishes will fail with ErrJetStreamContextClosed.
+//
+// If an error handler was provided, it will be called for each pending async
+// publish and PublishAsyncComplete will be closed.
+//
+// After completing JetStreamContext is still usable - internal subscription
+// will be recreated on next publish, but the acks from previous publishes will
+// be lost.
+func (js *jetStream) CleanupPublisher() {
+	js.cleanupReplySub()
+	js.publisher.Lock()
+	errCb := js.publisher.aecb
+	for id, paf := range js.publisher.acks {
+		paf.err = ErrJetStreamPublisherClosed
+		if paf.errCh != nil {
+			paf.errCh <- paf.err
+		}
+		if errCb != nil {
+			// call error handler after releasing the mutex to avoid contention
+			defer errCb(js, paf.msg, ErrJetStreamPublisherClosed)
+		}
+		delete(js.publisher.acks, id)
+	}
+	if js.publisher.doneCh != nil {
+		close(js.publisher.doneCh)
+		js.publisher.doneCh = nil
+	}
+	js.publisher.Unlock()
+}
+
 func (js *jetStream) cleanupReplySub() {
 	if js.publisher == nil {
 		return
diff --git a/jetstream/ordered.go b/jetstream/ordered.go
index 85b7ea9..199b8c3 100644
--- a/jetstream/ordered.go
+++ b/jetstream/ordered.go
@@ -32,13 +32,13 @@ type (
 		cfg               *OrderedConsumerConfig
 		stream            string
 		currentConsumer   *pullConsumer
-		currentSub        ConsumeContext
+		currentSub        *pullSubscription
 		cursor            cursor
 		namePrefix        string
 		serial            int
 		consumerType      consumerType
 		doReset           chan struct{}
-		resetInProgress   uint32
+		resetInProgress   atomic.Uint32
 		userErrHandler    ConsumeErrHandlerFunc
 		stopAfter         int
 		stopAfterMsgsLeft chan int
@@ -52,7 +52,7 @@ type (
 		consumer *orderedConsumer
 		opts     []PullMessagesOpt
 		done     chan struct{}
-		closed   uint32
+		closed   atomic.Uint32
 	}
 
 	cursor struct {
@@ -138,7 +138,7 @@ func (c *orderedConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt
 	if err != nil {
 		return nil, err
 	}
-	c.currentSub = cc
+	c.currentSub = cc.(*pullSubscription)
 
 	go func() {
 		for {
@@ -175,7 +175,7 @@ func (c *orderedConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt
 					c.errHandler(c.serial)(cc, err)
 				} else {
 					c.Lock()
-					c.currentSub = cc
+					c.currentSub = cc.(*pullSubscription)
 					c.Unlock()
 				}
 			case <-sub.done:
@@ -210,8 +210,8 @@ func (c *orderedConsumer) errHandler(serial int) func(cc ConsumeContext, err err
 			errors.Is(err, ErrConsumerDeleted) ||
 			errors.Is(err, errConnected) {
 			// only reset if serial matches the current consumer serial and there is no reset in progress
-			if serial == c.serial && atomic.LoadUint32(&c.resetInProgress) == 0 {
-				atomic.StoreUint32(&c.resetInProgress, 1)
+			if serial == c.serial && c.resetInProgress.Load() == 0 {
+				c.resetInProgress.Store(1)
 				c.doReset <- struct{}{}
 			}
 		}
@@ -256,7 +256,7 @@ func (c *orderedConsumer) Messages(opts ...PullMessagesOpt) (MessagesContext, er
 	if err != nil {
 		return nil, err
 	}
-	c.currentSub = cc
+	c.currentSub = cc.(*pullSubscription)
 
 	sub := &orderedSubscription{
 		consumer: c,
@@ -270,7 +270,7 @@ func (c *orderedConsumer) Messages(opts ...PullMessagesOpt) (MessagesContext, er
 
 func (s *orderedSubscription) Next() (Msg, error) {
 	for {
-		msg, err := s.consumer.currentSub.(*pullSubscription).Next()
+		msg, err := s.consumer.currentSub.Next()
 		if err != nil {
 			if errors.Is(err, ErrMsgIteratorClosed) {
 				s.Stop()
@@ -297,7 +297,7 @@ func (s *orderedSubscription) Next() (Msg, error) {
 			if err != nil {
 				return nil, err
 			}
-			s.consumer.currentSub = cc
+			s.consumer.currentSub = cc.(*pullSubscription)
 			continue
 		}
 
@@ -321,7 +321,7 @@ func (s *orderedSubscription) Next() (Msg, error) {
 			if err != nil {
 				return nil, err
 			}
-			s.consumer.currentSub = cc
+			s.consumer.currentSub = cc.(*pullSubscription)
 			continue
 		}
 		s.consumer.cursor.deliverSeq = dseq
@@ -331,7 +331,7 @@ func (s *orderedSubscription) Next() (Msg, error) {
 }
 
 func (s *orderedSubscription) Stop() {
-	if !atomic.CompareAndSwapUint32(&s.closed, 0, 1) {
+	if !s.closed.CompareAndSwap(0, 1) {
 		return
 	}
 	s.consumer.Lock()
@@ -343,7 +343,7 @@ func (s *orderedSubscription) Stop() {
 }
 
 func (s *orderedSubscription) Drain() {
-	if !atomic.CompareAndSwapUint32(&s.closed, 0, 1) {
+	if !s.closed.CompareAndSwap(0, 1) {
 		return
 	}
 	if s.consumer.currentSub != nil {
@@ -354,6 +354,37 @@ func (s *orderedSubscription) Drain() {
 	close(s.done)
 }
 
+// Closed returns a channel that is closed when the consuming is
+// fully stopped/drained. When the channel is closed, no more messages
+// will be received and processing is complete.
+func (s *orderedSubscription) Closed() <-chan struct{} {
+	s.consumer.Lock()
+	defer s.consumer.Unlock()
+	closedCh := make(chan struct{})
+
+	go func() {
+		for {
+			s.consumer.Lock()
+			if s.consumer.currentSub == nil {
+				return
+			}
+
+			closed := s.consumer.currentSub.Closed()
+			s.consumer.Unlock()
+
+			// wait until the underlying pull consumer is closed
+			<-closed
+			// if the subscription is closed and ordered consumer is closed as well,
+			// send a signal that the Consume() is fully stopped
+			if s.closed.Load() == 1 {
+				close(closedCh)
+				return
+			}
+		}
+	}()
+	return closedCh
+}
+
 // Fetch is used to retrieve up to a provided number of messages from a
 // stream. This method will always send a single request and wait until
 // either all messages are retrieved or request times out.
@@ -495,7 +526,7 @@ func serialNumberFromConsumer(name string) int {
 func (c *orderedConsumer) reset() error {
 	c.Lock()
 	defer c.Unlock()
-	defer atomic.StoreUint32(&c.resetInProgress, 0)
+	defer c.resetInProgress.Store(0)
 	if c.currentConsumer != nil {
 		c.currentConsumer.Lock()
 		if c.currentSub != nil {
@@ -524,7 +555,7 @@ func (c *orderedConsumer) reset() error {
 		cancel:          c.subscription.done,
 	}
 	err = retryWithBackoff(func(attempt int) (bool, error) {
-		isClosed := atomic.LoadUint32(&c.subscription.closed) == 1
+		isClosed := c.subscription.closed.Load() == 1
 		if isClosed {
 			return false, errOrderedConsumerClosed
 		}
@@ -583,7 +614,8 @@ func (c *orderedConsumer) getConsumerConfig() *ConsumerConfig {
 		cfg.InactiveThreshold = c.cfg.InactiveThreshold
 	}
 
-	if c.serial != 1 {
+	// if the cursor is not yet set, use the provided deliver policy
+	if c.cursor.streamSeq != 0 {
 		return cfg
 	}
 
diff --git a/jetstream/publish.go b/jetstream/publish.go
index 1a3a4fd..d08cdf1 100644
--- a/jetstream/publish.go
+++ b/jetstream/publish.go
@@ -1,4 +1,4 @@
-// Copyright 2022-2023 The NATS Authors
+// Copyright 2022-2024 The NATS Authors
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
@@ -467,11 +467,7 @@ func (js *jetStream) resetPendingAcksOnReconnect() {
 				paf.errCh <- paf.err
 			}
 			if errCb != nil {
-				js.publisher.Unlock()
-				// clear reply subject so that new one is created on republish
-				paf.msg.Reply = ""
-				errCb(js, paf.msg, nats.ErrDisconnected)
-				js.publisher.Lock()
+				defer errCb(js, paf.msg, nats.ErrDisconnected)
 			}
 			delete(js.publisher.acks, id)
 		}
diff --git a/jetstream/pull.go b/jetstream/pull.go
index a510c7c..5402828 100644
--- a/jetstream/pull.go
+++ b/jetstream/pull.go
@@ -59,6 +59,11 @@ type (
 		// Drain unsubscribes from the stream and cancels subscription.
 		// All messages that are already in the buffer will be processed in callback function.
 		Drain()
+
+		// Closed returns a channel that is closed when the consuming is
+		// fully stopped/drained. When the channel is closed, no more messages
+		// will be received and processing is complete.
+		Closed() <-chan struct{}
 	}
 
 	// MessageHandler is a handler function used as callback in [Consume].
@@ -125,6 +130,7 @@ type (
 		fetchNext         chan *pullRequest
 		consumeOpts       *consumeOpts
 		delivered         int
+		closedCh          chan struct{}
 	}
 
 	pendingMsgs struct {
@@ -257,6 +263,12 @@ func (p *pullConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt) (
 		return func(subject string) {
 			p.subs.Delete(sid)
 			sub.draining.CompareAndSwap(1, 0)
+			sub.Lock()
+			if sub.closedCh != nil {
+				close(sub.closedCh)
+				sub.closedCh = nil
+			}
+			sub.Unlock()
 		}
 	}(sub.id))
 
@@ -649,6 +661,24 @@ func (s *pullSubscription) Drain() {
 	}
 }
 
+// Closed returns a channel that is closed when consuming is
+// fully stopped/drained. When the channel is closed, no more messages
+// will be received and processing is complete.
+func (s *pullSubscription) Closed() <-chan struct{} {
+	s.Lock()
+	defer s.Unlock()
+	closedCh := s.closedCh
+	if closedCh == nil {
+		closedCh = make(chan struct{})
+		s.closedCh = closedCh
+	}
+	if !s.subscription.IsValid() {
+		close(s.closedCh)
+		s.closedCh = nil
+	}
+	return closedCh
+}
+
 // Fetch sends a single request to retrieve given number of messages.
 // It will wait up to provided expiry time if not all messages are available.
 func (p *pullConsumer) Fetch(batch int, opts ...FetchOpt) (MessageBatch, error) {
diff --git a/js.go b/js.go
index 5f8dfe3..c8b6268 100644
--- a/js.go
+++ b/js.go
@@ -1,4 +1,4 @@
-// Copyright 2020-2023 The NATS Authors
+// Copyright 2020-2024 The NATS Authors
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
@@ -58,6 +58,19 @@ type JetStream interface {
 	// PublishAsyncComplete returns a channel that will be closed when all outstanding messages are ack'd.
 	PublishAsyncComplete() <-chan struct{}
 
+	// CleanupPublisher will cleanup the publishing side of JetStreamContext.
+	//
+	// This will unsubscribe from the internal reply subject if needed.
+	// All pending async publishes will fail with ErrJetStreamPublisherClosed.
+	//
+	// If an error handler was provided, it will be called for each pending async
+	// publish and PublishAsyncComplete will be closed.
+	//
+	// After completing JetStreamContext is still usable - internal subscription
+	// will be recreated on next publish, but the acks from previous publishes will
+	// be lost.
+	CleanupPublisher()
+
 	// Subscribe creates an async Subscription for JetStream.
 	// The stream and consumer names can be provided with the nats.Bind() option.
 	// For creating an ephemeral (where the consumer name is picked by the server),
@@ -719,10 +732,7 @@ func (js *js) resetPendingAcksOnReconnect() {
 				paf.errCh <- paf.err
 			}
 			if errCb != nil {
-				// clear reply subject so that new one is created on republish
-				js.mu.Unlock()
-				errCb(js, paf.msg, ErrDisconnected)
-				js.mu.Lock()
+				defer errCb(js, paf.msg, ErrDisconnected)
 			}
 			delete(js.pafs, id)
 		}
@@ -734,6 +744,38 @@ func (js *js) resetPendingAcksOnReconnect() {
 	}
 }
 
+// CleanupPublisher will cleanup the publishing side of JetStreamContext.
+//
+// This will unsubscribe from the internal reply subject if needed.
+// All pending async publishes will fail with ErrJetStreamContextClosed.
+//
+// If an error handler was provided, it will be called for each pending async
+// publish and PublishAsyncComplete will be closed.
+//
+// After completing JetStreamContext is still usable - internal subscription
+// will be recreated on next publish, but the acks from previous publishes will
+// be lost.
+func (js *js) CleanupPublisher() {
+	js.cleanupReplySub()
+	js.mu.Lock()
+	errCb := js.opts.aecb
+	for id, paf := range js.pafs {
+		paf.err = ErrJetStreamPublisherClosed
+		if paf.errCh != nil {
+			paf.errCh <- paf.err
+		}
+		if errCb != nil {
+			defer errCb(js, paf.msg, ErrJetStreamPublisherClosed)
+		}
+		delete(js.pafs, id)
+	}
+	if js.dch != nil {
+		close(js.dch)
+		js.dch = nil
+	}
+	js.mu.Unlock()
+}
+
 func (js *js) cleanupReplySub() {
 	js.mu.Lock()
 	if js.rsub != nil {
@@ -2899,10 +2941,11 @@ func (sub *Subscription) Fetch(batch int, opts ...PullOpt) ([]*Msg, error) {
 			}
 
 			// Make our request expiration a bit shorter than the current timeout.
-			expires := ttl
-			if ttl >= 20*time.Millisecond {
-				expires = ttl - 10*time.Millisecond
+			expiresDiff := time.Duration(float64(ttl) * 0.1)
+			if expiresDiff > 5*time.Second {
+				expiresDiff = 5 * time.Second
 			}
+			expires := ttl - expiresDiff
 
 			nr.Batch = batch - len(msgs)
 			nr.Expires = expires
@@ -3166,10 +3209,11 @@ func (sub *Subscription) FetchBatch(batch int, opts ...PullOpt) (MessageBatch, e
 	ttl = time.Until(deadline)
 
 	// Make our request expiration a bit shorter than the current timeout.
-	expires := ttl
-	if ttl >= 20*time.Millisecond {
-		expires = ttl - 10*time.Millisecond
+	expiresDiff := time.Duration(float64(ttl) * 0.1)
+	if expiresDiff > 5*time.Second {
+		expiresDiff = 5 * time.Second
 	}
+	expires := ttl - expiresDiff
 
 	requestBatch := batch - len(result.msgs)
 	req := nextRequest{
diff --git a/jserrors.go b/jserrors.go
index 2d942e7..1c22d81 100644
--- a/jserrors.go
+++ b/jserrors.go
@@ -151,6 +151,9 @@ var (
 	// ErrSubscriptionClosed is returned when attempting to send pull request to a closed subscription
 	ErrSubscriptionClosed JetStreamError = &jsError{message: "subscription closed"}
 
+	// ErrJetStreamPublisherClosed is returned for each unfinished ack future when JetStream.Cleanup is called.
+	ErrJetStreamPublisherClosed JetStreamError = &jsError{message: "jetstream context closed"}
+
 	// Deprecated: ErrInvalidDurableName is no longer returned and will be removed in future releases.
 	// Use ErrInvalidConsumerName instead.
 	ErrInvalidDurableName = errors.New("nats: invalid durable name")


Test patch:
diff --git a/jetstream/test/jetstream_test.go b/jetstream/test/jetstream_test.go
index f5c9c8e..aa209aa 100644
--- a/jetstream/test/jetstream_test.go
+++ b/jetstream/test/jetstream_test.go
@@ -1,4 +1,4 @@
-// Copyright 2022-2023 The NATS Authors
+// Copyright 2022-2024 The NATS Authors
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
@@ -1963,3 +1963,125 @@ func TestConsumerConfigMatches(t *testing.T) {
 		t.Fatalf("ConsumerConfig doesn't match")
 	}
 }
+
+func TestJetStreamCleanupPublisher(t *testing.T) {
+
+	t.Run("cleanup js publisher", func(t *testing.T) {
+		s := RunBasicJetStreamServer()
+		defer shutdownJSServerAndRemoveStorage(t, s)
+
+		nc, js := jsClient(t, s)
+		defer nc.Close()
+
+		// Create a stream
+		if _, err := js.CreateStream(context.Background(), jetstream.StreamConfig{Name: "TEST", Subjects: []string{"FOO"}}); err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		numSubs := nc.NumSubscriptions()
+		if _, err := js.PublishAsync("FOO", []byte("hello")); err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		select {
+		case <-js.PublishAsyncComplete():
+		case <-time.After(5 * time.Second):
+			t.Fatalf("Did not receive completion signal")
+		}
+
+		if numSubs+1 != nc.NumSubscriptions() {
+			t.Fatalf("Expected an additional subscription after publish, got %d", nc.NumSubscriptions())
+		}
+
+		js.CleanupPublisher()
+
+		if numSubs != nc.NumSubscriptions() {
+			t.Fatalf("Expected subscriptions to be back to original count")
+		}
+	})
+
+	t.Run("cleanup js publisher, cancel pending acks", func(t *testing.T) {
+		s := RunBasicJetStreamServer()
+		defer shutdownJSServerAndRemoveStorage(t, s)
+
+		nc, err := nats.Connect(s.ClientURL())
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		defer nc.Close()
+
+		cbErr := make(chan error, 10)
+		js, err := jetstream.New(nc, jetstream.WithPublishAsyncErrHandler(func(js jetstream.JetStream, m *nats.Msg, err error) {
+			cbErr <- err
+		}))
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		// Create a stream with NoAck so that we can test that we cancel ack futures.
+		if _, err := js.CreateStream(context.Background(), jetstream.StreamConfig{Name: "TEST", Subjects: []string{"FOO"}, NoAck: true}); err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		numSubs := nc.NumSubscriptions()
+
+		var acks []jetstream.PubAckFuture
+		for i := 0; i < 10; i++ {
+			ack, err := js.PublishAsync("FOO", []byte("hello"))
+			if err != nil {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+			acks = append(acks, ack)
+		}
+
+		asyncComplete := js.PublishAsyncComplete()
+		select {
+		case <-asyncComplete:
+			t.Fatalf("Should not complete, NoAck is set")
+		case <-time.After(200 * time.Millisecond):
+		}
+
+		if numSubs+1 != nc.NumSubscriptions() {
+			t.Fatalf("Expected an additional subscription after publish, got %d", nc.NumSubscriptions())
+		}
+
+		js.CleanupPublisher()
+
+		if numSubs != nc.NumSubscriptions() {
+			t.Fatalf("Expected subscriptions to be back to original count")
+		}
+
+		// check that PublishAsyncComplete channel is closed
+		select {
+		case <-asyncComplete:
+		case <-time.After(5 * time.Second):
+			t.Fatalf("Did not receive completion signal")
+		}
+
+		// check that all ack futures are cancelled
+		for _, ack := range acks {
+			select {
+			case err := <-ack.Err():
+				if !errors.Is(err, jetstream.ErrJetStreamPublisherClosed) {
+					t.Fatalf("Expected JetStreamContextClosed error, got %v", err)
+				}
+			case <-ack.Ok():
+				t.Fatalf("Expected error on the ack future")
+			case <-time.After(200 * time.Millisecond):
+				t.Fatalf("Expected an error on the ack future")
+			}
+		}
+
+		// check that async error handler is called for each pending ack
+		for i := 0; i < 10; i++ {
+			select {
+			case err := <-cbErr:
+				if !errors.Is(err, jetstream.ErrJetStreamPublisherClosed) {
+					t.Fatalf("Expected JetStreamContextClosed error, got %v", err)
+				}
+			case <-time.After(200 * time.Millisecond):
+				t.Fatalf("Expected errors to be passed from the async handler")
+			}
+		}
+	})
+
+}
diff --git a/jetstream/test/ordered_test.go b/jetstream/test/ordered_test.go
index 522c921..6680955 100644
--- a/jetstream/test/ordered_test.go
+++ b/jetstream/test/ordered_test.go
@@ -578,6 +578,135 @@ func TestOrderedConsumerConsume(t *testing.T) {
 			time.Sleep(50 * time.Millisecond)
 		}
 	})
+
+	t.Run("wait for closed after drain", func(t *testing.T) {
+		for i := 0; i < 10; i++ {
+			t.Run(fmt.Sprintf("run %d", i), func(t *testing.T) {
+				srv := RunBasicJetStreamServer()
+				defer shutdownJSServerAndRemoveStorage(t, srv)
+				nc, err := nats.Connect(srv.ClientURL())
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+
+				js, err := jetstream.New(nc)
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				defer nc.Close()
+				ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+				defer cancel()
+				s, err := js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}})
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				c, err := s.OrderedConsumer(ctx, jetstream.OrderedConsumerConfig{})
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				msgs := make([]jetstream.Msg, 0)
+				lock := sync.Mutex{}
+				publishTestMsgs(t, js)
+				cc, err := c.Consume(func(msg jetstream.Msg) {
+					time.Sleep(50 * time.Millisecond)
+					msg.Ack()
+					lock.Lock()
+					msgs = append(msgs, msg)
+					lock.Unlock()
+				})
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				closed := cc.Closed()
+				time.Sleep(100 * time.Millisecond)
+				if err := s.DeleteConsumer(context.Background(), c.CachedInfo().Name); err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				publishTestMsgs(t, js)
+
+				// wait for the consumer to be recreated before calling drain
+				for i := 0; i < 5; i++ {
+					_, err = c.Info(ctx)
+					if err != nil {
+						if errors.Is(err, jetstream.ErrConsumerNotFound) {
+							time.Sleep(100 * time.Millisecond)
+							continue
+						}
+						t.Fatalf("Unexpected error: %v", err)
+					}
+					break
+				}
+
+				cc.Drain()
+
+				select {
+				case <-closed:
+				case <-time.After(5 * time.Second):
+					t.Fatalf("Timeout waiting for consume to be closed")
+				}
+
+				if len(msgs) != 2*len(testMsgs) {
+					t.Fatalf("Unexpected received message count after consume closed; want %d; got %d", 2*len(testMsgs), len(msgs))
+				}
+			})
+		}
+	})
+
+	t.Run("wait for closed on already closed consume", func(t *testing.T) {
+		for i := 0; i < 10; i++ {
+			t.Run(fmt.Sprintf("run %d", i), func(t *testing.T) {
+				srv := RunBasicJetStreamServer()
+				defer shutdownJSServerAndRemoveStorage(t, srv)
+				nc, err := nats.Connect(srv.ClientURL())
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+
+				js, err := jetstream.New(nc)
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				defer nc.Close()
+				ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+				defer cancel()
+				s, err := js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}})
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				c, err := s.OrderedConsumer(ctx, jetstream.OrderedConsumerConfig{})
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				msgs := make([]jetstream.Msg, 0)
+				lock := sync.Mutex{}
+				publishTestMsgs(t, js)
+				cc, err := c.Consume(func(msg jetstream.Msg) {
+					time.Sleep(50 * time.Millisecond)
+					msg.Ack()
+					lock.Lock()
+					msgs = append(msgs, msg)
+					lock.Unlock()
+				})
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				time.Sleep(100 * time.Millisecond)
+				if err := s.DeleteConsumer(context.Background(), c.CachedInfo().Name); err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+
+				cc.Stop()
+
+				time.Sleep(100 * time.Millisecond)
+
+				select {
+				case <-cc.Closed():
+				case <-time.After(5 * time.Second):
+					t.Fatalf("Timeout waiting for consume to be closed")
+				}
+			})
+		}
+	})
 }
 
 func TestOrderedConsumerMessages(t *testing.T) {
@@ -1119,6 +1248,71 @@ func TestOrderedConsumerFetch(t *testing.T) {
 		}
 	})
 
+	t.Run("with custom deliver policy", func(t *testing.T) {
+		srv := RunBasicJetStreamServer()
+		defer shutdownJSServerAndRemoveStorage(t, srv)
+		nc, err := nats.Connect(srv.ClientURL())
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		js, err := jetstream.New(nc)
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		defer nc.Close()
+
+		ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
+		defer cancel()
+		s, err := js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}})
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		msgs := make([]jetstream.Msg, 0)
+
+		for i := 0; i < 5; i++ {
+			if _, err := js.Publish(context.Background(), "FOO.A", []byte("msg")); err != nil {
+				t.Fatalf("Unexpected error during publish: %s", err)
+			}
+		}
+		for i := 0; i < 5; i++ {
+			if _, err := js.Publish(context.Background(), "FOO.B", []byte("msg")); err != nil {
+				t.Fatalf("Unexpected error during publish: %s", err)
+			}
+		}
+
+		c, err := s.OrderedConsumer(ctx, jetstream.OrderedConsumerConfig{
+			DeliverPolicy: jetstream.DeliverLastPerSubjectPolicy,
+		})
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		res, err := c.Fetch(int(c.CachedInfo().NumPending), jetstream.FetchMaxWait(1*time.Second))
+		if err != nil {
+			t.Fatalf("Unexpected error: %s", err)
+		}
+
+		for msg := range res.Messages() {
+			msgs = append(msgs, msg)
+		}
+
+		if res.Error() != nil {
+			t.Fatalf("Unexpected error: %s", err)
+		}
+
+		if len(msgs) != 2 {
+			t.Fatalf("Expected %d messages; got: %d", 2, len(msgs))
+		}
+		expectedSubjects := []string{"FOO.A", "FOO.B"}
+
+		for i := range msgs {
+			if msgs[i].Subject() != expectedSubjects[i] {
+				t.Fatalf("Expected subject: %s; got: %s", expectedSubjects[i], msgs[i].Subject())
+			}
+		}
+	})
+
 	t.Run("consumer used as consume", func(t *testing.T) {
 		srv := RunBasicJetStreamServer()
 		defer shutdownJSServerAndRemoveStorage(t, srv)
diff --git a/jetstream/test/pull_test.go b/jetstream/test/pull_test.go
index f35aae3..4042e52 100644
--- a/jetstream/test/pull_test.go
+++ b/jetstream/test/pull_test.go
@@ -2584,6 +2584,154 @@ func TestPullConsumerConsume(t *testing.T) {
 		cc.Drain()
 		wg.Wait()
 	})
+
+	t.Run("wait for closed after drain", func(t *testing.T) {
+		srv := RunBasicJetStreamServer()
+		defer shutdownJSServerAndRemoveStorage(t, srv)
+		nc, err := nats.Connect(srv.ClientURL())
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		js, err := jetstream.New(nc)
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		defer nc.Close()
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		s, err := js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}})
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		c, err := s.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{AckPolicy: jetstream.AckExplicitPolicy})
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		msgs := make([]jetstream.Msg, 0)
+		lock := sync.Mutex{}
+		publishTestMsgs(t, js)
+		cc, err := c.Consume(func(msg jetstream.Msg) {
+			time.Sleep(50 * time.Millisecond)
+			msg.Ack()
+			lock.Lock()
+			msgs = append(msgs, msg)
+			lock.Unlock()
+		})
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		closed := cc.Closed()
+		time.Sleep(100 * time.Millisecond)
+
+		cc.Drain()
+
+		select {
+		case <-closed:
+		case <-time.After(5 * time.Second):
+			t.Fatalf("Timeout waiting for consume to be closed")
+		}
+
+		if len(msgs) != len(testMsgs) {
+			t.Fatalf("Unexpected received message count after consume closed; want %d; got %d", len(testMsgs), len(msgs))
+		}
+	})
+
+	t.Run("wait for closed after stop", func(t *testing.T) {
+		srv := RunBasicJetStreamServer()
+		defer shutdownJSServerAndRemoveStorage(t, srv)
+		nc, err := nats.Connect(srv.ClientURL())
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		js, err := jetstream.New(nc)
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		defer nc.Close()
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		s, err := js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}})
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		c, err := s.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{AckPolicy: jetstream.AckExplicitPolicy})
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		msgs := make([]jetstream.Msg, 0)
+		lock := sync.Mutex{}
+		publishTestMsgs(t, js)
+		cc, err := c.Consume(func(msg jetstream.Msg) {
+			time.Sleep(50 * time.Millisecond)
+			msg.Ack()
+			lock.Lock()
+			msgs = append(msgs, msg)
+			lock.Unlock()
+		})
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		time.Sleep(100 * time.Millisecond)
+		closed := cc.Closed()
+
+		cc.Stop()
+
+		select {
+		case <-closed:
+		case <-time.After(5 * time.Second):
+			t.Fatalf("Timeout waiting for consume to be closed")
+		}
+
+		if len(msgs) < 1 || len(msgs) > 3 {
+			t.Fatalf("Unexpected received message count after consume closed; want 1-3; got %d", len(msgs))
+		}
+	})
+
+	t.Run("wait for closed on already closed consume", func(t *testing.T) {
+		srv := RunBasicJetStreamServer()
+		defer shutdownJSServerAndRemoveStorage(t, srv)
+		nc, err := nats.Connect(srv.ClientURL())
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		js, err := jetstream.New(nc)
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		defer nc.Close()
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		s, err := js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}})
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		c, err := s.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{AckPolicy: jetstream.AckExplicitPolicy})
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		publishTestMsgs(t, js)
+		cc, err := c.Consume(func(msg jetstream.Msg) {
+			time.Sleep(50 * time.Millisecond)
+			msg.Ack()
+		})
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		time.Sleep(100 * time.Millisecond)
+
+		cc.Stop()
+
+		time.Sleep(100 * time.Millisecond)
+
+		select {
+		case <-cc.Closed():
+		case <-time.After(5 * time.Second):
+			t.Fatalf("Timeout waiting for consume to be closed")
+		}
+	})
 }
 
 func TestPullConsumerConsume_WithCluster(t *testing.T) {
diff --git a/test/js_test.go b/test/js_test.go
index 0363706..fea098a 100644
--- a/test/js_test.go
+++ b/test/js_test.go
@@ -1357,7 +1357,7 @@ func TestPullSubscribeFetchBatchWithHeartbeat(t *testing.T) {
 	if msgs.Error() != nil {
 		t.Fatalf("Unexpected error: %s", msgs.Error())
 	}
-	if elapsed < 290*time.Millisecond {
+	if elapsed < 250*time.Millisecond {
 		t.Fatalf("Expected timeout after 300ms; got: %v", elapsed)
 	}
 
@@ -8207,6 +8207,128 @@ func TestJetStreamPublishAsyncPerf(t *testing.T) {
 	fmt.Printf("%.0f msgs/sec\n\n", float64(toSend)/tt.Seconds())
 }
 
+func TestJetStreamCleanupPublisher(t *testing.T) {
+
+	t.Run("cleanup js publisher", func(t *testing.T) {
+		s := RunBasicJetStreamServer()
+		defer shutdownJSServerAndRemoveStorage(t, s)
+
+		nc, js := jsClient(t, s)
+		defer nc.Close()
+
+		// Create a stream.
+		if _, err := js.AddStream(&nats.StreamConfig{Name: "TEST", Subjects: []string{"FOO"}}); err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		numSubs := nc.NumSubscriptions()
+		if _, err := js.PublishAsync("FOO", []byte("hello")); err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		select {
+		case <-js.PublishAsyncComplete():
+		case <-time.After(5 * time.Second):
+			t.Fatalf("Did not receive completion signal")
+		}
+
+		if numSubs+1 != nc.NumSubscriptions() {
+			t.Fatalf("Expected an additional subscription after publish, got %d", nc.NumSubscriptions())
+		}
+
+		js.CleanupPublisher()
+
+		if numSubs != nc.NumSubscriptions() {
+			t.Fatalf("Expected subscriptions to be back to original count")
+		}
+	})
+
+	t.Run("cleanup js publisher, cancel pending acks", func(t *testing.T) {
+		s := RunBasicJetStreamServer()
+		defer shutdownJSServerAndRemoveStorage(t, s)
+
+		nc, err := nats.Connect(s.ClientURL())
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		defer nc.Close()
+
+		cbErr := make(chan error, 10)
+		js, err := nc.JetStream(nats.PublishAsyncErrHandler(func(js nats.JetStream, m *nats.Msg, err error) {
+			cbErr <- err
+		}))
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		// Create a stream with NoAck so that we can test that we cancel ack futures.
+		if _, err := js.AddStream(&nats.StreamConfig{Name: "TEST", Subjects: []string{"FOO"}, NoAck: true}); err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		numSubs := nc.NumSubscriptions()
+
+		var acks []nats.PubAckFuture
+		for i := 0; i < 10; i++ {
+			ack, err := js.PublishAsync("FOO", []byte("hello"))
+			if err != nil {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+			acks = append(acks, ack)
+		}
+
+		asyncComplete := js.PublishAsyncComplete()
+		select {
+		case <-asyncComplete:
+			t.Fatalf("Should not complete, NoAck is set")
+		case <-time.After(200 * time.Millisecond):
+		}
+
+		if numSubs+1 != nc.NumSubscriptions() {
+			t.Fatalf("Expected an additional subscription after publish, got %d", nc.NumSubscriptions())
+		}
+
+		js.CleanupPublisher()
+
+		if numSubs != nc.NumSubscriptions() {
+			t.Fatalf("Expected subscriptions to be back to original count")
+		}
+
+		// check that PublishAsyncComplete channel is closed
+		select {
+		case <-asyncComplete:
+		case <-time.After(5 * time.Second):
+			t.Fatalf("Did not receive completion signal")
+		}
+
+		// check that all ack futures are cancelled
+		for _, ack := range acks {
+			select {
+			case err := <-ack.Err():
+				if !errors.Is(err, nats.ErrJetStreamPublisherClosed) {
+					t.Fatalf("Expected JetStreamContextClosed error, got %v", err)
+				}
+			case <-ack.Ok():
+				t.Fatalf("Expected error on the ack future")
+			case <-time.After(200 * time.Millisecond):
+				t.Fatalf("Expected an error on the ack future")
+			}
+		}
+
+		// check that async error handler is called for each pending ack
+		for i := 0; i < 10; i++ {
+			select {
+			case err := <-cbErr:
+				if !errors.Is(err, nats.ErrJetStreamPublisherClosed) {
+					t.Fatalf("Expected JetStreamContextClosed error, got %v", err)
+				}
+			case <-time.After(200 * time.Millisecond):
+				t.Fatalf("Expected errors to be passed from the async handler")
+			}
+		}
+	})
+
+}
+
 func TestJetStreamPublishExpectZero(t *testing.T) {
 	s := RunBasicJetStreamServer()
 	defer shutdownJSServerAndRemoveStorage(t, s)