From 8cc9037a597e854a4c29ea9420b89d927a4b995d Mon Sep 17 00:00:00 2001
From: Thomas Rooney <thomas@speakeasyapi.dev>
Date: Sun, 6 Oct 2024 21:06:57 +0100
Subject: [PATCH] chore: fix race condition

---
 index/extract_refs.go        |  9 ++++
 index/search_rolodex_test.go | 82 +++++++++++++++++++++---------------
 2 files changed, 56 insertions(+), 35 deletions(-)

diff --git a/index/extract_refs.go b/index/extract_refs.go
index b49c9a97..930c113a 100644
--- a/index/extract_refs.go
+++ b/index/extract_refs.go
@@ -640,7 +640,16 @@ func (index *SpecIndex) ExtractComponentsFromRefs(refs []*Reference) []*Referenc
 			index.refLock.Unlock()
 		} else {
 			index.refLock.Unlock()
+			// If it's local, this is safe to do unlocked
+			uri := strings.Split(ref.FullDefinition, "#/")
+			unsafeAsync := len(uri) == 2 && len(uri[0]) > 0
+			if unsafeAsync {
+				index.refLock.Lock()
+			}
 			located := index.FindComponent(ref.FullDefinition)
+			if unsafeAsync {
+				index.refLock.Unlock()
+			}
 			if located != nil {
 
 				// have we already mapped this?
diff --git a/index/search_rolodex_test.go b/index/search_rolodex_test.go
index f7d7898a..2f303bad 100644
--- a/index/search_rolodex_test.go
+++ b/index/search_rolodex_test.go
@@ -5,10 +5,12 @@ package index
 
 import (
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"github.com/vmware-labs/yaml-jsonpath/pkg/yamlpath"
 	"gopkg.in/yaml.v3"
 	"path/filepath"
 	"strings"
+	"sync"
 	"testing"
 )
 
@@ -129,52 +131,62 @@ func TestRolodex_FindNodeOrigin_ModifyLookup(t *testing.T) {
 }
 
 func TestSpecIndex_TestPathsAsRefWithFiles(t *testing.T) {
-
-	yml := `paths:
+	// We're TDD'ing some code that previously had a race condition.
+	// This test is to ensure that we don't regress.
+	wg := sync.WaitGroup{}
+	wg.Add(1000)
+	for i := 0; i < 1000; i++ {
+		go func(i int) {
+			defer wg.Done()
+			yml := `paths:
   /test:
     $ref: 'rolodex_test_data/paths/paths.yaml#/~1some~1path'
   /test-2:
     $ref: './rolodex_test_data/paths/paths.yaml#/~1some~1path'
 `
 
-	baseDir := "."
+			baseDir := "."
 
-	cf := CreateOpenAPIIndexConfig()
-	cf.BasePath = baseDir
-	cf.AvoidCircularReferenceCheck = true
+			cf := CreateOpenAPIIndexConfig()
+			cf.BasePath = baseDir
+			cf.AvoidCircularReferenceCheck = true
 
-	fileFS, err := NewLocalFSWithConfig(&LocalFSConfig{
-		BaseDirectory: baseDir,
-		IndexConfig:   cf,
-	})
-	if err != nil {
-		t.Fatal(err)
-	}
+			fileFS, err := NewLocalFSWithConfig(&LocalFSConfig{
+				BaseDirectory: baseDir,
+				IndexConfig:   cf,
+			})
+			if err != nil {
+				t.Fatal(err)
+			}
 
-	rolo := NewRolodex(cf)
-	rolo.AddLocalFS(baseDir, fileFS)
+			rolo := NewRolodex(cf)
+			rolo.AddLocalFS(baseDir, fileFS)
 
-	var rootNode yaml.Node
-	_ = yaml.Unmarshal([]byte(yml), &rootNode)
+			var rootNode yaml.Node
+			_ = yaml.Unmarshal([]byte(yml), &rootNode)
 
-	rolo.SetRootNode(&rootNode)
+			rolo.SetRootNode(&rootNode)
 
-	err = rolo.IndexTheRolodex()
-	assert.NoError(t, err)
-	rolo.Resolve()
+			err = rolo.IndexTheRolodex()
+			assert.NoError(t, err)
+			require.Len(t, rolo.indexes, 2)
+
+			rolo.Resolve()
 
-	assert.Len(t, rolo.indexes, 2)
-	assert.Len(t, rolo.GetCaughtErrors(), 0)
-
-	params := rolo.rootIndex.GetAllParametersFromOperations()
-	assert.Len(t, params, 2)
-	lookupPath, ok := params["/test"]
-	assert.True(t, ok)
-	lookupOperation, ok := lookupPath["get"]
-	assert.True(t, ok)
-	assert.Len(t, lookupOperation, 1)
-	lookupRef, ok := lookupOperation["../components.yaml#/components/parameters/SomeParam"]
-	assert.True(t, ok)
-	assert.Len(t, lookupRef, 1)
-	assert.Equal(t, lookupRef[0].Name, "SomeParam")
+			require.Len(t, rolo.GetCaughtErrors(), 0)
+
+			params := rolo.rootIndex.GetAllParametersFromOperations()
+			require.Len(t, params, 2)
+			lookupPath, ok := params["/test"]
+			require.True(t, ok)
+			lookupOperation, ok := lookupPath["get"]
+			require.True(t, ok)
+			require.Len(t, lookupOperation, 1)
+			lookupRef, ok := lookupOperation["../components.yaml#/components/parameters/SomeParam"]
+			require.True(t, ok)
+			require.Len(t, lookupRef, 1)
+			require.Equal(t, lookupRef[0].Name, "SomeParam")
+		}(i)
+	}
+	wg.Wait()
 }