Merge pull request #163 from maxmind/greg/eng-3415

Cache nested structures in Unmarshaler

Author: oschwald

CHANGELOG.md

@@ -1,5 +1,14 @@
 # CHANGELOG
 
+## 1.2.0
+
+* The `mmdbtype.Unmarshaler` now caches nested structures, maps and slices,
+  in addition to top-level values. This improves performance when loading
+  databases with shared nested data structures.
+* The zero value of `mmdbtype.Unmarshaler` is now documented as safe to use
+  for unmarshaling without caching enabled. Use `NewUnmarshaler()` when you
+  want caching.
+
 ## 1.1.0 (2025-10-08)
 
 * Removed unnecessary deep copies in inserter. GitHub #119.

go.mod

@@ -3,7 +3,7 @@ module github.com/maxmind/mmdbwriter
 go 1.24.0
 
 require (
-	github.com/oschwald/maxminddb-golang/v2 v2.0.0
+	github.com/oschwald/maxminddb-golang/v2 v2.1.0
 	github.com/stretchr/testify v1.11.1
 	go4.org/netipx v0.0.0-20231129151722-fdeea329fbba
 )

go.sum

@@ -2,6 +2,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/oschwald/maxminddb-golang/v2 v2.0.0 h1:Gyljxck1kHbBxDgLM++NfDWBqvu1pWWfT8XbosSo0bo=
 github.com/oschwald/maxminddb-golang/v2 v2.0.0/go.mod h1:gG4V88LsawPEqtbL1Veh1WRh+nVSYwXzJ1P5Fcn77g0=
+github.com/oschwald/maxminddb-golang/v2 v2.1.0 h1:2Iv7lmG9XtxuZA/jFAsd7LnZaC1E59pFsj5O/nU15pw=
+github.com/oschwald/maxminddb-golang/v2 v2.1.0/go.mod h1:gG4V88LsawPEqtbL1Veh1WRh+nVSYwXzJ1P5Fcn77g0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=

mmdbtype/types.go

@@ -295,6 +295,8 @@ func (t Int32) WriteTo(w writer) (int64, error) {
 }
 
 // Map is the MaxMind DB map type.
+//
+//nolint:recvcheck // preexisting/interface
 type Map map[String]DataType
 
 var _ DataType = Map(nil)
