Repository: nats-io/nats.go
Pull request #1690 (2024-07-31 13:30:18)

[ADDED] CleanupPublisher method for removing internal JetStream subscription
Signed-off-by: Piotr Piotrowski <piotr@synadia.com>


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 4696244..9f556bb 100644
--- a/jetstream/ordered.go
+++ b/jetstream/ordered.go
@@ -69,7 +69,10 @@ const (
 	consumerTypeFetch
 )
 
-var errOrderedSequenceMismatch = errors.New("sequence mismatch")
+var (
+	errOrderedSequenceMismatch = errors.New("sequence mismatch")
+	errOrderedConsumerClosed   = errors.New("ordered consumer closed")
+)
 
 // Consume can be used to continuously receive messages and handle them
 // with the provided callback function. Consume cannot be used concurrently
@@ -142,6 +145,9 @@ func (c *orderedConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt
 			select {
 			case <-c.doReset:
 				if err := c.reset(); err != nil {
+					if errors.Is(err, errOrderedConsumerClosed) {
+						continue
+					}
 					c.errHandler(c.serial)(c.currentSub, err)
 				}
 				if c.withStopAfter {
@@ -173,6 +179,12 @@ func (c *orderedConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt
 					c.Unlock()
 				}
 			case <-sub.done:
+				s := sub.consumer.currentSub
+				if s != nil {
+					sub.consumer.Lock()
+					s.Stop()
+					sub.consumer.Unlock()
+				}
 				return
 			case msgsLeft, ok := <-c.stopAfterMsgsLeft:
 				if !ok {
@@ -276,6 +288,9 @@ func (s *orderedSubscription) Next() (Msg, error) {
 				s.opts[len(s.opts)-1] = StopAfter(s.consumer.stopAfter)
 			}
 			if err := s.consumer.reset(); err != nil {
+				if errors.Is(err, errOrderedConsumerClosed) {
+					return nil, ErrMsgIteratorClosed
+				}
 				return nil, err
 			}
 			cc, err := s.consumer.currentConsumer.Messages(s.opts...)
@@ -297,6 +312,9 @@ func (s *orderedSubscription) Next() (Msg, error) {
 		dseq := meta.Sequence.Consumer
 		if dseq != s.consumer.cursor.deliverSeq+1 {
 			if err := s.consumer.reset(); err != nil {
+				if errors.Is(err, errOrderedConsumerClosed) {
+					return nil, ErrMsgIteratorClosed
+				}
 				return nil, err
 			}
 			cc, err := s.consumer.currentConsumer.Messages(s.opts...)
@@ -318,7 +336,9 @@ func (s *orderedSubscription) Stop() {
 	}
 	s.consumer.Lock()
 	defer s.consumer.Unlock()
-	s.consumer.currentSub.Stop()
+	if s.consumer.currentSub != nil {
+		s.consumer.currentSub.Stop()
+	}
 	close(s.done)
 }
 
@@ -326,9 +346,11 @@ func (s *orderedSubscription) Drain() {
 	if !atomic.CompareAndSwapUint32(&s.closed, 0, 1) {
 		return
 	}
-	s.consumer.currentConsumer.Lock()
-	defer s.consumer.currentConsumer.Unlock()
-	s.consumer.currentSub.Drain()
+	if s.consumer.currentSub != nil {
+		s.consumer.currentConsumer.Lock()
+		s.consumer.currentSub.Drain()
+		s.consumer.currentConsumer.Unlock()
+	}
 	close(s.done)
 }
 
@@ -504,7 +526,7 @@ func (c *orderedConsumer) reset() error {
 	err = retryWithBackoff(func(attempt int) (bool, error) {
 		isClosed := atomic.LoadUint32(&c.subscription.closed) == 1
 		if isClosed {
-			return false, nil
+			return false, errOrderedConsumerClosed
 		}
 		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 		defer cancel()
@@ -512,7 +534,6 @@ func (c *orderedConsumer) reset() error {
 		if err != nil {
 			return true, err
 		}
-		c.currentConsumer = cons.(*pullConsumer)
 		return false, nil
 	}, backoffOpts)
 	if err != nil {
@@ -562,7 +583,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/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/kv_test.go b/jetstream/test/kv_test.go
index 26a038c..010a7d5 100644
--- a/jetstream/test/kv_test.go
+++ b/jetstream/test/kv_test.go
@@ -967,6 +967,7 @@ func TestKeyValueListKeys(t *testing.T) {
 
 func TestKeyValueCrossAccounts(t *testing.T) {
 	conf := createConfFile(t, []byte(`
+		listen: 127.0.0.1:-1
         jetstream: enabled
         accounts: {
            A: {
diff --git a/jetstream/test/ordered_test.go b/jetstream/test/ordered_test.go
index 01a67a8..2a7960d 100644
--- a/jetstream/test/ordered_test.go
+++ b/jetstream/test/ordered_test.go
@@ -539,6 +539,45 @@ func TestOrderedConsumerConsume(t *testing.T) {
 		cc.Drain()
 		wg.Wait()
 	})
+
+	t.Run("stop consume during reset", 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)
+		}
+		for i := 0; i < 10; i++ {
+			c, err := s.OrderedConsumer(context.Background(), jetstream.OrderedConsumerConfig{})
+			if err != nil {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+			cc, err := c.Consume(func(msg jetstream.Msg) {
+				msg.Ack()
+			})
+			if err != nil {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+			if err := s.DeleteConsumer(context.Background(), c.CachedInfo().Name); err != nil {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+			cc.Stop()
+			time.Sleep(50 * time.Millisecond)
+		}
+	})
 }
 
 func TestOrderedConsumerMessages(t *testing.T) {
@@ -1080,6 +1119,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/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)