From aa527c5992c8a35bc031a38203ef5962651ce9a3 Mon Sep 17 00:00:00 2001
From: "felix.fengmin" <felix.fengmin@bytedance.com>
Date: Wed, 8 Nov 2023 11:12:45 +0800
Subject: [PATCH] optimize: use lockfree to register new modules

---
 loader/stubs.go      | 34 ++++++++++++++++++++++++-------
 loader/stubs_test.go | 48 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 7 deletions(-)
 create mode 100644 loader/stubs_test.go

diff --git a/loader/stubs.go b/loader/stubs.go
index 8377649b7..80f8de836 100644
--- a/loader/stubs.go
+++ b/loader/stubs.go
@@ -17,7 +17,8 @@
 package loader
 
 import (
-    `sync`
+    "sync/atomic"
+    "unsafe"
     _ `unsafe`
 )
 
@@ -25,16 +26,35 @@ import (
 //goland:noinspection GoUnusedGlobalVariable
 var lastmoduledatap *moduledata
 
-var moduledataMux sync.Mutex
-
 func registerModule(mod *moduledata) {
-    moduledataMux.Lock()
-    lastmoduledatap.next = mod
-    lastmoduledatap = mod
-    moduledataMux.Unlock()
+    registerModuleLockFree(&lastmoduledatap, mod)
 }
 
 //go:linkname moduledataverify1 runtime.moduledataverify1
 func moduledataverify1(_ *moduledata)
 
+func registerModuleLockFree(tail **moduledata, mod *moduledata) {
+    for {
+        oldTail := loadModule(tail)
+        if casModule(tail, oldTail, mod) {
+            storeModule(&oldTail.next, mod)
+            break
+        }
+    }
+}
+
+func loadModule(p **moduledata) *moduledata {
+    return (*moduledata)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(p))))
+}
 
+func storeModule(p **moduledata, value *moduledata) {
+    atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(p)), unsafe.Pointer(value))
+}
+
+func casModule(p **moduledata, oldValue *moduledata, newValue *moduledata) bool {
+    return atomic.CompareAndSwapPointer(
+        (*unsafe.Pointer)(unsafe.Pointer(p)),
+        unsafe.Pointer(oldValue),
+        unsafe.Pointer(newValue),
+    )
+}
diff --git a/loader/stubs_test.go b/loader/stubs_test.go
new file mode 100644
index 000000000..45bc6e35b
--- /dev/null
+++ b/loader/stubs_test.go
@@ -0,0 +1,48 @@
+/*
+ * Copyright 2023 ByteDance Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package loader
+
+import (
+    "testing"
+    "sync"
+)
+
+func Test_registerModuleLockFree(t *testing.T) {
+    n, parallel := 1000, 8
+    head := moduledata{}
+    tail := &head
+    wg := sync.WaitGroup{}
+    wg.Add(parallel)
+    filler := func(n int) {
+        defer wg.Done()
+        for i := 0; i < n; i++ {
+            m := &moduledata{}
+            registerModuleLockFree(&tail, m)
+        }
+    }
+    for i := 0; i < parallel; i++ {
+        go filler(n)
+    }
+    wg.Wait()
+    i := 0
+    for p := head.next; p != nil; p = p.next {
+        i += 1
+    }
+    if i != parallel * n {
+        t.Errorf("got %v, expected %v", i, parallel * n)
+    }
+}