Repository: nats-io/nats.go
Pull request #1548 (2024-02-07 20:23:45)

[ADDED] FetchHeartbeat option for Fetch and FetchBytes
Signed-off-by: Piotr Piotrowski <piotr@synadia.com>


Patch:
diff --git a/jetstream/consumer.go b/jetstream/consumer.go
index 3023413..7792ced 100644
--- a/jetstream/consumer.go
+++ b/jetstream/consumer.go
@@ -52,6 +52,12 @@ type (
 		// defaults to 30 seconds and can be configured using FetchMaxWait
 		// option.
 		//
+		// By default, Fetch uses a 5s idle heartbeat for requests longer than
+		// 10 seconds. For shorter requests, the idle heartbeat is disabled.
+		// This can be configured using FetchHeartbeat option. If a client does
+		// not receive a heartbeat message from a stream for more than 2 times
+		// the idle heartbeat setting, Fetch will return [ErrNoHeartbeat].
+		//
 		// Fetch is non-blocking and returns MessageBatch, exposing a channel
 		// for delivered messages.
 		//
@@ -65,6 +71,12 @@ type (
 		// timeout defaults to 30 seconds and can be configured using
 		// FetchMaxWait option.
 		//
+		// By default, FetchBytes uses a 5s idle heartbeat for requests longer than
+		// 10 seconds. For shorter requests, the idle heartbeat is disabled.
+		// This can be configured using FetchHeartbeat option. If a client does
+		// not receive a heartbeat message from a stream for more than 2 times
+		// the idle heartbeat setting, Fetch will return ErrNoHeartbeat.
+		//
 		// FetchBytes is non-blocking and returns MessageBatch, exposing a channel
 		// for delivered messages.
 		//
@@ -75,9 +87,7 @@ type (
 		// FetchNoWait is used to retrieve up to a provided number of messages
 		// from a stream. Unlike Fetch, FetchNoWait will only deliver messages
 		// that are currently available in the stream and will not wait for new
-		// messages to arrive, even if batch size is not met. FetchNoWait
-		// timeout defaults to 30 seconds and can be configured using
-		// FetchMaxWait option.
+		// messages to arrive, even if batch size is not met.
 		//
 		// FetchNoWait is non-blocking and returns MessageBatch, exposing a
 		// channel for delivered messages.
diff --git a/jetstream/errors.go b/jetstream/errors.go
index bebc1d0..5ed5176 100644
--- a/jetstream/errors.go
+++ b/jetstream/errors.go
@@ -67,7 +67,7 @@ const (
 )
 
 var (
-	// API errors
+	// JetStream API errors
 
 	// ErrJetStreamNotEnabled is an error returned when JetStream is not
 	// enabled.
@@ -284,6 +284,10 @@ var (
 	// name.
 	ErrInvalidKey JetStreamError = &jsError{message: "invalid key"}
 
+	// ErrBucketExists is returned when attempting to create a bucket that
+	// already exists and has a different configuration.
+	ErrBucketExists JetStreamError = &jsError{message: "bucket name already in use"}
+
 	// ErrBucketNotFound is returned when attempting to access a bucket that
 	// does not exist.
 	ErrBucketNotFound JetStreamError = &jsError{message: "bucket not found"}
diff --git a/jetstream/jetstream_options.go b/jetstream/jetstream_options.go
index 53fbd71..699a901 100644
--- a/jetstream/jetstream_options.go
+++ b/jetstream/jetstream_options.go
@@ -260,6 +260,8 @@ func WithMessagesErrOnMissingHeartbeat(hbErr bool) PullMessagesOpt {
 }
 
 // FetchMaxWait sets custom timeout for fetching predefined batch of messages.
+//
+// If not provided, a default of 30 seconds will be used.
 func FetchMaxWait(timeout time.Duration) FetchOpt {
 	return func(req *pullRequest) error {
 		if timeout <= 0 {
@@ -270,6 +272,24 @@ func FetchMaxWait(timeout time.Duration) FetchOpt {
 	}
 }
 
+// FetchHeartbeat sets custom heartbeat for individual fetch request. If a
+// client does not receive a heartbeat message from a stream for more than 2
+// times the idle heartbeat setting, Fetch will return [ErrNoHeartbeat].
+//
+// Heartbeat value has to be lower than FetchMaxWait / 2.
+//
+// If not provided, heartbeat will is set to 5s for requests with FetchMaxWait > 30s
+// and disabled otherwise.
+func FetchHeartbeat(hb time.Duration) FetchOpt {
+	return func(req *pullRequest) error {
+		if hb <= 0 {
+			return fmt.Errorf("%w: timeout value must be greater than 0", ErrInvalidOption)
+		}
+		req.Heartbeat = hb
+		return nil
+	}
+}
+
 // WithDeletedDetails can be used to display the information about messages
 // deleted from a stream on a stream info request
 func WithDeletedDetails(deletedDetails bool) StreamInfoOpt {
diff --git a/jetstream/kv.go b/jetstream/kv.go
index 9371f6f..42a86c5 100644
--- a/jetstream/kv.go
+++ b/jetstream/kv.go
@@ -31,16 +31,35 @@ type (
 	// KeyValueManager is used to manage KeyValue stores. It provides methods to
 	// create, delete, and retrieve KeyValue stores.
 	KeyValueManager interface {
-		// KeyValue will lookup and bind to an existing KeyValue store. If the
-		// KeyValue store with given name does not exist, ErrBucketNotFound will
-		// be returned.
+		// KeyValue will lookup and bind to an existing KeyValue store.
+		//
+		// If the KeyValue store with given name does not exist,
+		// ErrBucketNotFound will be returned.
 		KeyValue(ctx context.Context, bucket string) (KeyValue, error)
 
 		// CreateKeyValue will create a KeyValue store with the given
 		// configuration.
+		//
+		// If a KeyValue store with the same name already exists and the
+		// configuration is different, ErrBucketExists will be returned.
 		CreateKeyValue(ctx context.Context, cfg KeyValueConfig) (KeyValue, error)
 
+		// UpdateKeyValue will update an existing KeyValue store with the given
+		// configuration.
+		//
+		// If a KeyValue store with the given name does not exist, ErrBucketNotFound
+		// will be returned.
+		UpdateKeyValue(ctx context.Context, cfg KeyValueConfig) (KeyValue, error)
+
+		// CreateOrUpdateKeyValue will create a KeyValue store if it does not
+		// exist or update an existing KeyValue store with the given
+		// configuration (if possible).
+		CreateOrUpdateKeyValue(ctx context.Context, cfg KeyValueConfig) (KeyValue, error)
+
 		// DeleteKeyValue will delete this KeyValue store.
+		//
+		// If the KeyValue store with given name does not exist,
+		// ErrBucketNotFound will be returned.
 		DeleteKeyValue(ctx context.Context, bucket string) error
 
 		// KeyValueStoreNames is used to retrieve a list of key value store
@@ -457,20 +476,81 @@ func (js *jetStream) KeyValue(ctx context.Context, bucket string) (KeyValue, err
 	return mapStreamToKVS(js, pushJS, stream), nil
 }
 
-// CreateKeyValue will create a KeyValue store with the following configuration.
 func (js *jetStream) CreateKeyValue(ctx context.Context, cfg KeyValueConfig) (KeyValue, error) {
+	scfg, err := js.prepareKeyValueConfig(ctx, cfg)
+	if err != nil {
+		return nil, err
+	}
+
+	stream, err := js.CreateStream(ctx, scfg)
+	if err != nil {
+		if errors.Is(err, ErrStreamNameAlreadyInUse) {
+			// errors are joined so that backwards compatibility is retained
+			// and previous checks for ErrStreamNameAlreadyInUse will still work.
+			err = errors.Join(fmt.Errorf("%w: %s", ErrBucketExists, cfg.Bucket), err)
+		}
+		return nil, err
+	}
+	pushJS, err := js.legacyJetStream()
+	if err != nil {
+		return nil, err
+	}
+
+	return mapStreamToKVS(js, pushJS, stream), nil
+}
+
+func (js *jetStream) UpdateKeyValue(ctx context.Context, cfg KeyValueConfig) (KeyValue, error) {
+	scfg, err := js.prepareKeyValueConfig(ctx, cfg)
+	if err != nil {
+		return nil, err
+	}
+
+	stream, err := js.UpdateStream(ctx, scfg)
+	if err != nil {
+		if errors.Is(err, ErrStreamNotFound) {
+			err = fmt.Errorf("%w: %s", ErrBucketNotFound, cfg.Bucket)
+		}
+		return nil, err
+	}
+	pushJS, err := js.legacyJetStream()
+	if err != nil {
+		return nil, err
+	}
+
+	return mapStreamToKVS(js, pushJS, stream), nil
+}
+
+func (js *jetStream) CreateOrUpdateKeyValue(ctx context.Context, cfg KeyValueConfig) (KeyValue, error) {
+	scfg, err := js.prepareKeyValueConfig(ctx, cfg)
+	if err != nil {
+		return nil, err
+	}
+
+	stream, err := js.CreateOrUpdateStream(ctx, scfg)
+	if err != nil {
+		return nil, err
+	}
+	pushJS, err := js.legacyJetStream()
+	if err != nil {
+		return nil, err
+	}
+
+	return mapStreamToKVS(js, pushJS, stream), nil
+}
+
+func (js *jetStream) prepareKeyValueConfig(ctx context.Context, cfg KeyValueConfig) (StreamConfig, error) {
 	if !validBucketRe.MatchString(cfg.Bucket) {
-		return nil, ErrInvalidBucketName
+		return StreamConfig{}, ErrInvalidBucketName
 	}
 	if _, err := js.AccountInfo(ctx); err != nil {
-		return nil, err
+		return StreamConfig{}, err
 	}
 
 	// Default to 1 for history. Max is 64 for now.
 	history := int64(1)
 	if cfg.History > 0 {
 		if cfg.History > KeyValueMaxHistory {
-			return nil, ErrHistoryTooLarge
+			return StreamConfig{}, ErrHistoryTooLarge
 		}
 		history = int64(cfg.History)
 	}
@@ -551,16 +631,7 @@ func (js *jetStream) CreateKeyValue(ctx context.Context, cfg KeyValueConfig) (Ke
 		scfg.Subjects = []string{fmt.Sprintf(kvSubjectsTmpl, cfg.Bucket)}
 	}
 
-	stream, err := js.CreateStream(ctx, scfg)
-	if err != nil {
-		return nil, err
-	}
-	pushJS, err := js.legacyJetStream()
-	if err != nil {
-		return nil, err
-	}
-
-	return mapStreamToKVS(js, pushJS, stream), nil
+	return scfg, nil
 }
 
 // DeleteKeyValue will delete this KeyValue store (JetStream stream).
@@ -569,7 +640,13 @@ func (js *jetStream) DeleteKeyValue(ctx context.Context, bucket string) error {
 		return ErrInvalidBucketName
 	}
 	stream := fmt.Sprintf(kvBucketNameTmpl, bucket)
-	return js.DeleteStream(ctx, stream)
+	if err := js.DeleteStream(ctx, stream); err != nil {
+		if errors.Is(err, ErrStreamNotFound) {
+			err = errors.Join(fmt.Errorf("%w: %s", ErrBucketNotFound, bucket), err)
+		}
+		return err
+	}
+	return nil
 }
 
 // KeyValueStoreNames is used to retrieve a list of key value store names
diff --git a/jetstream/object.go b/jetstream/object.go
index cc931e3..32715d7 100644
--- a/jetstream/object.go
+++ b/jetstream/object.go
@@ -116,6 +116,7 @@ type (
 		Storage     StorageType   `json:"storage,omitempty"`
 		Replicas    int           `json:"num_replicas,omitempty"`
 		Placement   *Placement    `json:"placement,omitempty"`
+		Compression bool          `json:"compression,omitempty"`
 
 		// Bucket-specific metadata
 		// NOTE: Metadata requires nats-server v2.10.0+
@@ -151,6 +152,8 @@ type (
 		BackingStore() string
 		// Metadata is the user supplied metadata for the bucket
 		Metadata() map[string]string
+		// IsCompressed indicates if the data is compressed on disk
+		IsCompressed() bool
 	}
 
 	// ObjectMetaOptions
@@ -271,7 +274,10 @@ func (js *jetStream) CreateObjectStore(ctx context.Context, cfg ObjectStoreConfi
 	if maxBytes == 0 {
 		maxBytes = -1
 	}
-
+	var compression StoreCompression
+	if cfg.Compression {
+		compression = S2Compression
+	}
 	scfg := StreamConfig{
 		Name:        fmt.Sprintf(objNameTmpl, name),
 		Description: cfg.Description,
@@ -285,6 +291,7 @@ func (js *jetStream) CreateObjectStore(ctx context.Context, cfg ObjectStoreConfi
 		AllowRollup: true,
 		AllowDirect: true,
 		Metadata:    cfg.Metadata,
+		Compression: compression,
 	}
 
 	// Create our stream.
@@ -1101,6 +1108,9 @@ func (s *ObjectBucketStatus) Metadata() map[string]string { return s.nfo.Config.
 // StreamInfo is the stream info retrieved to create the status
 func (s *ObjectBucketStatus) StreamInfo() *StreamInfo { return s.nfo }
 
+// IsCompressed indicates if the data is compressed on disk
+func (s *ObjectBucketStatus) IsCompressed() bool { return s.nfo.Config.Compression != NoCompression }
+
 // Status retrieves run-time status about a bucket
 func (obs *obs) Status(ctx context.Context) (ObjectStoreStatus, error) {
 	nfo, err := obs.stream.Info(ctx)
diff --git a/jetstream/pull.go b/jetstream/pull.go
index f7c1f05..bb5479a 100644
--- a/jetstream/pull.go
+++ b/jetstream/pull.go
@@ -729,17 +729,26 @@ func (s *pullSubscription) Drain() {
 // It will wait up to provided expiry time if not all messages are available.
 func (p *pullConsumer) Fetch(batch int, opts ...FetchOpt) (MessageBatch, error) {
 	req := &pullRequest{
-		Batch:   batch,
-		Expires: DefaultExpires,
+		Batch:     batch,
+		Expires:   DefaultExpires,
+		Heartbeat: unset,
 	}
 	for _, opt := range opts {
 		if err := opt(req); err != nil {
 			return nil, err
 		}
 	}
-	// for longer pulls, set heartbeat value
-	if req.Expires >= 10*time.Second {
-		req.Heartbeat = 5 * time.Second
+	// if heartbeat was not explicitly set, set it to 5 seconds for longer pulls
+	// and disable it for shorter pulls
+	if req.Heartbeat == unset {
+		if req.Expires >= 10*time.Second {
+			req.Heartbeat = 5 * time.Second
+		} else {
+			req.Heartbeat = 0
+		}
+	}
+	if req.Expires < 2*req.Heartbeat {
+		return nil, fmt.Errorf("%w: expiry time should be at least 2 times the heartbeat", ErrInvalidOption)
 	}
 
 	return p.fetch(req)
@@ -748,26 +757,35 @@ func (p *pullConsumer) Fetch(batch int, opts ...FetchOpt) (MessageBatch, error)
 // FetchBytes is used to retrieve up to a provided bytes from the stream.
 func (p *pullConsumer) FetchBytes(maxBytes int, opts ...FetchOpt) (MessageBatch, error) {
 	req := &pullRequest{
-		Batch:    1000000,
-		MaxBytes: maxBytes,
-		Expires:  DefaultExpires,
+		Batch:     1000000,
+		MaxBytes:  maxBytes,
+		Expires:   DefaultExpires,
+		Heartbeat: unset,
 	}
 	for _, opt := range opts {
 		if err := opt(req); err != nil {
 			return nil, err
 		}
 	}
-	// for longer pulls, set heartbeat value
-	if req.Expires >= 10*time.Second {
-		req.Heartbeat = 5 * time.Second
+	// if heartbeat was not explicitly set, set it to 5 seconds for longer pulls
+	// and disable it for shorter pulls
+	if req.Heartbeat == unset {
+		if req.Expires >= 10*time.Second {
+			req.Heartbeat = 5 * time.Second
+		} else {
+			req.Heartbeat = 0
+		}
+	}
+	if req.Expires < 2*req.Heartbeat {
+		return nil, fmt.Errorf("%w: expiry time should be at least 2 times the heartbeat", ErrInvalidOption)
 	}
 
 	return p.fetch(req)
 }
 
 // FetchNoWait sends a single request to retrieve given number of messages.
-// If there are any messages available at the time of sending request,
-// FetchNoWait will return immediately.
+// FetchNoWait will only return messages that are available at the time of the
+// request. It will not wait for more messages to arrive.
 func (p *pullConsumer) FetchNoWait(batch int) (MessageBatch, error) {
 	req := &pullRequest{
 		Batch:  batch,
@@ -842,6 +860,10 @@ func (p *pullConsumer) fetch(req *pullRequest) (MessageBatch, error) {
 					return
 				}
 				p.Unlock()
+			case err := <-sub.errs:
+				res.err = err
+				res.done = true
+				return
 			case <-time.After(req.Expires + 1*time.Second):
 				res.done = true
 				return
diff --git a/object.go b/object.go
index 72e1ef4..f3f5d97 100644
--- a/object.go
+++ b/object.go
@@ -153,6 +153,9 @@ type ObjectStoreConfig struct {
 	// Bucket-specific metadata
 	// NOTE: Metadata requires nats-server v2.10.0+
 	Metadata map[string]string `json:"metadata,omitempty"`
+	// Enable underlying stream compression.
+	// NOTE: Compression is supported for nats-server 2.10.0+
+	Compression bool
 }
 
 type ObjectStoreStatus interface {
@@ -174,6 +177,8 @@ type ObjectStoreStatus interface {
 	BackingStore() string
 	// Metadata is the user supplied metadata for the bucket
 	Metadata() map[string]string
+	// IsCompressed indicates if the data is compressed on disk
+	IsCompressed() bool
 }
 
 // ObjectMetaOptions
@@ -266,7 +271,10 @@ func (js *js) CreateObjectStore(cfg *ObjectStoreConfig) (ObjectStore, error) {
 	if maxBytes == 0 {
 		maxBytes = -1
 	}
-
+	var compression StoreCompression
+	if cfg.Compression {
+		compression = S2Compression
+	}
 	scfg := &StreamConfig{
 		Name:        fmt.Sprintf(objNameTmpl, name),
 		Description: cfg.Description,
@@ -280,6 +288,7 @@ func (js *js) CreateObjectStore(cfg *ObjectStoreConfig) (ObjectStore, error) {
 		AllowRollup: true,
 		AllowDirect: true,
 		Metadata:    cfg.Metadata,
+		Compression: compression,
 	}
 
 	// Create our stream.
@@ -1224,6 +1233,9 @@ func (s *ObjectBucketStatus) Metadata() map[string]string { return s.nfo.Config.
 // StreamInfo is the stream info retrieved to create the status
 func (s *ObjectBucketStatus) StreamInfo() *StreamInfo { return s.nfo }
 
+// IsCompressed indicates if the data is compressed on disk
+func (s *ObjectBucketStatus) IsCompressed() bool { return s.nfo.Config.Compression != NoCompression }
+
 // Status retrieves run-time status about a bucket
 func (obs *obs) Status() (ObjectStoreStatus, error) {
 	nfo, err := obs.js.StreamInfo(obs.stream)


Test patch:
diff --git a/jetstream/test/kv_test.go b/jetstream/test/kv_test.go
index 6dda860..f820535 100644
--- a/jetstream/test/kv_test.go
+++ b/jetstream/test/kv_test.go
@@ -1,4 +1,4 @@
-// Copyright 2023 The NATS Authors
+// Copyright 2023-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
@@ -113,6 +113,69 @@ func TestKeyValueBasics(t *testing.T) {
 	}
 }
 
+func TestCreateKeyValue(t *testing.T) {
+	s := RunBasicJetStreamServer()
+	defer shutdownJSServerAndRemoveStorage(t, s)
+
+	nc, js := jsClient(t, s)
+	defer nc.Close()
+	ctx := context.Background()
+
+	// invalid bucket name
+	_, err := js.CreateKeyValue(ctx, jetstream.KeyValueConfig{Bucket: "TEST.", Description: "Test KV"})
+	expectErr(t, err, jetstream.ErrInvalidBucketName)
+
+	_, err = js.CreateKeyValue(ctx, jetstream.KeyValueConfig{Bucket: "TEST", Description: "Test KV"})
+	expectOk(t, err)
+
+	// Check that we can't overwrite existing bucket.
+	_, err = js.CreateKeyValue(ctx, jetstream.KeyValueConfig{Bucket: "TEST", Description: "New KV"})
+	expectErr(t, err, jetstream.ErrBucketExists)
+
+	// assert that we're backwards compatible
+	expectErr(t, err, jetstream.ErrStreamNameAlreadyInUse)
+}
+
+func TestUpdateKeyValue(t *testing.T) {
+	s := RunBasicJetStreamServer()
+	defer shutdownJSServerAndRemoveStorage(t, s)
+
+	nc, js := jsClient(t, s)
+	defer nc.Close()
+	ctx := context.Background()
+
+	// cannot update a non-existing bucket
+	_, err := js.UpdateKeyValue(ctx, jetstream.KeyValueConfig{Bucket: "TEST", Description: "Test KV"})
+	expectErr(t, err, jetstream.ErrBucketNotFound)
+
+	_, err = js.CreateKeyValue(ctx, jetstream.KeyValueConfig{Bucket: "TEST", Description: "Test KV"})
+	expectOk(t, err)
+
+	// update the bucket
+	_, err = js.UpdateKeyValue(ctx, jetstream.KeyValueConfig{Bucket: "TEST", Description: "New KV"})
+	expectOk(t, err)
+}
+
+func TestCreateOrUpdateKeyValue(t *testing.T) {
+	s := RunBasicJetStreamServer()
+	defer shutdownJSServerAndRemoveStorage(t, s)
+
+	nc, js := jsClient(t, s)
+	defer nc.Close()
+	ctx := context.Background()
+
+	// invalid bucket name
+	_, err := js.CreateOrUpdateKeyValue(ctx, jetstream.KeyValueConfig{Bucket: "TEST.", Description: "Test KV"})
+	expectErr(t, err, jetstream.ErrInvalidBucketName)
+
+	_, err = js.CreateOrUpdateKeyValue(ctx, jetstream.KeyValueConfig{Bucket: "TEST", Description: "Test KV"})
+	expectOk(t, err)
+
+	// update the bucket
+	_, err = js.CreateOrUpdateKeyValue(ctx, jetstream.KeyValueConfig{Bucket: "TEST", Description: "New KV"})
+	expectOk(t, err)
+}
+
 func TestKeyValueHistory(t *testing.T) {
 	s := RunBasicJetStreamServer()
 	defer shutdownJSServerAndRemoveStorage(t, s)
@@ -565,8 +628,15 @@ func TestKeyValueDeleteStore(t *testing.T) {
 	err = js.DeleteKeyValue(ctx, "WATCH")
 	expectOk(t, err)
 
+	// delete again should fail
+	err = js.DeleteKeyValue(ctx, "WATCH")
+	expectErr(t, err, jetstream.ErrBucketNotFound)
+
+	// check that we're backwards compatible
+	expectErr(t, err, jetstream.ErrStreamNotFound)
+
 	_, err = js.KeyValue(ctx, "WATCH")
-	expectErr(t, err)
+	expectErr(t, err, jetstream.ErrBucketNotFound)
 }
 
 func TestKeyValueDeleteVsPurge(t *testing.T) {
@@ -1474,7 +1544,7 @@ func expectErr(t *testing.T, err error, expected ...error) {
 		return
 	}
 	for _, e := range expected {
-		if err == e || strings.Contains(e.Error(), err.Error()) {
+		if errors.Is(err, e) {
 			return
 		}
 	}
diff --git a/jetstream/test/object_test.go b/jetstream/test/object_test.go
index 309f462..7b9d85f 100644
--- a/jetstream/test/object_test.go
+++ b/jetstream/test/object_test.go
@@ -1106,3 +1106,38 @@ func TestObjectStoreGetObjectContextTimeout(t *testing.T) {
 	expectErr(t, err, nats.ErrTimeout)
 	r.Close()
 }
+
+func TestObjectStoreCompression(t *testing.T) {
+	s := RunBasicJetStreamServer()
+	defer shutdownJSServerAndRemoveStorage(t, s)
+
+	nc, js := jsClient(t, s)
+	defer nc.Close()
+	ctx := context.Background()
+
+	objCompressed, err := js.CreateObjectStore(ctx, jetstream.ObjectStoreConfig{
+		Bucket:      "A",
+		Compression: true,
+	})
+	if err != nil {
+		t.Fatalf("Error creating object store: %v", err)
+	}
+
+	status, err := objCompressed.Status(ctx)
+	if err != nil {
+		t.Fatalf("Error getting bucket status: %v", err)
+	}
+
+	if !status.IsCompressed() {
+		t.Fatalf("Expected bucket to be compressed")
+	}
+
+	objStream, err := js.Stream(ctx, "OBJ_A")
+	if err != nil {
+		t.Fatalf("Error getting stream info: %v", err)
+	}
+
+	if objStream.CachedInfo().Config.Compression != jetstream.S2Compression {
+		t.Fatalf("Expected stream to be compressed with S2")
+	}
+}
diff --git a/jetstream/test/pull_test.go b/jetstream/test/pull_test.go
index 300d544..b46697e 100644
--- a/jetstream/test/pull_test.go
+++ b/jetstream/test/pull_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
@@ -351,6 +351,130 @@ func TestPullConsumerFetch(t *testing.T) {
 			t.Fatalf("Expected error: %v; got: %v", jetstream.ErrInvalidOption, err)
 		}
 	})
+
+	t.Run("with missing heartbeat", 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)
+		}
+
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		js, err := jetstream.New(nc)
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		defer nc.Close()
+
+		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, nc)
+		// fetch 5 messages, should return normally
+		msgs, err := c.Fetch(5, jetstream.FetchHeartbeat(50*time.Millisecond))
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		var i int
+		for range msgs.Messages() {
+			i++
+		}
+		if i != len(testMsgs) {
+			t.Fatalf("Expected 5 messages; got: %d", i)
+		}
+		if msgs.Error() != nil {
+			t.Fatalf("Unexpected error during fetch: %v", msgs.Error())
+		}
+
+		// fetch again, should timeout without any error
+		msgs, err = c.Fetch(5, jetstream.FetchHeartbeat(50*time.Millisecond), jetstream.FetchMaxWait(200*time.Millisecond))
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		select {
+		case _, ok := <-msgs.Messages():
+			if ok {
+				t.Fatalf("Expected channel to be closed")
+			}
+		case <-time.After(1 * time.Second):
+			t.Fatalf("Expected channel to be closed")
+		}
+		if msgs.Error() != nil {
+			t.Fatalf("Unexpected error during fetch: %v", msgs.Error())
+		}
+
+		// delete the consumer, at this point server should stop sending heartbeats for pull requests
+		if err := s.DeleteConsumer(ctx, c.CachedInfo().Name); err != nil {
+			t.Fatalf("Error deleting consumer: %s", err)
+		}
+		msgs, err = c.Fetch(5, jetstream.FetchHeartbeat(50*time.Millisecond))
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		select {
+		case _, ok := <-msgs.Messages():
+			if ok {
+				t.Fatalf("Expected channel to be closed")
+			}
+		case <-time.After(1 * time.Second):
+			t.Fatalf("Expected channel to be closed")
+		}
+		if !errors.Is(msgs.Error(), jetstream.ErrNoHeartbeat) {
+			t.Fatalf("Expected error: %v; got: %v", jetstream.ErrNoHeartbeat, err)
+		}
+	})
+
+	t.Run("with invalid heartbeat value", 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)
+		}
+
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		js, err := jetstream.New(nc)
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		defer nc.Close()
+
+		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)
+		}
+
+		// default expiry (30s), hb too large
+		_, err = c.Fetch(5, jetstream.FetchHeartbeat(20*time.Second))
+		if !errors.Is(err, jetstream.ErrInvalidOption) {
+			t.Fatalf("Expected error: %v; got: %v", jetstream.ErrInvalidOption, err)
+		}
+
+		// custom expiry, hb too large
+		_, err = c.Fetch(5, jetstream.FetchHeartbeat(2*time.Second), jetstream.FetchMaxWait(3*time.Second))
+		if !errors.Is(err, jetstream.ErrInvalidOption) {
+			t.Fatalf("Expected error: %v; got: %v", jetstream.ErrInvalidOption, err)
+		}
+
+		// negative heartbeat
+		_, err = c.Fetch(5, jetstream.FetchHeartbeat(-2*time.Second))
+		if !errors.Is(err, jetstream.ErrInvalidOption) {
+			t.Fatalf("Expected error: %v; got: %v", jetstream.ErrInvalidOption, err)
+		}
+	})
 }
 
 func TestPullConsumerFetchBytes(t *testing.T) {
@@ -542,6 +666,114 @@ func TestPullConsumerFetchBytes(t *testing.T) {
 			t.Fatalf("Unexpected error during fetch: %v", msgs.Error())
 		}
 	})
+
+	t.Run("with missing heartbeat", 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)
+		}
+
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		js, err := jetstream.New(nc)
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		defer nc.Close()
+
+		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)
+		}
+
+		// fetch again, should timeout without any error
+		msgs, err := c.FetchBytes(5, jetstream.FetchHeartbeat(50*time.Millisecond), jetstream.FetchMaxWait(200*time.Millisecond))
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		select {
+		case _, ok := <-msgs.Messages():
+			if ok {
+				t.Fatalf("Expected channel to be closed")
+			}
+		case <-time.After(1 * time.Second):
+			t.Fatalf("Expected channel to be closed")
+		}
+		if msgs.Error() != nil {
+			t.Fatalf("Unexpected error during fetch: %v", msgs.Error())
+		}
+
+		// delete the consumer, at this point server should stop sending heartbeats for pull requests
+		if err := s.DeleteConsumer(ctx, c.CachedInfo().Name); err != nil {
+			t.Fatalf("Error deleting consumer: %s", err)
+		}
+		msgs, err = c.FetchBytes(5, jetstream.FetchHeartbeat(50*time.Millisecond))
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		select {
+		case _, ok := <-msgs.Messages():
+			if ok {
+				t.Fatalf("Expected channel to be closed")
+			}
+		case <-time.After(1 * time.Second):
+			t.Fatalf("Expected channel to be closed")
+		}
+		if !errors.Is(msgs.Error(), jetstream.ErrNoHeartbeat) {
+			t.Fatalf("Expected error: %v; got: %v", jetstream.ErrNoHeartbeat, err)
+		}
+	})
+
+	t.Run("with invalid heartbeat value", 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)
+		}
+
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		js, err := jetstream.New(nc)
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		defer nc.Close()
+
+		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)
+		}
+
+		// default expiry (30s), hb too large
+		_, err = c.FetchBytes(5, jetstream.FetchHeartbeat(20*time.Second))
+		if !errors.Is(err, jetstream.ErrInvalidOption) {
+			t.Fatalf("Expected error: %v; got: %v", jetstream.ErrInvalidOption, err)
+		}
+
+		// custom expiry, hb too large
+		_, err = c.FetchBytes(5, jetstream.FetchHeartbeat(2*time.Second), jetstream.FetchMaxWait(3*time.Second))
+		if !errors.Is(err, jetstream.ErrInvalidOption) {
+			t.Fatalf("Expected error: %v; got: %v", jetstream.ErrInvalidOption, err)
+		}
+
+		// negative heartbeat
+		_, err = c.FetchBytes(5, jetstream.FetchHeartbeat(-2*time.Second))
+		if !errors.Is(err, jetstream.ErrInvalidOption) {
+			t.Fatalf("Expected error: %v; got: %v", jetstream.ErrInvalidOption, err)
+		}
+	})
 }
 
 func TestPullConsumerFetch_WithCluster(t *testing.T) {
diff --git a/test/object_test.go b/test/object_test.go
index f6a1fc1..616cc1d 100644
--- a/test/object_test.go
+++ b/test/object_test.go
@@ -1087,3 +1087,37 @@ func TestObjectStoreGetObjectContextTimeout(t *testing.T) {
 	expectErr(t, err, nats.ErrTimeout)
 	r.Close()
 }
+
+func TestObjectStoreCompression(t *testing.T) {
+	s := RunBasicJetStreamServer()
+	defer shutdownJSServerAndRemoveStorage(t, s)
+
+	nc, js := jsClient(t, s)
+	defer nc.Close()
+
+	obj, err := js.CreateObjectStore(&nats.ObjectStoreConfig{
+		Bucket:      "A",
+		Compression: true,
+	})
+	if err != nil {
+		t.Fatalf("Error creating object store: %v", err)
+	}
+
+	status, err := obj.Status()
+	if err != nil {
+		t.Fatalf("Error getting bucket status: %v", err)
+	}
+
+	if !status.IsCompressed() {
+		t.Fatalf("Expected bucket to be compressed")
+	}
+
+	objStream, err := js.StreamInfo("OBJ_A")
+	if err != nil {
+		t.Fatalf("Error getting stream info: %v", err)
+	}
+
+	if objStream.Config.Compression != nats.S2Compression {
+		t.Fatalf("Expected stream to be compressed with S2")
+	}
+}