Repository: bytedance/sonic
Pull request #91 (2021-09-09 04:45:11)

fix: potential ProgramCache OOM under high-load
A lot of threads rushing into the compiler may cause OOM. This PR fixes this problem.


Patch:
diff --git a/decoder/pools.go b/decoder/pools.go
index e0c68fc..a36cc25 100644
--- a/decoder/pools.go
+++ b/decoder/pools.go
@@ -86,24 +86,20 @@ func referenceFields(v *caching.FieldMap) int64 {
     return int64(uintptr(unsafe.Pointer(v)))
 }
 
-func findOrCompile(vt *rt.GoType) (_Decoder, error) {
-    var ex error
-    var fn _Decoder
-    var pp _Program
-    var fv interface{}
-
-    /* fast path: the program is in the cache */
-    if fv = programCache.Get(vt); fv != nil {
-        return fv.(_Decoder), nil
+func makeDecoder(vt *rt.GoType) (interface{}, error) {
+    if pp, err := make(_Compiler).compile(vt.Pack()); err != nil {
+        return nil, err
+    } else {
+        return newAssembler(pp).Load(), nil
     }
+}
 
-    /* slow path: not found, compile the type on the fly */
-    if pp, ex = make(_Compiler).compile(vt.Pack()); ex != nil {
-        return nil, ex
+func findOrCompile(vt *rt.GoType) (_Decoder, error) {
+    if val := programCache.Get(vt); val != nil {
+        return val.(_Decoder), nil
+    } else if ret, err := programCache.Compute(vt, makeDecoder); err == nil {
+        return ret.(_Decoder), nil
+    } else {
+        return nil, err
     }
-
-    /* link the program, and put it into cache */
-    fn = newAssembler(pp).Load()
-    programCache.Put(vt, fn)
-    return fn, nil
 }
diff --git a/encoder/pools.go b/encoder/pools.go
index fbaf751..c9257b8 100644
--- a/encoder/pools.go
+++ b/encoder/pools.go
@@ -95,24 +95,20 @@ func freeBuffer(p *bytes.Buffer) {
     bufferPool.Put(p)
 }
 
-func findOrCompile(vt *rt.GoType) (_Encoder, error) {
-    var ex error
-    var fn _Encoder
-    var pp _Program
-    var fv interface{}
-
-    /* fast path: the program is in the cache */
-    if fv = programCache.Get(vt); fv != nil {
-        return fv.(_Encoder), nil
+func makeEncoder(vt *rt.GoType) (interface{}, error) {
+    if pp, err := newCompiler().compile(vt.Pack()); err != nil {
+        return nil, err
+    } else {
+        return newAssembler(pp).Load(), nil
     }
+}
 
-    /* slow path: not found, compile the type on the fly */
-    if pp, ex = newCompiler().compile(vt.Pack()); ex != nil {
-        return nil, ex
+func findOrCompile(vt *rt.GoType) (_Encoder, error) {
+    if val := programCache.Get(vt); val != nil {
+        return val.(_Encoder), nil
+    } else if ret, err := programCache.Compute(vt, makeEncoder); err == nil {
+        return ret.(_Encoder), nil
+    } else {
+        return nil, err
     }
-
-    /* link the program, and put it into cache */
-    fn = newAssembler(pp).Load()
-    programCache.Put(vt, fn)
-    return fn, nil
 }
diff --git a/internal/caching/fcache.go b/internal/caching/fcache.go
index 54fd0f1..6fbc6ca 100644
--- a/internal/caching/fcache.go
+++ b/internal/caching/fcache.go
@@ -60,7 +60,7 @@ func (self *FieldMap) At(p uint64) *FieldEntry {
 }
 
 // Get searches FieldMap by name. JIT generated assembly does NOT call this
-// function, rather it implements it's own version directly in assembly. So
+// function, rather it implements its own version directly in assembly. So
 // we must ensure this function stays in sync with the JIT generated one.
 func (self *FieldMap) Get(name string) int {
     h := StrHash(name)
diff --git a/internal/caching/pcache.go b/internal/caching/pcache.go
index a42ec05..42c3413 100644
--- a/internal/caching/pcache.go
+++ b/internal/caching/pcache.go
@@ -149,8 +149,25 @@ func (self *ProgramCache) Get(vt *rt.GoType) interface{} {
     return (*_ProgramMap)(atomic.LoadPointer(&self.p)).get(vt)
 }
 
-func (self *ProgramCache) Put(vt *rt.GoType, fn interface{}) {
+func (self *ProgramCache) Compute(vt *rt.GoType, compute func(*rt.GoType) (interface{}, error)) (interface{}, error) {
+    var err error
+    var val interface{}
+
+    /* use defer to prevent inlining of this function */
     self.m.Lock()
-    atomic.StorePointer(&self.p, unsafe.Pointer((*_ProgramMap)(atomic.LoadPointer(&self.p)).add(vt, fn)))
-    self.m.Unlock()
+    defer self.m.Unlock()
+
+    /* double check with write lock held */
+    if val = self.Get(vt); val != nil {
+        return val, nil
+    }
+
+    /* compute the value */
+    if val, err = compute(vt); err != nil {
+        return nil, err
+    }
+
+    /* update the RCU cache */
+    atomic.StorePointer(&self.p, unsafe.Pointer((*_ProgramMap)(atomic.LoadPointer(&self.p)).add(vt, val)))
+    return val, nil
 }


Test patch:
diff --git a/internal/caching/pcache_test.go b/internal/caching/pcache_test.go
index 688b8a4..b2b82aa 100644
--- a/internal/caching/pcache_test.go
+++ b/internal/caching/pcache_test.go
@@ -36,7 +36,9 @@ func TestPcacheRace(t *testing.T) {
         var k = map[string]interface{}{}
         <- start
         for i:=0; i<100; i++ {
-            pc.Put(rt.UnpackEface(k).Type, k)
+            _, _ = pc.Compute(rt.UnpackEface(k).Type, func(_ *rt.GoType) (interface{}, error) {
+                return map[string]interface{}{}, nil
+            })
         }
     }()