use a sync.Map for storing buckets

Signed-off-by: jkoberg <[email protected]>
micro · Oct 20, 2023 · 9463c82 · 9463c82
1 parent 3882f06
commit 9463c82
Showing 1 changed file with 13 additions and 10 deletions.
diff --git a/v4/store/nats-js/nats.go b/v4/store/nats-js/nats.go
@@ -34,7 +34,7 @@ type natsStore struct {
 
 	conn    *nats.Conn
 	js      nats.JetStreamContext
-	buckets map[string]nats.ObjectStore
+	buckets *sync.Map
 }
 
 func init() {
@@ -55,7 +55,7 @@ func NewStore(opts ...store.Option) store.Store {
 		opts:            options,
 		jsopts:          []nats.JSOpt{},
 		objStoreConfigs: []*nats.ObjectStoreConfig{},
-		buckets:         map[string]nats.ObjectStore{},
+		buckets:         &sync.Map{},
 		storageType:     nats.FileStorage,
 	}
 
@@ -103,7 +103,7 @@ func (n *natsStore) Init(opts ...store.Option) error {
 		if err != nil {
 			return errors.Wrapf(err, "Failed to create bucket (%s)", cfg.Bucket)
 		}
-		n.buckets[cfg.Bucket] = store
+		n.buckets.Store(cfg.Bucket, store)
 	}
 
 	return nil
@@ -178,10 +178,11 @@ func (n *natsStore) Read(key string, opts ...store.ReadOption) ([]*store.Record,
 		opt.Table = n.opts.Table
 	}
 
-	bucket, ok := n.buckets[opt.Database]
+	b, ok := n.buckets.Load(opt.Database)
 	if !ok {
 		return nil, ErrBucketNotFound
 	}
+	bucket := b.(nats.ObjectStore)
 
 	var keys []string
 
@@ -291,8 +292,8 @@ func (n *natsStore) Write(r *store.Record, opts ...store.WriteOption) error {
 		opt.Table = n.opts.Table
 	}
 
-	store, ok := n.buckets[opt.Database]
-
+	s, ok := n.buckets.Load(opt.Database)
+	store, _ := s.(nats.ObjectStore)
 	// Create new bucket if not exists
 	if !ok {
 		var err error
@@ -344,18 +345,19 @@ func (n *natsStore) Delete(key string, opts ...store.DeleteOption) error {
 	}
 
 	if opt.Table == "DELETE_BUCKET" {
-		delete(n.buckets, key)
+		n.buckets.Delete(key)
 		if err := n.js.DeleteObjectStore(key); err != nil {
 			return errors.Wrap(err, "Failed to delete bucket")
 		}
 		return nil
 	}
 
-	store, ok := n.buckets[opt.Database]
+	s, ok := n.buckets.Load(opt.Database)
 	if !ok {
 		return ErrBucketNotFound
 	}
 
+	store := s.(nats.ObjectStore)
 	if err := store.Delete(getKey(key, opt.Table)); err != nil {
 		return errors.Wrap(err, "Failed to delete data")
 	}
@@ -381,10 +383,11 @@ func (n *natsStore) List(opts ...store.ListOption) ([]string, error) {
 		opt.Table = n.opts.Table
 	}
 
-	store, ok := n.buckets[opt.Database]
+	s, ok := n.buckets.Load(opt.Database)
 	if !ok {
 		return nil, ErrBucketNotFound
 	}
+	store := s.(nats.ObjectStore)
 
 	objects, err := store.List()
 	if err != nil {
@@ -445,7 +448,7 @@ func (n *natsStore) createNewBucket(name string) (nats.ObjectStore, error) {
 	if err != nil {
 		return nil, errors.Wrapf(err, "Failed to create new bucket (%s)", name)
 	}
-	n.buckets[name] = store
+	n.buckets.Store(name, store)
 	return store, err
 }