From 3f06fecca3a2c90ee9591ad12ecdee84e69b0ffc Mon Sep 17 00:00:00 2001
From: Anatolie Lupacescu <anatolie@ssvlabs.io>
Date: Thu, 10 Oct 2024 12:31:51 +0300
Subject: [PATCH] syncmap fixes

---
 registry/storage/shares.go | 45 ++++++++++++++++++--------------------
 1 file changed, 21 insertions(+), 24 deletions(-)

diff --git a/registry/storage/shares.go b/registry/storage/shares.go
index cac12d02fd..d031a0d1bb 100644
--- a/registry/storage/shares.go
+++ b/registry/storage/shares.go
@@ -56,7 +56,7 @@ type sharesStorage struct {
 	prefix         []byte
 	shares         *sync.Map
 	validatorStore *validatorStore
-	dbmu           sync.Mutex
+	mu             sync.Mutex
 }
 
 type storageOperator struct {
@@ -132,8 +132,8 @@ func NewSharesStorage(logger *zap.Logger, db basedb.Database, prefix []byte) (Sh
 
 // load reads all shares from db.
 func (s *sharesStorage) load() error {
-	s.dbmu.Lock()
-	defer s.dbmu.Unlock()
+	s.mu.Lock()
+	defer s.mu.Unlock()
 
 	return s.db.GetAll(append(s.prefix, sharesPrefix...), func(i int, obj basedb.Obj) error {
 		val := &storageShare{}
@@ -207,19 +207,17 @@ func (s *sharesStorage) Save(rw basedb.ReadWriter, shares ...*types.SSVShare) er
 		}
 	}
 
-	err := func() error {
-		s.dbmu.Lock()
-		defer s.dbmu.Unlock()
+	s.mu.Lock()
+	defer s.mu.Unlock()
 
-		return s.db.Using(rw).SetMany(s.prefix, len(shares), func(i int) (basedb.Obj, error) {
-			share := specShareToStorageShare(shares[i])
-			value, err := share.Encode()
-			if err != nil {
-				return basedb.Obj{}, fmt.Errorf("failed to serialize share: %w", err)
-			}
-			return basedb.Obj{Key: s.storageKey(share.ValidatorPubKey[:]), Value: value}, nil
-		})
-	}()
+	err := s.db.SetMany(s.prefix, len(shares), func(i int) (basedb.Obj, error) {
+		share := specShareToStorageShare(shares[i])
+		value, err := share.Encode()
+		if err != nil {
+			return basedb.Obj{}, fmt.Errorf("failed to serialize share: %w", err)
+		}
+		return basedb.Obj{Key: s.storageKey(share.ValidatorPubKey[:]), Value: value}, nil
+	})
 	if err != nil {
 		return err
 	}
@@ -312,8 +310,8 @@ func (s *sharesStorage) storageShareToSpecShare(share *storageShare) (*types.SSV
 }
 
 func (s *sharesStorage) Delete(rw basedb.ReadWriter, pubKey []byte) error {
-	s.dbmu.Lock()
-	defer s.dbmu.Unlock()
+	s.mu.Lock()
+	defer s.mu.Unlock()
 
 	// Delete the share from the database
 	if err := s.db.Using(rw).Delete(s.prefix, s.storageKey(pubKey)); err != nil {
@@ -321,16 +319,15 @@ func (s *sharesStorage) Delete(rw basedb.ReadWriter, pubKey []byte) error {
 	}
 
 	key := hex.EncodeToString(pubKey)
-	val, found := s.shares.Load(key)
+
+	// Remove the share from local storage map
+	val, found := s.shares.LoadAndDelete(key)
 	if !found {
 		return nil
 	}
 
 	share := val.(*types.SSVShare)
 
-	// Remove the share from local storage map
-	s.shares.Delete(key)
-
 	// Remove the share from the validator store. This method will handle its own locking.
 	if err := s.validatorStore.handleShareRemoved(share); err != nil {
 		return err
@@ -365,8 +362,8 @@ func (s *sharesStorage) UpdateValidatorsMetadata(data map[spectypes.ValidatorPK]
 
 // Drop deletes all shares.
 func (s *sharesStorage) Drop() error {
-	s.dbmu.Lock()
-	defer s.dbmu.Unlock()
+	s.mu.Lock()
+	defer s.mu.Unlock()
 
 	err := s.db.DropPrefix(bytes.Join(
 		[][]byte{s.prefix, sharesPrefix, []byte("/")},
@@ -376,7 +373,7 @@ func (s *sharesStorage) Drop() error {
 		return err
 	}
 
-	s.shares = new(sync.Map) // make(map[string]*types.SSVShare)
+	s.shares = new(sync.Map)
 	s.validatorStore.handleDrop()
 	return nil
 }