btf: fix race due to concurrent read access

btf/btf.go

@@ -66,15 +66,30 @@ func (s *immutableTypes) typeByID(id TypeID) (Type, bool) {
 // mutableTypes is a set of types which may be changed.
 type mutableTypes struct {
 	imm           immutableTypes
+	mu            *sync.RWMutex   // protects copies below
 	copies        map[Type]Type   // map[orig]copy
-	copiedTypeIDs map[Type]TypeID //map[copy]origID
+	copiedTypeIDs map[Type]TypeID // map[copy]origID
 }
 
 // add a type to the set of mutable types.
 //
 // Copies type and all of its children once. Repeated calls with the same type
 // do not copy again.
 func (mt *mutableTypes) add(typ Type, typeIDs map[Type]TypeID) Type {
+	mt.mu.RLock()
+	cpy, ok := mt.copies[typ]
+	mt.mu.RUnlock()
+
+	if ok {
+		// Fast path: the type has been copied before.
+		return cpy
+	}
+
+	// modifyGraphPreorder copies the type graph node by node, so we can't drop
+	// the lock in between.
+	mt.mu.Lock()
+	defer mt.mu.Unlock()
+
 	return modifyGraphPreorder(typ, func(t Type) (Type, bool) {
 		cpy, ok := mt.copies[t]
 		if ok {
@@ -98,6 +113,7 @@ func (mt *mutableTypes) add(typ Type, typeIDs map[Type]TypeID) Type {
 func (mt *mutableTypes) copy() mutableTypes {
 	mtCopy := mutableTypes{
 		mt.imm,
+		&sync.RWMutex{},
 		make(map[Type]Type, len(mt.copies)),
 		make(map[Type]TypeID, len(mt.copiedTypeIDs)),
 	}
@@ -122,6 +138,9 @@ func (mt *mutableTypes) typeID(typ Type) (TypeID, error) {
 		return 0, nil
 	}
 
+	mt.mu.RLock()
+	defer mt.mu.RUnlock()
+
 	id, ok := mt.copiedTypeIDs[typ]
 	if !ok {
 		return 0, fmt.Errorf("no ID for type %s: %w", typ, ErrNotFound)
@@ -343,6 +362,7 @@ func loadRawSpec(btf io.ReaderAt, bo binary.ByteOrder, base *Spec) (*Spec, error
 				typesByName,
 				bo,
 			},
+			&sync.RWMutex{},
 			make(map[Type]Type),
 			make(map[Type]TypeID),
 		},

btf/btf_test.go

@@ -7,7 +7,9 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"runtime"
 	"sync"
+	"sync/atomic"
 	"testing"
 
 	"github.com/go-quicktest/qt"
@@ -525,6 +527,35 @@ func TestFixupDatasecLayout(t *testing.T) {
 	qt.Assert(t, qt.Equals(ds.Vars[5].Offset, 32))
 }
 
+func TestSpecConcurrentAccess(t *testing.T) {
+	spec := vmlinuxTestdataSpec(t)
+
+	maxprocs := runtime.GOMAXPROCS(0)
+	if maxprocs < 2 {
+		t.Error("GOMAXPROCS is lower than 2:", maxprocs)
+	}
+
+	var cond atomic.Int64
+	var wg sync.WaitGroup
+	for i := 0; i < maxprocs; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			cond.Add(1)
+			for cond.Load() != int64(maxprocs) {
+				// Spin to increase the chances of a race.
+			}
+
+			_, _ = spec.AnyTypeByName("gov_update_cpu_data")
+		}()
+
+		// Try to get the Goroutines scheduled and spinning.
+		runtime.Gosched()
+	}
+	wg.Wait()
+}
+
 func BenchmarkSpecCopy(b *testing.B) {
 	spec := vmlinuxTestdataSpec(b)
 	b.ResetTimer()

run-tests.sh

@@ -68,12 +68,12 @@ if [[ "${1:-}" = "--exec-vm" ]]; then
   fi
 
   for ((i = 0; i < 3; i++)); do
-    if ! $sudo virtme-run --kimg "${input}/boot/vmlinuz" --memory 768M --pwd \
+    if ! $sudo virtme-run --kimg "${input}/boot/vmlinuz" --cpus 2 --memory 768M --pwd \
       --rwdir="${testdir}=${testdir}" \
       --rodir=/run/input="${input}" \
       --rwdir=/run/output="${output}" \
-      --script-sh "$(quote_env "${preserved_env[@]}") \"$script\" --exec-test $cmd" \
-      --kopt possible_cpus=2; then # need at least two CPUs for some tests
+      --script-sh "$(quote_env "${preserved_env[@]}") \"$script\" \
+      --exec-test $cmd"; then
       exit 23
     fi