fix: potential ProgramCache OOM under high-load

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
+    }
+}
+
+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
     }
-
-    /* slow path: not found, compile the type on the fly */
-    if pp, ex = make(_Compiler).compile(vt.Pack()); ex != nil {
-        return nil, ex
-    }
-
-    /* link the program, and put it into cache */
-    fn = newAssembler(pp).Load()
-    programCache.Put(vt, fn)
-    return fn, nil
 }

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
+    }
+}
+
+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
     }
-
-    /* slow path: not found, compile the type on the fly */
-    if pp, ex = newCompiler().compile(vt.Pack()); ex != nil {
-        return nil, ex
-    }
-
-    /* link the program, and put it into cache */
-    fn = newAssembler(pp).Load()
-    programCache.Put(vt, fn)
-    return fn, nil
 }

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)

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
 }

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
+            })
         }
     }()