@@ -341,6 +343,11 @@ func (t Map) typeNum() typeNum {
 
 // UnmarshalMaxMindDB implements the mmdbdata.Unmarshaler interface.
 func (t *Map) UnmarshalMaxMindDB(decoder *mmdbdata.Decoder) error {
+	return t.unmarshalMaxMindDB(decoder, nil)
+}
+
+// unmarshalMaxMindDB is the internal implementation that supports caching.
+func (t *Map) unmarshalMaxMindDB(decoder *mmdbdata.Decoder, cache map[uint]DataType) error {
 	iter, size, err := decoder.ReadMap()
 	if err != nil {
 		return fmt.Errorf("reading Map: %w", err)
@@ -352,7 +359,7 @@ func (t *Map) UnmarshalMaxMindDB(decoder *mmdbdata.Decoder) error {
 			return iterErr
 		}
 
-		value, err := decodeDataTypeValue(decoder)
+		value, err := decodeDataTypeValue(decoder, cache)
 		if err != nil {
 			return err
 		}
@@ -511,6 +518,8 @@ func (t Pointer) WriteTo(w writer) (int64, error) {
 }
 
 // Slice is the MaxMind DB array type.
+//
+//nolint:recvcheck // preexisting/interface
 type Slice []DataType
 
 var _ DataType = Slice(nil)
@@ -557,6 +566,11 @@ func (t Slice) typeNum() typeNum {
 
 // UnmarshalMaxMindDB implements the mmdbdata.Unmarshaler interface.
 func (t *Slice) UnmarshalMaxMindDB(decoder *mmdbdata.Decoder) error {
+	return t.unmarshalMaxMindDB(decoder, nil)
+}
+
+// unmarshalMaxMindDB is the internal implementation that supports caching.
+func (t *Slice) unmarshalMaxMindDB(decoder *mmdbdata.Decoder, cache map[uint]DataType) error {
 	iter, size, err := decoder.ReadSlice()
 	if err != nil {
 		return fmt.Errorf("reading Slice: %w", err)
@@ -568,7 +582,7 @@ func (t *Slice) UnmarshalMaxMindDB(decoder *mmdbdata.Decoder) error {
 			return iterErr
 		}
 
-		value, err := decodeDataTypeValue(decoder)
+		value, err := decodeDataTypeValue(decoder, cache)
 		if err != nil {
 			return err
 		}
@@ -939,63 +953,98 @@ func writeCtrlByte(w writer, t DataType) (int64, error) {
 	return numBytes, nil
 }
 
+// isCacheableKind returns true if the given kind is worth caching. Currently,
+// we have primarily found a benefit with containers, not scalar values.
+// Potentially, longer strings may also benefit from caching, although it
+// may be even better to just intern them, either here or directly in
+// the maxminddb reader.
+func isCacheableKind(kind mmdbdata.Kind) bool {
+	return kind == mmdbdata.KindMap || kind == mmdbdata.KindSlice
+}
+
 // decodeDataTypeValue decodes a value from the decoder and returns the appropriate DataType.
-func decodeDataTypeValue(decoder *mmdbdata.Decoder) (DataType, error) {
+// If cache is provided (non-nil), it will check for cached values at the current decoder offset
+// and store newly decoded container types (Map, Slice) in the cache. Simple scalar types
+// are not cached as they are cheap to decode and caching them would waste memory.
+func decodeDataTypeValue(decoder *mmdbdata.Decoder, cache map[uint]DataType) (DataType, error) {
 	kind, err := decoder.PeekKind()
 	if err != nil {
 		return nil, fmt.Errorf("peeking kind: %w", err)
 	}
 
+	// Only check cache if provided and the type is worth caching
+	// This avoids unnecessary map lookups for scalar types in tight loops
+	useCache := cache != nil && isCacheableKind(kind)
+	var offset uint
+	if useCache {
+		offset = decoder.Offset()
+		if cached, ok := cache[offset]; ok {
+			return cached, nil
+		}
+	}
+
+	var value DataType
 	switch kind {
 	case mmdbdata.KindString:
-		var value String
-		err := value.UnmarshalMaxMindDB(decoder)
-		return value, err
+		var v String
+		err = v.UnmarshalMaxMindDB(decoder)
+		value = v
 	case mmdbdata.KindFloat64:
-		var value Float64
-		err := value.UnmarshalMaxMindDB(decoder)
-		return value, err
+		var v Float64
+		err = v.UnmarshalMaxMindDB(decoder)
+		value = v
 	case mmdbdata.KindBytes:
-		var value Bytes
-		err := value.UnmarshalMaxMindDB(decoder)
-		return value, err
+		var v Bytes
+		err = v.UnmarshalMaxMindDB(decoder)
+		value = v
 	case mmdbdata.KindUint16:
-		var value Uint16
-		err := value.UnmarshalMaxMindDB(decoder)
-		return value, err
+		var v Uint16
+		err = v.UnmarshalMaxMindDB(decoder)
+		value = v
 	case mmdbdata.KindUint32:
-		var value Uint32
-		err := value.UnmarshalMaxMindDB(decoder)
-		return value, err
+		var v Uint32
+		err = v.UnmarshalMaxMindDB(decoder)
+		value = v
 	case mmdbdata.KindInt32:
-		var value Int32
-		err := value.UnmarshalMaxMindDB(decoder)
-		return value, err
+		var v Int32
+		err = v.UnmarshalMaxMindDB(decoder)
+		value = v
 	case mmdbdata.KindUint64:
-		var value Uint64
-		err := value.UnmarshalMaxMindDB(decoder)
-		return value, err
+		var v Uint64
+		err = v.UnmarshalMaxMindDB(decoder)
+		value = v
 	case mmdbdata.KindUint128:
-		var value Uint128
-		err := value.UnmarshalMaxMindDB(decoder)
-		return &value, err // Return pointer for Uint128
+		var v Uint128
+		err = v.UnmarshalMaxMindDB(decoder)
+		value = &v // Return pointer for Uint128
 	case mmdbdata.KindBool:
-		var value Bool
-		err := value.UnmarshalMaxMindDB(decoder)
-		return value, err
+		var v Bool
+		err = v.UnmarshalMaxMindDB(decoder)
+		value = v
 	case mmdbdata.KindFloat32:
-		var value Float32
-		err := value.UnmarshalMaxMindDB(decoder)
-		return value, err
+		var v Float32
+		err = v.UnmarshalMaxMindDB(decoder)
+		value = v
 	case mmdbdata.KindMap:
-		var value Map
-		err := value.UnmarshalMaxMindDB(decoder)
-		return value, err
+		var v Map
+		err = v.unmarshalMaxMindDB(decoder, cache)
+		value = v
 	case mmdbdata.KindSlice:
-		var value Slice
-		err := value.UnmarshalMaxMindDB(decoder)
-		return value, err
+		var v Slice
+		err = v.unmarshalMaxMindDB(decoder, cache)
+		value = v
 	default:
 		return nil, fmt.Errorf("unsupported data type: %v", kind)
 	}
+
+	if err != nil {
+		return nil, err
+	}
+
+	// Store the decoded value in cache.
+	if useCache {
+		cache[offset] = value
+	}
+
+	return value, nil
 }

mmdbtype/unmarshaler.go

@@ -7,12 +7,23 @@ import (
 // Unmarshaler implements the mmdbdata.Unmarshaler interface for converting
 // MMDB data back into mmdbtype.DataType values. This is used when loading
 // existing MMDB files to reconstruct the original data structures.
+//
+// The Unmarshaler caches decoded complex types (Map, Slice, Uint128) at all
+// nesting levels to improve performance when loading databases with shared
+// nested data structures. Simple scalar types are not cached as they are
+// cheap to decode.
+//
+// The zero value of Unmarshaler is safe to use and will unmarshal data
+// without caching. Use NewUnmarshaler() to create an Unmarshaler with
+// caching enabled for better performance when loading full databases.
 type Unmarshaler struct {
 	cache  map[uint]DataType
 	result DataType
 }
 
-// NewUnmarshaler creates a new Unmarshaler for converting MMDB data to mmdbtype values.
+// NewUnmarshaler creates a new Unmarshaler with caching enabled for converting
+// MMDB data to mmdbtype values. The cache improves performance when loading
+// databases with shared data structures by avoiding redundant decoding.
 func NewUnmarshaler() *Unmarshaler {
 	return &Unmarshaler{
 		cache: map[uint]DataType{},
@@ -21,18 +32,11 @@ func NewUnmarshaler() *Unmarshaler {
 
 // UnmarshalMaxMindDB implements the mmdbdata.Unmarshaler interface.
 func (u *Unmarshaler) UnmarshalMaxMindDB(decoder *mmdbdata.Decoder) error {
-	offset := decoder.Offset()
-	if cached, ok := u.cache[offset]; ok {
-		u.result = cached
-		return nil
-	}
-
-	value, err := decodeDataTypeValue(decoder)
+	value, err := decodeDataTypeValue(decoder, u.cache)
 	if err != nil {
 		return err
 	}
 
-	u.cache[offset] = value
 	u.result = value
 	return nil
 }

mmdbtype/unmarshaler_cache_test.go

@@ -0,0 +1,61 @@
+package mmdbtype
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/oschwald/maxminddb-golang/v2/mmdbdata"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestUnmarshalerCache verifies that the Unmarshaler has a cache field
+// and that NewUnmarshaler initializes it correctly.
+func TestUnmarshalerCache(t *testing.T) {
+	u := NewUnmarshaler()
+	assert.NotNil(t, u)
+
+	// The cache should be initialized (non-nil)
+	assert.NotNil(t, u.cache)
+}
+
+// TestZeroUnmarshaler verifies that the zero value of Unmarshaler is safe to use
+// and correctly unmarshals data without caching enabled.
+func TestZeroUnmarshaler(t *testing.T) {
+	// Create test data with nested structures
+	testData := Map{
+		"outer": Map{
+			"inner": Slice{
+				String("value1"),
+				String("value2"),
+				Map{
+					"deep": String("nested"),
+					"num":  Uint32(42),
+				},
+			},
+			"other": Uint64(100),
+		},
+		"simple": String("test"),
+	}
+
+	// Encode to bytes
+	var buf bytes.Buffer
+	dw := &dataWriter{Buffer: &buf}
+	_, err := testData.WriteTo(dw)
+	require.NoError(t, err)
+
+	var zeroUnmarshaler Unmarshaler
+	decoder := mmdbdata.NewDecoder(buf.Bytes(), 0)
+	err = zeroUnmarshaler.UnmarshalMaxMindDB(decoder)
+	require.NoError(t, err)
+
+	decoded := zeroUnmarshaler.Result()
+	require.NotNil(t, decoded)
+
+	// Verify the decoded data matches the original
+	assert.True(t, testData.Equal(decoded), "decoded data should match original")
+
+	// Verify we can call Clear on zero value (should be safe)
+	zeroUnmarshaler.Clear()
+	assert.Nil(t, zeroUnmarshaler.Result())
+}