opt:(loader) lock-free impl for registerModule()

bytedance · Nov 7, 2023 · 0b5832a · 0b5832a
1 parent ba39bc4
commit 0b5832a
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 35 deletions.
diff --git a/loader/loader_go117_test.go b/loader/loader_go117_test.go
@@ -23,40 +23,15 @@ import (
     `runtime`
     `runtime/debug`
     `strconv`
+    `sync`
     `testing`
     `unsafe`
 
     `github.com/bytedance/sonic/internal/rt`
     `github.com/stretchr/testify/require`
 )
 
-func TestLoad(t *testing.T) {
-    // defer func() {
-    //     if r := recover(); r != nil {
-    //         runtime.GC()
-    //         if r != "hook1" {
-    //             t.Fatal("not right panic:" + r.(string))
-    //         }
-    //     } else {
-    //         t.Fatal("not panic")
-    //     }
-    // }()
-
-    var hstr string
-
-    type TestFunc func(i *int, hook func(i *int)) int
-    var hook = func(i *int) {
-        runtime.GC()
-        debug.FreeOSMemory()
-        hstr = ("hook" + strconv.Itoa(*i))
-        runtime.GC()
-        debug.FreeOSMemory()
-    }
-    // var f TestFunc = func(i *int, hook func(i *int)) int {
-    //     var t = *i
-    //     hook(i)
-    //     return t + *i
-    // }
+func makeFn() ([]byte, *Func) {
     bc := []byte {
         0x48, 0x83, 0xec, 0x18,         // (0x00) subq $24, %rsp
         0x48, 0x89, 0x6c, 0x24, 0x10,   // (0x04) movq %rbp, 16(%rsp)
@@ -119,7 +94,39 @@ func TestLoad(t *testing.T) {
     locals.AddField(false)
     fn.LocalsPointerMaps = locals.Build()
 
-    rets := Load(bc, []Func{fn}, "dummy_module", []string{"github.com/bytedance/sonic/dummy.go"})
+    return bc, &fn
+}
+
+func TestLoad(t *testing.T) {
+    // defer func() {
+    //     if r := recover(); r != nil {
+    //         runtime.GC()
+    //         if r != "hook1" {
+    //             t.Fatal("not right panic:" + r.(string))
+    //         }
+    //     } else {
+    //         t.Fatal("not panic")
+    //     }
+    // }()
+
+    var hstr string
+
+    type TestFunc func(i *int, hook func(i *int)) int
+    var hook = func(i *int) {
+        runtime.GC()
+        debug.FreeOSMemory()
+        hstr = ("hook" + strconv.Itoa(*i))
+        runtime.GC()
+        debug.FreeOSMemory()
+    }
+    // var f TestFunc = func(i *int, hook func(i *int)) int {
+    //     var t = *i
+    //     hook(i)
+    //     return t + *i
+    // }
+
+    bc, fn := makeFn()
+    rets := Load(bc, []Func{*fn}, "dummy_module", []string{"github.com/bytedance/sonic/dummy.go"})
     println("func address ", *(*unsafe.Pointer)(rets[0]))
     // for k, _ := range moduleCache.m {
     //     spew.Dump(k)
@@ -137,3 +144,17 @@ func TestLoad(t *testing.T) {
     require.Equal(t, "github.com/bytedance/sonic/dummy.go", file)
     require.Equal(t, 0, line)
 }
+
+func TestLoaderRace(t *testing.T) {
+    N := 100
+    wg := sync.WaitGroup{}
+    for i:=0; i<N; i++ {
+        wg.Add(1)
+        go func(i int) {
+            defer wg.Done()
+            bc, fn := makeFn()
+            _ = Load(bc, []Func{*fn}, "dummy"+strconv.Itoa(i), []string{"github.com/bytedance/sonic/dummy"+strconv.Itoa(i)+".go"})
+        }(i)
+    }
+    wg.Wait()
+}
diff --git a/loader/stubs.go b/loader/stubs.go
@@ -17,21 +17,31 @@
 package loader
 
 import (
-    `sync`
+    `sync/atomic`
+    `unsafe`
     _ `unsafe`
 )
 
 //go:linkname lastmoduledatap runtime.lastmoduledatap
 //goland:noinspection GoUnusedGlobalVariable
 var lastmoduledatap *moduledata
 
-var moduledataMux sync.Mutex
-
 func registerModule(mod *moduledata) {
-    moduledataMux.Lock()
-    lastmoduledatap.next = mod
-    lastmoduledatap = mod
-    moduledataMux.Unlock()
+load:
+    old := (*moduledata)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&lastmoduledatap))))
+    next := (*moduledata)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&old.next))))
+    // load and swap old.next
+    if !atomic.CompareAndSwapPointer((*unsafe.Pointer)(unsafe.Pointer(&old.next)), 
+        unsafe.Pointer(next), unsafe.Pointer(mod)){
+        goto load
+    }
+    // sucussefully exchange old.next, then load and swap lastmoduledatap
+    if !atomic.CompareAndSwapPointer((*unsafe.Pointer)(unsafe.Pointer(&lastmoduledatap)), 
+        unsafe.Pointer(old), unsafe.Pointer(mod)) {
+        // exchange fail, recover old.next...
+        atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(&old.next)), unsafe.Pointer(next))
+        goto load
+    }
 }
 
 //go:linkname moduledataverify1 runtime.moduledataverify1