fix(modelchecker): make protopath cache safe for concurrent readers

modelchecker/protopath.go

@@ -8,29 +8,50 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"sync"
 )
 
 var re = regexp.MustCompile(`Stmts\[\d+\]`)
 
 type ProtoPath struct {
-	// TODO(jayaprabhakar): A quick hack, fix this. It is safe because this field is immutable.
+	mu       sync.RWMutex
 	filesMap map[*ast.File]map[string]proto.Message
 }
+
 var protoPathInstance = &ProtoPath{filesMap: make(map[*ast.File]map[string]proto.Message)}
 
 func GetProtoFieldByPath(file *ast.File, location string) proto.Message {
-	if protoPathInstance.filesMap[file] == nil {
-		protoPathInstance.filesMap[file] = make(map[string]proto.Message)
-	} else if val, ok := protoPathInstance.filesMap[file][location]; ok {
-		return val
+	// Fast path: read-locked cache lookup. Parallel callers (e.g. parallel
+	// simulation workers) hit this most of the time once the cache is warm.
+	protoPathInstance.mu.RLock()
+	if inner, ok := protoPathInstance.filesMap[file]; ok {
+		if val, hit := inner[location]; hit {
+			protoPathInstance.mu.RUnlock()
+			return val
+		}
 	}
+	protoPathInstance.mu.RUnlock()
+
+	// Slow path: compute outside the lock (reflection on the read-only AST
+	// is safe without the lock), then insert under a write lock with a
+	// re-check in case a concurrent caller already populated the entry.
 	field := GetFieldByPath(file, location)
-	if field == nil {
-		protoPathInstance.filesMap[file][location] = nil
-		return nil
+	var protobuf proto.Message
+	if field != nil {
+		protobuf = convertToProto(field.Elem().Interface(), field.Type())
+	}
+
+	protoPathInstance.mu.Lock()
+	defer protoPathInstance.mu.Unlock()
+	inner, ok := protoPathInstance.filesMap[file]
+	if !ok {
+		inner = make(map[string]proto.Message)
+		protoPathInstance.filesMap[file] = inner
+	}
+	if existing, hit := inner[location]; hit {
+		return existing
 	}
-	protobuf := convertToProto(field.Elem().Interface(), field.Type())
-	protoPathInstance.filesMap[file][location] = protobuf
+	inner[location] = protobuf
 	return protobuf
 }
 

modelchecker/protopath_test.go

@@ -5,6 +5,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/protobuf/encoding/protojson"
+	"sync"
 	"testing"
 )
 
@@ -77,6 +78,45 @@ func TestEndOfBlock(t *testing.T) {
 		EndOfBlock("Actions[0].Block.Stmts[0].AnyStmt.Block.Stmts[0].IfStmt.Branches[0].Block.Stmts[0]"))
 }
 
+// TestGetProtoFieldByPath_Concurrent exercises the cache under parallel reads
+// + writes. Run with `go test -race` to catch any locking regression. The
+// shared protoPathInstance cache is hit by every Processor; parallel
+// simulation workers must not race on it.
+func TestGetProtoFieldByPath_Concurrent(t *testing.T) {
+	file, err := readFileToAst()
+	require.Nil(t, err)
+
+	paths := []string{
+		"Actions[0]",
+		"Actions[0].Block",
+		"Actions[0].Block.Stmts[0]",
+		"Actions[0].Block.Stmts[0].AnyStmt",
+		"Actions[0].Block.Stmts[0].AnyStmt.Block",
+		"Actions[1]",
+		"Actions[1].Block",
+		"Actions[1].Block.Stmts[0]",
+		"Actions[1].Block.Stmts[0].AnyStmt",
+		"NonExistentPath[42]", // exercises the nil-result cache entry
+	}
+
+	const workers = 16
+	const iterations = 200
+
+	var wg sync.WaitGroup
+	wg.Add(workers)
+	for w := 0; w < workers; w++ {
+		go func() {
+			defer wg.Done()
+			for i := 0; i < iterations; i++ {
+				for _, p := range paths {
+					_ = GetProtoFieldByPath(file, p)
+				}
+			}
+		}()
+	}
+	wg.Wait()
+}
+
 func readFileToAst() (*ast.File, error) {
 	jsonFile := `
 {