From e810946c7096b669195abd9f0ec280168f7385ea Mon Sep 17 00:00:00 2001
From: Mustafa Gezen <mustafa@ctrliq.com>
Date: Tue, 16 Aug 2022 15:45:07 +0200
Subject: [PATCH] Switch to sync.Map for keywarming

Signed-off-by: Mustafa Gezen <mustafa@ctrliq.com>
---
 peridot/keykeeper/v1/keywarming.go | 28 ++++++----------
 peridot/keykeeper/v1/server.go     | 54 ++++++------------------------
 2 files changed, 22 insertions(+), 60 deletions(-)

diff --git a/peridot/keykeeper/v1/keywarming.go b/peridot/keykeeper/v1/keywarming.go
index f95ce133..5ff67e64 100644
--- a/peridot/keykeeper/v1/keywarming.go
+++ b/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{
-			keyUuid: db.ID,
-			gpgId:   gpgKey.GetHexKeyID(),
-		}
+	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
diff --git a/peridot/keykeeper/v1/server.go b/peridot/keykeeper/v1/server.go
index c2ba2c6e..9d14225b 100644
--- a/peridot/keykeeper/v1/server.go
+++ b/peridot/keykeeper/v1/server.go
@@ -57,45 +57,17 @@ 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
 
-	log           *logrus.Logger
-	db            peridotdb.Access
-	storage       lookaside.Storage
-	worker        worker.Worker
-	temporal      client.Client
-	stores        map[string]store.Store
-	keys          map[string]*LoadedKey
-	keyImportLock *MapStringLock
-	defaultStore  string
+	log          *logrus.Logger
+	db           peridotdb.Access
+	storage      lookaside.Storage
+	worker       worker.Worker
+	temporal     client.Client
+	stores       map[string]store.Store
+	keys         *sync.Map
+	defaultStore string
 }
 
 func NewServer(db peridotdb.Access, c client.Client) (*Server, error) {
@@ -116,13 +88,9 @@ func NewServer(db peridotdb.Access, c client.Client) (*Server, error) {
 		worker: worker.New(c, TaskQueue, worker.Options{
 			DeadlockDetectionTimeout: 15 * time.Minute,
 		}),
-		temporal: c,
-		stores:   map[string]store.Store{"awssm": sm},
-		keys:     map[string]*LoadedKey{},
-		keyImportLock: &MapStringLock{
-			RWMutex: &sync.RWMutex{},
-			m:       map[string]*sync.Mutex{},
-		},
+		temporal:     c,
+		stores:       map[string]store.Store{"awssm": sm},
+		keys:         &sync.Map{},
 		defaultStore: "awssm",
 	}, nil
 }