Repository: nats-io/nats.go (go)

Issue:
Add stream compression option for object store buckets
### Proposed change

An S2 compression option is available for Jetstream streams and KV buckets. However, it is still missing for the object store. We should add it.

### Use case

Storing compressible content in the object store.

### Contribution

The change seems very simple as  #1451 can be used as blueprint. I can do that.

Pull request description:
 compression option for object store
This PR resolves #1536. It is mainly a copy of how the same was realized for KV.

I would be very happy to see this merged. I am about to use the NATS object store for a large amount of textual data. It seems, S2 compression saves about 50 % of disk storage.

Patch:
diff --git a/encoders/builtin/default_enc.go b/encoders/builtin/default_enc.go
index 7e72963..c1d0f6f 100644
--- a/encoders/builtin/default_enc.go
+++ b/encoders/builtin/default_enc.go
@@ -24,7 +24,7 @@ import (
 // DefaultEncoder implementation for EncodedConn.
 // This encoder will leave []byte and string untouched, but will attempt to
 // turn numbers into appropriate strings that can be decoded. It will also
-// propely encoded and decode bools. If will encode a struct, but if you want
+// properly encoded and decode bools. If will encode a struct, but if you want
 // to properly handle structures you should use JsonEncoder.
 type DefaultEncoder struct {
 	// Empty
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/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/stream_config.go b/jetstream/stream_config.go
index e0133d8..dd1f9d9 100644
--- a/jetstream/stream_config.go
+++ b/jetstream/stream_config.go
@@ -123,7 +123,7 @@ type (
 		//
 		// If set to true, publish methods from the JetStream client will not
 		// work as expected, since they rely on acknowledgements. Core NATS
-		// publish methods should be used instead. Note tha this will make
+		// publish methods should be used instead. Note that this will make
 		// message delivery less reliable.
 		NoAck bool `json:"no_ack,omitempty"`
 
diff --git a/object.go b/object.go
index 0c6647e..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.
@@ -384,7 +393,7 @@ func (obs *obs) Put(meta *ObjectMeta, r io.Reader, opts ...ObjectOpt) (*ObjectIn
 		case <-time.After(obs.js.opts.wait):
 		}
 		if err := obs.js.purgeStream(obs.stream, &StreamPurgeRequest{Subject: chunkSubj}); err != nil {
-			return fmt.Errorf("could not cleanup bucket after erronous put operation: %w", err)
+			return fmt.Errorf("could not cleanup bucket after erroneous put operation: %w", err)
 		}
 		return nil
 	}
@@ -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/example_test.go b/example_test.go
index 26327f3..6aa9363 100644
--- a/example_test.go
+++ b/example_test.go
@@ -612,7 +612,7 @@ func ExampleSubOpt() {
 		nats.MaxDeliver(1),
 		// Activate Flow control algorithm from the server.
 		nats.EnableFlowControl(),
-		// Track heartbeats from the server fro missed sequences.
+		// Track heartbeats from the server for missed sequences.
 		nats.IdleHeartbeat(500*time.Millisecond),
 	)
 
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/micro/test/service_test.go b/micro/test/service_test.go
index 30cfa45..b9e0049 100644
--- a/micro/test/service_test.go
+++ b/micro/test/service_test.go
@@ -181,7 +181,7 @@ func TestServiceBasics(t *testing.T) {
 
 	// Services should process 50 requests total
 	if requestsNum != 50 {
-		t.Fatalf("Expected a total fo 50 requests processed, got: %d", requestsNum)
+		t.Fatalf("Expected a total of 50 requests processed, got: %d", requestsNum)
 	}
 	// Reset stats for a service
 	svcs[0].Reset()
diff --git a/test/conn_test.go b/test/conn_test.go
index 588acc2..e621c90 100644
--- a/test/conn_test.go
+++ b/test/conn_test.go
@@ -2890,7 +2890,7 @@ func TestTLSHandshakeFirst(t *testing.T) {
 	tlsConf.ServerName = "localhost"
 
 	// Start a mockup server that will do the TLS handshake first
-	// and then send the INFO protcol.
+	// and then send the INFO protocol.
 	l, e := net.Listen("tcp", ":0")
 	if e != nil {
 		t.Fatal("Could not listen on an ephemeral port")
diff --git a/test/js_test.go b/test/js_test.go
index 631eda6..4ae2915 100644
--- a/test/js_test.go
+++ b/test/js_test.go
@@ -8622,7 +8622,7 @@ func TestJetStreamMsgSubjectRewrite(t *testing.T) {
 		t.Fatalf("Subject should be %q, got %q", "foo", msg.Subject)
 	}
 	if string(msg.Data) != "msg" {
-		t.Fatalf("Unexepcted data: %q", msg.Data)
+		t.Fatalf("Unexpected data: %q", msg.Data)
 	}
 }
 
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")
+	}
+}