Switch to sync.Map for keywarming

Signed-off-by: Mustafa Gezen <mustafa@ctrliq.com>

peridot/keykeeper/v1/keywarming.go

@@ -38,8 +38,6 @@ import (
 	"fmt"
 	"github.com/ProtonMail/gopenpgp/v2/crypto"
 	"github.com/google/uuid"
-	"io/ioutil"
-	"os"
 	"os/exec"
 	"peridot.resf.org/peridot/db/models"
 	"peridot.resf.org/utils"
@@ -88,13 +86,10 @@ func (s *Server) importGpgKey(armoredKey string) error {
 // WarmGPGKey warms up a specific GPG key
 // This involves shelling out to GPG to import the key
 func (s *Server) WarmGPGKey(key string, armoredKey string, gpgKey *crypto.Key, db *models.Key) (*LoadedKey, error) {
-	s.keyImportLock.ReadLock(key)
-	defer s.keyImportLock.ReadUnlock(key)
-
-	cachedKey := s.keys[key]
+	cachedKeyAny, ok := s.keys.Load(key)
 	// This means that the key is already loaded
-	if cachedKey != nil {
-		return cachedKey, nil
+	if ok {
+		return cachedKeyAny.(*LoadedKey), nil
 	}
 
 	err := s.importGpgKey(armoredKey)
@@ -102,21 +97,20 @@ func (s *Server) WarmGPGKey(key string, armoredKey string, gpgKey *crypto.Key, d
 		return nil, err
 	}
 
-	if cachedKey == nil {
-		s.keys[key] = &LoadedKey{
+	cachedKey := &LoadedKey{
 		keyUuid: db.ID,
 		gpgId:   gpgKey.GetHexKeyID(),
 	}
-	}
+	s.keys.Store(key, cachedKey)
 
-	return s.keys[key], nil
+	return cachedKey, nil
 }
 
 // EnsureGPGKey ensures that the key is loaded
 func (s *Server) EnsureGPGKey(key string) (*LoadedKey, error) {
-	cachedKey := s.keys[key]
-	if cachedKey != nil {
-		return cachedKey, nil
+	cachedKeyAny, ok := s.keys.Load(key)
+	if ok {
+		return cachedKeyAny.(*LoadedKey), nil
 	}
 
 	// Key not found in cache, fetch from database

peridot/keykeeper/v1/server.go

@@ -57,33 +57,6 @@ import (
 
 const TaskQueue = "keykeeper"
 
-type MapStringLock struct {
-	*sync.RWMutex
-	m map[string]*sync.Mutex
-}
-
-func (m *MapStringLock) ReadLock(key string) {
-	m.RLock()
-	defer m.RUnlock()
-	if m.m[key] == nil {
-		m.Lock()
-		m.m[key] = &sync.Mutex{}
-		m.Unlock()
-	}
-	m.m[key].Lock()
-}
-
-func (m *MapStringLock) ReadUnlock(key string) {
-	m.RLock()
-	defer m.RUnlock()
-	if m.m[key] == nil {
-		m.Lock()
-		m.m[key] = &sync.Mutex{}
-		m.Unlock()
-	}
-	m.m[key].Unlock()
-}
-
 type Server struct {
 	keykeeperpb.UnimplementedKeykeeperServiceServer
 
@@ -93,8 +66,7 @@ type Server struct {
 	worker       worker.Worker
 	temporal     client.Client
 	stores       map[string]store.Store
-	keys          map[string]*LoadedKey
-	keyImportLock *MapStringLock
+	keys         *sync.Map
 	defaultStore string
 }
 
@@ -118,11 +90,7 @@ func NewServer(db peridotdb.Access, c client.Client) (*Server, error) {
 		}),
 		temporal:     c,
 		stores:       map[string]store.Store{"awssm": sm},
-		keys:     map[string]*LoadedKey{},
-		keyImportLock: &MapStringLock{
-			RWMutex: &sync.RWMutex{},
-			m:       map[string]*sync.Mutex{},
-		},
+		keys:         &sync.Map{},
 		defaultStore: "awssm",
 	}, nil
 }