feat: Pretouch recursively for large/deep struct

bytedance · Nov 24, 2021 · 5220d69 · 5220d69
1 parent 9a95e9d
commit 5220d69
Show file tree

Hide file tree

Showing 9 changed files with 243 additions and 25 deletions.
diff --git a/decoder/assembler_test.go b/decoder/assembler_test.go
@@ -670,7 +670,7 @@ type JsonStruct struct {
 func TestAssembler_DecodeStruct(t *testing.T) {
     var v JsonStruct
     s := `{"A": 123, "B": "asdf", "C": {"qwer": 4567}, "D": [1, 2, 3, 4, 5]}`
-    p, err := make(_Compiler).compile(reflect.TypeOf(v))
+    p, err := newCompiler().compile(reflect.TypeOf(v))
     require.NoError(t, err)
     k := new(_Stack)
     a := newAssembler(p)
@@ -693,7 +693,7 @@ type Tx struct {
 func TestAssembler_DecodeStruct_SinglePrivateField(t *testing.T) {
     var v Tx
     s := `{"x": 1}`
-    p, err := make(_Compiler).compile(reflect.TypeOf(v))
+    p, err := newCompiler().compile(reflect.TypeOf(v))
     require.NoError(t, err)
     k := new(_Stack)
     a := newAssembler(p)
@@ -707,7 +707,7 @@ func TestAssembler_DecodeStruct_SinglePrivateField(t *testing.T) {
 func TestAssembler_DecodeByteSlice_Bin(t *testing.T) {
     var v []byte
     s := `"aGVsbG8sIHdvcmxk"`
-    p, err := make(_Compiler).compile(reflect.TypeOf(v))
+    p, err := newCompiler().compile(reflect.TypeOf(v))
     require.NoError(t, err)
     k := new(_Stack)
     a := newAssembler(p)
@@ -721,7 +721,7 @@ func TestAssembler_DecodeByteSlice_Bin(t *testing.T) {
 func TestAssembler_DecodeByteSlice_List(t *testing.T) {
     var v []byte
     s := `[104, 101, 108, 108, 111, 44, 32, 119, 111, 114, 108, 100]`
-    p, err := make(_Compiler).compile(reflect.TypeOf(v))
+    p, err := newCompiler().compile(reflect.TypeOf(v))
     require.NoError(t, err)
     k := new(_Stack)
     a := newAssembler(p)

diff --git a/decoder/compiler.go b/decoder/compiler.go
@@ -474,10 +474,31 @@ func (self _Program) disassemble() string {
     return strings.Join(append(ret, "\tend"), "\n")
 }
 
-type (
-	_Compiler map[reflect.Type]bool
+type _CompileOpt uint64
+
+const (
+    // ALL uncompiled types will recored in _Compiler's rec map.
+    _RecordUncompiled _CompileOpt = 1 << iota
 )
 
+type _Compiler struct {
+    opts _CompileOpt
+    tab  map[reflect.Type]bool
+    rec  map[reflect.Type]bool
+}
+
+func newCompiler() *_Compiler {
+    return &_Compiler {
+        tab: map[reflect.Type]bool{},
+    }
+}
+
+func (self *_Compiler) recordUncompiled() *_Compiler {
+    self.opts |= _RecordUncompiled
+    self.rec = map[reflect.Type]bool{}
+    return self
+}
+
 func (self _Compiler) rescue(ep *error) {
     if val := recover(); val != nil {
         if err, ok := val.(error); ok {
@@ -495,7 +516,7 @@ func (self _Compiler) compile(vt reflect.Type) (ret _Program, err error) {
 }
 
 func (self _Compiler) compileOne(p *_Program, sp int, vt reflect.Type) {
-    ok := self[vt]
+    ok := self.tab[vt]
     pt := reflect.PtrTo(vt)
 
     /* check for recursive nesting */
@@ -533,9 +554,9 @@ func (self _Compiler) compileOne(p *_Program, sp int, vt reflect.Type) {
 
     /* enter the recursion */
     p.add(_OP_lspace)
-    self[vt] = true
+    self.tab[vt] = true
     self.compileOps(p, sp, vt)
-    delete(self, vt)
+    delete(self.tab, vt)
 }
 
 func (self _Compiler) compileOps(p *_Program, sp int, vt reflect.Type) {
@@ -793,6 +814,9 @@ func (self _Compiler) compileStringBody(p *_Program) {
 func (self _Compiler) compileStruct(p *_Program, sp int, vt reflect.Type) {
     if sp >= _MAX_STACK || p.pc() >= _MAX_ILBUF {
         p.rtt(_OP_recurse, vt)
+        if (self.opts & _RecordUncompiled) != 0 {
+            self.rec[vt] = true
+        }
     } else {
         self.compileStructBody(p, sp, vt)
     }

diff --git a/decoder/compiler_test.go b/decoder/compiler_test.go
@@ -24,7 +24,7 @@ import (
 )
 
 func TestCompiler_Compile(t *testing.T) {
-    prg, err := make(_Compiler).compile(reflect.TypeOf(TwitterStruct{}))
+    prg, err := newCompiler().compile(reflect.TypeOf(TwitterStruct{}))
     assert.Nil(t, err)
     prg.disassemble()
 }
diff --git a/decoder/decoder.go b/decoder/decoder.go
@@ -105,8 +105,57 @@ func (self *Decoder) DisallowUnknownFields() {
 }
 
 // Pretouch compiles vt ahead-of-time to avoid JIT compilation on-the-fly, in
-// order to reduce the first-hit latency.
-func Pretouch(vt reflect.Type) (err error) {
-    _, err = findOrCompile(rt.UnpackType(vt))
-    return
+// order to reduce the first-hit latency. Depth is the recursive times.
+func Pretouch(vt reflect.Type, depth ...int) (err error) {
+    var dep int
+    for _, arg := range depth {
+        dep = arg
+        break
+    }
+    return pretouchRec(map[reflect.Type]bool{vt:true}, dep)
+}
+
+func pretouchType(_vt reflect.Type, depth int) (map[reflect.Type]bool, error) {
+    compiler := newCompiler().recordUncompiled()
+    if depth > 0 { // need recursive pretouch
+        compiler = compiler.recordUncompiled()
+    }
+
+    /* compile function */
+    decoder := func(vt *rt.GoType) (interface{}, error) {
+        if pp, err := compiler.compile(_vt); err != nil {
+            return nil, err
+        } else {
+            return newAssembler(pp).Load(), nil
+        }
+    }
+
+    /* find or compile */
+    vt := rt.UnpackType(_vt)
+    if val := programCache.Get(vt); val != nil {
+        return nil, nil
+    } else if _, err := programCache.Compute(vt, decoder); err == nil {
+        return compiler.rec, nil
+    } else {
+        return nil, err
+    }
+}
+
+func pretouchRec(vtm map[reflect.Type]bool, depth int) (err error) {
+    if depth < 0 || len(vtm) == 0 {
+        return nil
+    }
+
+    next := make(map[reflect.Type]bool)
+    for vt, _ := range(vtm) {
+        sub, err := pretouchType(vt, depth)
+        if err != nil {
+            return err
+        }
+        for svt, _ := range(sub) {
+            next[svt] = true
+        }
+    }
+
+    return pretouchRec(next, depth - 1)
 }
diff --git a/decoder/pools.go b/decoder/pools.go
@@ -115,7 +115,7 @@ func referenceFields(v *caching.FieldMap) int64 {
 }
 
 func makeDecoder(vt *rt.GoType) (interface{}, error) {
-    if pp, err := make(_Compiler).compile(vt.Pack()); err != nil {
+    if pp, err := newCompiler().compile(vt.Pack()); err != nil {
         return nil, err
     } else {
         return newAssembler(pp).Load(), nil

diff --git a/encoder/compiler.go b/encoder/compiler.go
@@ -370,9 +370,18 @@ func (self _Program) disassemble() string {
     return strings.Join(append(ret, "\tend"), "\n")
 }
 
+type _CompileOpt uint64
+
+const (
+    // ALL uncompiled types will recored in _Compiler's rec map.
+    _RecordUncompiled _CompileOpt = 1 << iota
+)
+
 type _Compiler struct {
-    pv  bool
-    tab map[reflect.Type]bool
+    opts _CompileOpt
+    pv   bool
+    tab  map[reflect.Type]bool
+    rec  map[reflect.Type]bool
 }
 
 func newCompiler() *_Compiler {
@@ -381,6 +390,12 @@ func newCompiler() *_Compiler {
     }
 }
 
+func (self *_Compiler) recordUncompiled() *_Compiler {
+    self.opts |= _RecordUncompiled
+    self.rec = map[reflect.Type]bool{}
+    return self
+}
+
 func (self *_Compiler) rescue(ep *error) {
     if val := recover(); val != nil {
         if err, ok := val.(error); ok {
@@ -645,6 +660,9 @@ func (self *_Compiler) compileString(p *_Program, vt reflect.Type) {
 func (self *_Compiler) compileStruct(p *_Program, sp int, vt reflect.Type) {
     if sp >= _MAX_STACK || p.pc() >= _MAX_ILBUF {
         p.rtt(_OP_recurse, vt)
+        if (self.opts & _RecordUncompiled) != 0 {
+            self.rec[vt] = true
+        }
     } else {
         self.compileStructBody(p, sp, vt)
     }

diff --git a/encoder/encoder.go b/encoder/encoder.go
@@ -149,8 +149,57 @@ func EncodeIndented(val interface{}, prefix string, indent string, opts Options)
 }
 
 // Pretouch compiles vt ahead-of-time to avoid JIT compilation on-the-fly, in
-// order to reduce the first-hit latency.
-func Pretouch(vt reflect.Type) (err error) {
-    _, err = findOrCompile(rt.UnpackType(vt))
-    return
+// order to reduce the first-hit latency. Depth is the recursive times.
+func Pretouch(vt reflect.Type, depth ...int) (err error) {
+    var dep int
+    for _, arg := range depth {
+        dep = arg
+        break
+    }
+    return pretouchRec(map[reflect.Type]bool{vt:true}, dep)
+}
+
+func pretouchType(_vt reflect.Type, depth int) (map[reflect.Type]bool, error) {
+    compiler := newCompiler()
+    if depth > 0 { // need recursive pretouch
+        compiler = compiler.recordUncompiled()
+    }
+
+    /* compile function */
+    encoder := func(vt *rt.GoType) (interface{}, error) {
+        if pp, err := compiler.compile(_vt); err != nil {
+            return nil, err
+        } else {
+            return newAssembler(pp).Load(), nil
+        }
+    }
+
+    /* find or compile */
+    vt := rt.UnpackType(_vt)
+    if val := programCache.Get(vt); val != nil {
+        return nil, nil
+    } else if _, err := programCache.Compute(vt, encoder); err == nil {
+        return compiler.rec, nil
+    } else {
+        return nil, err
+    }
 }
+
+func pretouchRec(vtm map[reflect.Type]bool, depth int) (err error) {
+    if depth < 0 || len(vtm) == 0 {
+        return nil
+    }
+
+    next := make(map[reflect.Type]bool)
+    for vt, _ := range(vtm) {
+        sub, err := pretouchType(vt, depth)
+        if err != nil {
+            return err
+        }
+        for svt, _ := range(sub) {
+            next[svt] = true
+        }
+    }
+
+    return pretouchRec(next, depth - 1)
+}
diff --git a/issue_test/issue138_test.go b/issue_test/issue138_test.go
@@ -0,0 +1,78 @@
+/*
+ * Copyright 2021 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 issue_test
+
+import (
+    `fmt`
+    `time`
+    `testing`
+    `reflect`
+
+    `github.com/stretchr/testify/require`
+    `github.com/bytedance/sonic`
+)
+
+type Issue138_DeepStruct struct {
+    L0 struct {
+    L1 struct {
+    L2 struct {
+    L3 struct {
+    L4 struct {
+    L5 struct {
+    L6 struct {
+    L7 struct {
+    L8 struct {
+    L9 struct {
+    L10 struct {
+    L11 struct {
+    L12 struct {
+    L13 struct {
+    L14 struct {
+    L15 struct {
+    L16 struct {
+    L17 struct {
+    L18 struct {
+    L19 struct {
+    L20 struct {
+    L21 struct {
+    L22 struct {
+        A int
+        B string
+        C []float64
+        E map[string]bool
+        F *Issue138_DeepStruct
+    }}}}}}}}}}}}}}}}}}}}}}}
+}
+
+func testPretouchTime(depth int) {
+    start := time.Now()
+    sonic.Pretouch(reflect.TypeOf(Issue138_DeepStruct{}), depth)
+    elapsed := time.Since(start)
+    fmt.Printf("Pretouch with recursive depth %d, time is %s\n", depth, elapsed)
+}
+
+func TestIssue138_PretouchTime(t *testing.T) {
+    testPretouchTime(1)
+    var obj Issue138_DeepStruct
+    start := time.Now()
+    data, err := sonic.Marshal(obj)
+    err = sonic.Unmarshal([]byte(data), &obj)
+    elapsed := time.Since(start)
+    fmt.Printf("Marshal and unmarshal time is %s\n", elapsed)
+    require.NoError(t, err)
+}
+
diff --git a/sonic.go b/sonic.go
@@ -73,11 +73,11 @@ func UnmarshalString(buf string, val interface{}) error {
 }
 
 // Pretouch compiles vt ahead-of-time to avoid JIT compilation on-the-fly, in
-// order to reduce the first-hit latency.
-func Pretouch(vt reflect.Type) error {
-    if err := encoder.Pretouch(vt); err != nil {
+// order to reduce the first-hit latency. Depth is the recursive times.
+func Pretouch(vt reflect.Type, depth ...int) error {
+    if err := encoder.Pretouch(vt, depth...); err != nil {
         return err
-    } else if err = decoder.Pretouch(vt); err != nil {
+    } else if err = decoder.Pretouch(vt, depth...); err != nil {
         return err
     } else {
         return nil