From a04dab5b2a3eea0e29da8eea3d4e9c3097ddc2a0 Mon Sep 17 00:00:00 2001
From: rsteube <rsteube@users.noreply.github.com>
Date: Wed, 6 Sep 2023 09:53:56 +0200
Subject: [PATCH] storage: fix concurrent map read/write

---
 .github/workflows/go.yml |  3 +++
 carapace.go              |  6 +++++-
 storage.go               | 15 +++++++++------
 storage_test.go          | 20 ++++++++++++++++++++
 4 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml
index 2a7f1426..e81f84b4 100644
--- a/.github/workflows/go.yml
+++ b/.github/workflows/go.yml
@@ -31,6 +31,9 @@ jobs:
       - name: Test
         run: mkdir .cover && CARAPACE_COVERDIR="$(pwd)/.cover" go test -v -coverpkg ./... -coverprofile=unit.cov ./... ./example-nonposix/...
 
+      - name: Bench
+        run: go test -bench ./...
+
       - name: Convert coverage
         run: go tool covdata textfmt -i .cover/ -o integration.cov 
 
diff --git a/carapace.go b/carapace.go
index 8e81670a..714139f6 100644
--- a/carapace.go
+++ b/carapace.go
@@ -72,7 +72,11 @@ func (c Carapace) DashAnyCompletion(action Action) {
 
 // FlagCompletion defines completion for flags using a map consisting of name and Action.
 func (c Carapace) FlagCompletion(actions ActionMap) {
-	if e := storage.get(c.cmd); e.flag == nil {
+	e := storage.get(c.cmd)
+	e.flagMutex.Lock()
+	defer e.flagMutex.Unlock()
+
+	if e.flag == nil {
 		e.flag = actions
 	} else {
 		for name, action := range actions {
diff --git a/storage.go b/storage.go
index c46fa1a7..fe4b8c22 100644
--- a/storage.go
+++ b/storage.go
@@ -15,6 +15,7 @@ import (
 
 type entry struct {
 	flag          ActionMap
+	flagMutex     sync.RWMutex
 	positional    []Action
 	positionalAny Action
 	dash          []Action
@@ -27,20 +28,22 @@ type entry struct {
 
 type _storage map[*cobra.Command]*entry
 
-var storageMutex sync.Mutex
+var storageMutex sync.RWMutex
 
-func (s _storage) get(cmd *cobra.Command) (e *entry) {
-	var ok bool
-	if e, ok = s[cmd]; !ok {
+func (s _storage) get(cmd *cobra.Command) *entry {
+	storageMutex.RLock()
+	e, ok := s[cmd]
+	storageMutex.RUnlock()
+
+	if !ok {
 		storageMutex.Lock()
 		defer storageMutex.Unlock()
-
 		if e, ok = s[cmd]; !ok {
 			e = &entry{}
 			s[cmd] = e
 		}
 	}
-	return
+	return e
 }
 
 var bridgeMutex sync.Mutex
diff --git a/storage_test.go b/storage_test.go
index 46486331..729baf3b 100644
--- a/storage_test.go
+++ b/storage_test.go
@@ -75,3 +75,23 @@ func TestCheck(t *testing.T) {
 		t.Error("check should fail")
 	}
 }
+
+// BenchmarkStorage tests for concurrent map read/write
+func BenchmarkStorage(b *testing.B) {
+	cmd := &cobra.Command{}
+	cmd2 := &cobra.Command{}
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			Gen(cmd).FlagCompletion(ActionMap{
+				"flag1": ActionValues("a", "b"),
+			})
+			Gen(cmd).PositionalCompletion(ActionValues("a", "b"))
+
+			Gen(cmd2).FlagCompletion(ActionMap{
+				"flag2": ActionValues("a", "b"),
+			})
+			Gen(cmd2).PositionalCompletion(ActionValues("a", "b"))
+		}
+	})
+
+}