Merge pull request #68 from Basekick-Labs/perf/intern-dict-pool

perf: pool and pre-allocate interned-string dict (#66)

Author: xe-nvdk

CHANGELOG.md

@@ -5,6 +5,7 @@
 - **decode:** `readCode` bsr fast path — when decoding from a byte slice, reads directly from the underlying array instead of dispatching through the `io.ByteReader` interface; eliminates ~900M interface calls/sec at Arc's throughput ([#57](https://github.com/Basekick-Labs/msgpack/issues/57)) (StructUnmarshal **-7.5%**, StructUnmarshalPartially **-6.1%**)
 - **decode:** `PeekCode` bsr fast path — peeks directly at `bsr.data[bsr.pos]` instead of `ReadByte` + `UnreadByte` (two interface calls) ([#59](https://github.com/Basekick-Labs/msgpack/issues/59))
 - **encode:** pool `OmitEmpty` filtered field slices via `sync.Pool` — when fields are actually omitted, the allocated `[]*field` slice is now returned to a pool for reuse instead of being GC'd ([#58](https://github.com/Basekick-Labs/msgpack/issues/58))
+- **encode/decode:** pool and pre-allocate interned-string dict — `SetInternedStringsDictCap(n)` pre-sizes the dict to avoid map rehashing and slice growth; pooled encoders/decoders now reuse dict storage across `Reset()` (cleared in place) instead of discarding it, and `Put*()` drops oversized dicts to keep the pool lean ([#66](https://github.com/Basekick-Labs/msgpack/issues/66))
 
 ---
 

decode.go

@@ -53,6 +53,18 @@ func PutDecoder(dec *Decoder) {
 	} else if dec.buf != nil {
 		dec.buf = dec.buf[:0]
 	}
+	// Drop the interned-string dict if we own it and it grew large, so pool
+	// entries don't permanently retain memory from a one-off large interning
+	// session. A caller-owned dict is always dropped so we never hold a
+	// reference to caller memory across pool round-trips. We check cap(),
+	// not len(), because PutDecoder can see a truncated slice (len=0) whose
+	// backing array is still large.
+	if !dec.dictOwned {
+		dec.dict = nil
+	} else if cap(dec.dict) > internDictPoolCap {
+		dec.dict = nil
+		dec.dictOwned = false
+	}
 	decPool.Put(dec)
 }
 
@@ -82,6 +94,8 @@ type Decoder struct {
 	rec        []byte
 	dict       []string
 	flags      uint32
+	internCap  int  // initial capacity hint for the interned-string dict
+	dictOwned  bool // true when dict was lazily allocated by us, safe to mutate/pool
 }
 
 // NewDecoder returns a new decoder that reads from r.
@@ -102,24 +116,34 @@ func (d *Decoder) Reset(r io.Reader) {
 }
 
 // ResetDict is like Reset, but also resets the dict.
+//
+// A non-nil dict replaces the current dict; the decoder will not append to
+// it or otherwise mutate it (caller retains ownership). A nil dict keeps
+// any internally-owned dict storage for reuse, truncated to empty —
+// subsequent interned decodes skip the slice allocation.
 func (d *Decoder) ResetDict(r io.Reader, dict []string) {
 	d.ResetReader(r)
 	d.flags = 0
 	d.structTag = ""
-	d.dict = dict
+	if dict != nil {
+		d.dict = dict
+		d.dictOwned = false
+	}
 }
 
 func (d *Decoder) WithDict(dict []string, fn func(*Decoder) error) error {
-	oldDict := d.dict
+	oldDict, oldOwned := d.dict, d.dictOwned
 	d.dict = dict
+	d.dictOwned = false
 	err := fn(d)
 	d.dict = oldDict
+	d.dictOwned = oldOwned
 	return err
 }
 
 func (d *Decoder) ResetReader(r io.Reader) {
 	d.mapDecoder = nil
-	d.dict = nil
+	d.releaseOrTruncateDict()
 
 	if br, ok := r.(bufReader); ok {
 		d.r = br
@@ -144,7 +168,18 @@ func (d *Decoder) ResetBytes(data []byte) {
 	d.mapDecoder = nil
 	d.flags = 0
 	d.structTag = ""
-	d.dict = nil
+	d.releaseOrTruncateDict()
+}
+
+// releaseOrTruncateDict reuses the dict storage if we allocated it
+// ourselves; otherwise it drops the reference so a caller-supplied dict is
+// never aliased or appended into by a subsequent reset.
+func (d *Decoder) releaseOrTruncateDict() {
+	if d.dictOwned {
+		d.dict = d.dict[:0]
+	} else {
+		d.dict = nil
+	}
 }
 
 func (d *Decoder) SetMapDecoder(fn func(*Decoder) (interface{}, error)) {
@@ -187,6 +222,27 @@ func (d *Decoder) UseInternedStrings(on bool) {
 	}
 }
 
+// SetInternedStringsDictCap sets an initial capacity hint for the
+// interned-string dict, avoiding slice growth as entries are appended.
+// n is clamped to [0, maxDictLen]; 0 restores lazy allocation.
+//
+// The hint is consulted the next time the dict is allocated — typically on
+// the first interned decode after construction, Reset, or ResetDict. Call
+// it before decoding to guarantee it takes effect.
+//
+// When the decoder is managed by GetDecoder/PutDecoder, PutDecoder drops
+// dicts whose capacity exceeds an internal pool-retention threshold so a
+// single oversized session doesn't permanently bloat pool memory. Setting
+// n above that threshold forfeits cross-Put reuse of the dict.
+func (d *Decoder) SetInternedStringsDictCap(n int) {
+	if n < 0 {
+		n = 0
+	} else if n > maxDictLen {
+		n = maxDictLen
+	}
+	d.internCap = n
+}
+
 // UsePreallocateValues enables preallocating values in chunks
 func (d *Decoder) UsePreallocateValues(on bool) {
 	if on {

encode.go

@@ -58,6 +58,16 @@ func PutEncoder(enc *Encoder) {
 	if cap(enc.wbuf) > 32*1024 {
 		enc.wbuf = nil
 	}
+	// Drop the interned-string dict if we own it and it grew large, so pool
+	// entries don't permanently retain memory from a one-off large interning
+	// session. A caller-owned dict is always dropped so we never hold a
+	// reference to caller memory across pool round-trips.
+	if !enc.dictOwned {
+		enc.dict = nil
+	} else if len(enc.dict) > internDictPoolCap {
+		enc.dict = nil
+		enc.dictOwned = false
+	}
 	encPool.Put(enc)
 }
 
@@ -94,6 +104,8 @@ type Encoder struct {
 	buf       []byte
 	timeBuf   []byte
 	flags     uint32
+	internCap int  // initial capacity hint for the interned-string dict
+	dictOwned bool // true when dict was lazily allocated by us, safe to mutate/pool
 }
 
 // NewEncoder returns a new encoder that writes to w.
@@ -116,23 +128,33 @@ func (e *Encoder) Reset(w io.Writer) {
 }
 
 // ResetDict is like Reset, but also resets the dict.
+//
+// A non-nil dict replaces the current dict; the encoder will not mutate it
+// (caller retains ownership). A nil dict keeps any internally-owned dict
+// storage for reuse, cleared to empty — subsequent interned encodes skip
+// the map allocation.
 func (e *Encoder) ResetDict(w io.Writer, dict map[string]int) {
 	e.ResetWriter(w)
 	e.flags = 0
 	e.structTag = ""
-	e.dict = dict
+	if dict != nil {
+		e.dict = dict
+		e.dictOwned = false
+	}
 }
 
 func (e *Encoder) WithDict(dict map[string]int, fn func(*Encoder) error) error {
-	oldDict := e.dict
+	oldDict, oldOwned := e.dict, e.dictOwned
 	e.dict = dict
+	e.dictOwned = false
 	err := fn(e)
 	e.dict = oldDict
+	e.dictOwned = oldOwned
 	return err
 }
 
 func (e *Encoder) ResetWriter(w io.Writer) {
-	e.dict = nil
+	e.releaseOrClearDict()
 	if bw, ok := w.(writer); ok {
 		e.w = bw
 	} else if w == nil {
@@ -150,7 +172,18 @@ func (e *Encoder) resetForMarshal() {
 	e.w = &e.bsw
 	e.flags = 0
 	e.structTag = ""
-	e.dict = nil
+	e.releaseOrClearDict()
+}
+
+// releaseOrClearDict reuses the dict storage if we allocated it ourselves;
+// otherwise it drops the reference so a caller-supplied dict is never
+// mutated by a subsequent reset.
+func (e *Encoder) releaseOrClearDict() {
+	if e.dictOwned {
+		clear(e.dict)
+	} else {
+		e.dict = nil
+	}
 }
 
 // SetSortMapKeys causes the Encoder to encode map keys in increasing order.
@@ -220,6 +253,27 @@ func (e *Encoder) UseInternedStrings(on bool) {
 	}
 }
 
+// SetInternedStringsDictCap sets an initial capacity hint for the
+// interned-string dict, avoiding map rehashing as entries are added.
+// n is clamped to [0, maxDictLen]; 0 restores lazy allocation.
+//
+// The hint is consulted the next time the dict is allocated — typically on
+// the first interned encode after construction, Reset, or ResetDict. Call
+// it before encoding to guarantee it takes effect.
+//
+// When the encoder is managed by GetEncoder/PutEncoder, PutEncoder drops
+// dicts whose length exceeds an internal pool-retention threshold so a
+// single oversized session doesn't permanently bloat pool memory. Setting
+// n above that threshold forfeits cross-Put reuse of the dict.
+func (e *Encoder) SetInternedStringsDictCap(n int) {
+	if n < 0 {
+		n = 0
+	} else if n > maxDictLen {
+		n = maxDictLen
+	}
+	e.internCap = n
+}
+
 func (e *Encoder) Encode(v interface{}) error {
 	switch v := v.(type) {
 	case nil:

intern.go

@@ -11,6 +11,12 @@ import (
 const (
 	minInternedStringLen = 3
 	maxDictLen           = math.MaxUint16
+
+	// internDictPoolCap is the threshold above which a pooled Encoder/Decoder
+	// drops its interned-string dict in Put*() instead of retaining it for
+	// reuse. Mirrors the wbuf/buf cap-drop pattern: keeps hot, small dicts
+	// warm without letting a one-off large session bloat pool memory.
+	internDictPoolCap = 4096
 )
 
 var internedStringExtID = int8(math.MinInt8)
@@ -63,7 +69,8 @@ func (e *Encoder) encodeInternedString(s string, intern bool) error {
 
 	if intern && len(s) >= minInternedStringLen && len(e.dict) < maxDictLen {
 		if e.dict == nil {
-			e.dict = make(map[string]int)
+			e.dict = make(map[string]int, e.internCap)
+			e.dictOwned = true
 		}
 		idx := len(e.dict)
 		e.dict[s] = idx
@@ -227,6 +234,12 @@ func (d *Decoder) decodeInternedStringWithLen(n int, intern bool) (string, error
 	}
 
 	if intern && len(s) >= minInternedStringLen && len(d.dict) < maxDictLen {
+		if d.dict == nil {
+			if d.internCap > 0 {
+				d.dict = make([]string, 0, d.internCap)
+			}
+			d.dictOwned = true
+		}
 		d.dict = append(d.dict, s)
 	}