storage: fix concurrent map read/write

.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 
 

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 {

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

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"))
+		}
+	})
+
+}