Repository: nats-io/nats.go
Pull request #1843 (2025-04-02 11:10:01)

[ADDED] DefaultTimeout option for JetStream API requests
This adds a `WithDefaultTimeout` option when creating `JetStream`. This timeout is used when a context without deadline is passed to JetStream requests.

Signed-off-by: Piotr Piotrowski <piotr@synadia.com>


Patch:
diff --git a/jetstream/jetstream.go b/jetstream/jetstream.go
index fdc0956..a99c2a0 100644
--- a/jetstream/jetstream.go
+++ b/jetstream/jetstream.go
@@ -304,6 +304,11 @@ type (
 		// Domain is the domain name token used when sending JetStream requests.
 		Domain string
 
+		// DefaultTimeout is the default timeout used for JetStream API requests.
+		// This applies when the context passed to JetStream methods does not have
+		// a deadline set.
+		DefaultTimeout time.Duration
+
 		publisherOpts asyncPublisherOpts
 
 		// this is the actual prefix used in the API requests
@@ -407,6 +412,7 @@ func New(nc *nats.Conn, opts ...JetStreamOpt) (JetStream, error) {
 		publisherOpts: asyncPublisherOpts{
 			maxpa: defaultAsyncPubAckInflight,
 		},
+		DefaultTimeout: defaultAPITimeout,
 	}
 	setReplyPrefix(nc, &jsOpts)
 	for _, opt := range opts {
@@ -451,7 +457,8 @@ func NewWithAPIPrefix(nc *nats.Conn, apiPrefix string, opts ...JetStreamOpt) (Je
 		publisherOpts: asyncPublisherOpts{
 			maxpa: defaultAsyncPubAckInflight,
 		},
-		APIPrefix: apiPrefix,
+		APIPrefix:      apiPrefix,
+		DefaultTimeout: defaultAPITimeout,
 	}
 	setReplyPrefix(nc, &jsOpts)
 	for _, opt := range opts {
@@ -488,7 +495,8 @@ func NewWithDomain(nc *nats.Conn, domain string, opts ...JetStreamOpt) (JetStrea
 		publisherOpts: asyncPublisherOpts{
 			maxpa: defaultAsyncPubAckInflight,
 		},
-		Domain: domain,
+		Domain:         domain,
+		DefaultTimeout: defaultAPITimeout,
 	}
 	setReplyPrefix(nc, &jsOpts)
 	for _, opt := range opts {
@@ -1089,7 +1097,7 @@ func (js *jetStream) wrapContextWithoutDeadline(ctx context.Context) (context.Co
 	if _, ok := ctx.Deadline(); ok {
 		return ctx, nil
 	}
-	return context.WithTimeout(ctx, defaultAPITimeout)
+	return context.WithTimeout(ctx, js.opts.DefaultTimeout)
 }
 
 // CleanupPublisher will cleanup the publishing side of JetStreamContext.
diff --git a/jetstream/jetstream_options.go b/jetstream/jetstream_options.go
index 579a6a7..40cb2a4 100644
--- a/jetstream/jetstream_options.go
+++ b/jetstream/jetstream_options.go
@@ -65,6 +65,19 @@ func WithPublishAsyncTimeout(dur time.Duration) JetStreamOpt {
 	}
 }
 
+// WithDefaultTimeout sets the default timeout for JetStream API requests.
+// It is used when context used for the request does not have a deadline set.
+// If not provided, a default of 5 seconds will be used.
+func WithDefaultTimeout(timeout time.Duration) JetStreamOpt {
+	return func(opts *JetStreamOptions) error {
+		if timeout <= 0 {
+			return fmt.Errorf("%w: timeout value must be greater than 0", ErrInvalidOption)
+		}
+		opts.DefaultTimeout = timeout
+		return nil
+	}
+}
+
 // WithPurgeSubject sets a specific subject for which messages on a stream will
 // be purged
 func WithPurgeSubject(subject string) StreamPurgeOpt {
diff --git a/js.go b/js.go
index 749fb23..4c956f1 100644
--- a/js.go
+++ b/js.go
@@ -3073,6 +3073,11 @@ func (sub *Subscription) Fetch(batch int, opts ...PullOpt) ([]*Msg, error) {
 		}
 	}
 	var hbTimer *time.Timer
+	defer func() {
+		if hbTimer != nil {
+			hbTimer.Stop()
+		}
+	}()
 	var hbErr error
 	sub.mu.Lock()
 	subClosed := sub.closed || sub.draining
@@ -3081,6 +3086,7 @@ func (sub *Subscription) Fetch(batch int, opts ...PullOpt) ([]*Msg, error) {
 		err = errors.Join(ErrBadSubscription, ErrSubscriptionClosed)
 	}
 	hbLock := sync.Mutex{}
+	var disconnected atomic.Bool
 	if err == nil && len(msgs) < batch && !subClosed {
 		// For batch real size of 1, it does not make sense to set no_wait in
 		// the request.
@@ -3135,7 +3141,17 @@ func (sub *Subscription) Fetch(batch int, opts ...PullOpt) ([]*Msg, error) {
 			}
 			return nil
 		}
-
+		connStatusChanged := nc.StatusChanged()
+		go func() {
+			select {
+			case <-ctx.Done():
+				return
+			case <-connStatusChanged:
+				disconnected.Store(true)
+				cancel()
+				return
+			}
+		}()
 		err = sendReq()
 		for err == nil && len(msgs) < batch {
 			// Ask for next message and wait if there are no messages
@@ -3162,9 +3178,6 @@ func (sub *Subscription) Fetch(batch int, opts ...PullOpt) ([]*Msg, error) {
 				}
 			}
 		}
-		if hbTimer != nil {
-			hbTimer.Stop()
-		}
 	}
 	// If there is at least a message added to msgs, then need to return OK and no error
 	if err != nil && len(msgs) == 0 {
@@ -3173,6 +3186,9 @@ func (sub *Subscription) Fetch(batch int, opts ...PullOpt) ([]*Msg, error) {
 		if hbErr != nil {
 			return nil, hbErr
 		}
+		if disconnected.Load() {
+			return nil, ErrFetchDisconnected
+		}
 		return nil, o.checkCtxErr(err)
 	}
 	return msgs, nil
@@ -3383,6 +3399,18 @@ func (sub *Subscription) FetchBatch(batch int, opts ...PullOpt) (MessageBatch, e
 	}
 	expires := ttl - expiresDiff
 
+	connStatusChanged := nc.StatusChanged()
+	var disconnected atomic.Bool
+	go func() {
+		select {
+		case <-ctx.Done():
+			return
+		case <-connStatusChanged:
+			disconnected.Store(true)
+			cancel()
+			return
+		}
+	}()
 	requestBatch := batch - len(result.msgs)
 	req := nextRequest{
 		Expires:   expires,
@@ -3407,6 +3435,11 @@ func (sub *Subscription) FetchBatch(batch int, opts ...PullOpt) (MessageBatch, e
 		return result, nil
 	}
 	var hbTimer *time.Timer
+	defer func() {
+		if hbTimer != nil {
+			hbTimer.Stop()
+		}
+	}()
 	var hbErr error
 	if o.hb > 0 {
 		hbTimer = time.AfterFunc(2*o.hb, func() {
@@ -3453,6 +3486,8 @@ func (sub *Subscription) FetchBatch(batch int, opts ...PullOpt) (MessageBatch, e
 			result.Lock()
 			if hbErr != nil {
 				result.err = hbErr
+			} else if disconnected.Load() {
+				result.err = ErrFetchDisconnected
 			} else {
 				result.err = o.checkCtxErr(err)
 			}
diff --git a/jserrors.go b/jserrors.go
index 65e78cf..6d72bbc 100644
--- a/jserrors.go
+++ b/jserrors.go
@@ -164,6 +164,10 @@ var (
 	// ErrTooManyStalledMsgs is returned when too many outstanding async
 	// messages are waiting for ack.
 	ErrTooManyStalledMsgs JetStreamError = &jsError{message: "stalled with too many outstanding async published messages"}
+
+	// ErrFetchDisconnected is returned when the connection to the server is lost
+	// while waiting for messages to be delivered on PullSubscribe.
+	ErrFetchDisconnected = &jsError{message: "disconnected during fetch"}
 )
 
 // Error code represents JetStream error codes returned by the API
diff --git a/nats.go b/nats.go
index 3f9a8ba..085138e 100644
--- a/nats.go
+++ b/nats.go
@@ -2224,6 +2224,7 @@ func (nc *Conn) ForceReconnect() error {
 		// even if we're in the middle of a backoff
 		if nc.rqch != nil {
 			close(nc.rqch)
+			nc.rqch = nil
 		}
 		return nil
 	}
@@ -2843,6 +2844,16 @@ func (nc *Conn) doReconnect(err error, forceReconnect bool) {
 	var rt *time.Timer
 	// Channel used to kick routine out of sleep when conn is closed.
 	rqch := nc.rqch
+
+	// if rqch is nil, we need to set it up to signal
+	// the reconnect loop to reconnect immediately
+	// this means that `ForceReconnect` was called
+	// before entering doReconnect
+	if rqch == nil {
+		rqch = make(chan struct{})
+		close(rqch)
+	}
+
 	// Counter that is increased when the whole list of servers has been tried.
 	var wlf int
 
@@ -4600,6 +4611,8 @@ func (s *Subscription) StatusChanged(statuses ...SubStatus) <-chan SubStatus {
 		statuses = []SubStatus{SubscriptionActive, SubscriptionDraining, SubscriptionClosed, SubscriptionSlowConsumer}
 	}
 	ch := make(chan SubStatus, 10)
+	s.mu.Lock()
+	defer s.mu.Unlock()
 	for _, status := range statuses {
 		s.registerStatusChangeListener(status, ch)
 		// initial status
@@ -4613,9 +4626,8 @@ func (s *Subscription) StatusChanged(statuses ...SubStatus) <-chan SubStatus {
 // registerStatusChangeListener registers a channel waiting for a specific status change event.
 // Status change events are non-blocking - if no receiver is waiting for the status change,
 // it will not be sent on the channel. Closed channels are ignored.
+// Lock should be held entering.
 func (s *Subscription) registerStatusChangeListener(status SubStatus, ch chan SubStatus) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
 	if s.statListeners == nil {
 		s.statListeners = make(map[chan SubStatus][]SubStatus)
 	}
@@ -5893,6 +5905,8 @@ func (nc *Conn) StatusChanged(statuses ...Status) chan Status {
 		statuses = []Status{CONNECTED, RECONNECTING, DISCONNECTED, CLOSED}
 	}
 	ch := make(chan Status, 10)
+	nc.mu.Lock()
+	defer nc.mu.Unlock()
 	for _, s := range statuses {
 		nc.registerStatusChangeListener(s, ch)
 	}
@@ -5902,9 +5916,8 @@ func (nc *Conn) StatusChanged(statuses ...Status) chan Status {
 // registerStatusChangeListener registers a channel waiting for a specific status change event.
 // Status change events are non-blocking - if no receiver is waiting for the status change,
 // it will not be sent on the channel. Closed channels are ignored.
+// The lock should be held entering.
 func (nc *Conn) registerStatusChangeListener(status Status, ch chan Status) {
-	nc.mu.Lock()
-	defer nc.mu.Unlock()
 	if nc.statListeners == nil {
 		nc.statListeners = make(map[Status][]chan Status)
 	}


Test patch:
diff --git a/jetstream/test/publish_test.go b/jetstream/test/publish_test.go
index f581fc8..64904bd 100644
--- a/jetstream/test/publish_test.go
+++ b/jetstream/test/publish_test.go
@@ -774,6 +774,41 @@ func TestPublish(t *testing.T) {
 	}
 }
 
+func TestPublishTimeout(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, jetstream.WithDefaultTimeout(200*time.Millisecond))
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	defer nc.Close()
+
+	// create stream with no ack to force timeout
+	_, err = js.CreateStream(context.Background(), jetstream.StreamConfig{
+		Name:     "foo",
+		Subjects: []string{"FOO.*"},
+		NoAck:    true,
+	})
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	now := time.Now()
+	_, err = js.Publish(context.Background(), "FOO.1", []byte("msg"))
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("Expected deadline exceeded error; got: %v", err)
+	}
+	since := time.Since(now)
+	if since < 200*time.Millisecond || since > 500*time.Millisecond {
+		t.Fatalf("Expected timeout to be around 200ms; got: %v", since)
+	}
+}
+
 func TestPublishMsgAsync(t *testing.T) {
 	type publishConfig struct {
 		msg              *nats.Msg
diff --git a/test/js_test.go b/test/js_test.go
index 94702a8..769ad71 100644
--- a/test/js_test.go
+++ b/test/js_test.go
@@ -11152,3 +11152,65 @@ func TestJetStreamSubscribeConsumerCreateTimeout(t *testing.T) {
 		t.Fatalf("Expected error")
 	}
 }
+
+func TestJetStreamPullSubscribeFetchErrOnReconnect(t *testing.T) {
+	srv := RunBasicJetStreamServer()
+	defer shutdownJSServerAndRemoveStorage(t, srv)
+	nc, js := jsClient(t, srv)
+	defer nc.Close()
+	_, err := js.AddStream(&nats.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}})
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	sub, err := js.PullSubscribe("FOO.123", "bar")
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	defer sub.Unsubscribe()
+	errs := make(chan error, 1)
+	go func() {
+		time.Sleep(100 * time.Millisecond)
+		errs <- nc.ForceReconnect()
+	}()
+	_, err = sub.Fetch(1, nats.MaxWait(time.Second))
+	if !errors.Is(err, nats.ErrFetchDisconnected) {
+		t.Fatalf("Expected error: %v; got: %v", nats.ErrFetchDisconnected, err)
+	}
+	if err := <-errs; err != nil {
+		t.Fatalf("Error on reconnect: %v", err)
+	}
+}
+
+func TestJetStreamPullSubscribeFetchBatchErrOnReconnect(t *testing.T) {
+	srv := RunBasicJetStreamServer()
+	defer shutdownJSServerAndRemoveStorage(t, srv)
+	nc, js := jsClient(t, srv)
+	defer nc.Close()
+	_, err := js.AddStream(&nats.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}})
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	sub, err := js.PullSubscribe("FOO.123", "bar")
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	defer sub.Unsubscribe()
+	errs := make(chan error, 1)
+	go func() {
+		time.Sleep(100 * time.Millisecond)
+		errs <- nc.ForceReconnect()
+	}()
+	msgs, err := sub.FetchBatch(1, nats.MaxWait(time.Second), nats.PullHeartbeat(100*time.Millisecond))
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	for range msgs.Messages() {
+		t.Fatalf("Expected no messages, got one")
+	}
+	if !errors.Is(msgs.Error(), nats.ErrFetchDisconnected) {
+		t.Fatalf("Expected error: %v; got: %v", nats.ErrFetchDisconnected, msgs.Error())
+	}
+	if err := <-errs; err != nil {
+		t.Fatalf("Error on reconnect: %v", err)
+	}
+}
diff --git a/test/reconnect_test.go b/test/reconnect_test.go
index 9fc1b23..1b30b5b 100644
--- a/test/reconnect_test.go
+++ b/test/reconnect_test.go
@@ -1018,6 +1018,55 @@ func TestForceReconnectDisallowReconnect(t *testing.T) {
 
 }
 
+func TestForceReconnectSubsequentCalls(t *testing.T) {
+	s := RunDefaultServer()
+	defer s.Shutdown()
+
+	nc, err := nats.Connect(s.ClientURL(), nats.ReconnectWait(10*time.Second))
+	if err != nil {
+		t.Fatalf("Unexpected error on connect: %v", err)
+	}
+	defer nc.Close()
+
+	statusCh := nc.StatusChanged(nats.RECONNECTING, nats.CONNECTED)
+	defer close(statusCh)
+	newStatus := make(chan nats.Status, 10)
+	// non-blocking channel, so we need to be constantly listening
+	go func() {
+		for {
+			s, ok := <-statusCh
+			if !ok {
+				return
+			}
+			newStatus <- s
+		}
+	}()
+
+	for range 10 {
+		err = nc.ForceReconnect()
+		if err != nil {
+			t.Fatalf("Unexpected error on reconnect: %v", err)
+		}
+	}
+
+	WaitOnChannel(t, newStatus, nats.RECONNECTING)
+	WaitOnChannel(t, newStatus, nats.CONNECTED)
+
+	// check that we did not try to reconnect again
+	select {
+	case <-newStatus:
+		t.Fatal("Should not have received a new status")
+	case <-time.After(200 * time.Millisecond):
+	}
+
+	// now force a reconnect again
+	if err := nc.ForceReconnect(); err != nil {
+		t.Fatalf("Unexpected error on reconnect: %v", err)
+	}
+	WaitOnChannel(t, newStatus, nats.RECONNECTING)
+	WaitOnChannel(t, newStatus, nats.CONNECTED)
+}
+
 func TestAuthExpiredForceReconnect(t *testing.T) {
 	ts := runTrustServer()
 	defer ts.Shutdown()