WIP

Author: oschwald

internal/decoder/decoder.go

@@ -12,8 +12,14 @@ import (
 	"github.com/oschwald/maxminddb-golang/v2/internal/mmdberrors"
 )
 
+type Cacher interface {
+	Load(key uint) (any, bool)
+	Store(key uint, value any)
+}
+
 type Decoder struct {
 	buffer []byte
+	cache  Cacher
 }
 
 type dataType int
@@ -44,8 +50,8 @@ const (
 	maximumDataStructureDepth = 512
 )
 
-func New(buffer []byte) Decoder {
-	return Decoder{buffer: buffer}
+func New(buffer []byte, cache Cacher) Decoder {
+	return Decoder{buffer: buffer, cache: cache}
 }
 
 func (d *Decoder) Decode(offset uint, v any) error {
@@ -348,7 +354,10 @@ func (d *Decoder) decodeFromTypeToDeserializer(
 		v, offset := d.decodeInt(size, offset)
 		return offset, dser.Int32(int32(v))
 	case _String:
-		v, offset := d.decodeString(size, offset)
+		v, offset, err := d.decodeString(size, offset)
+		if err != nil {
+			return 0, err
+		}
 		return offset, dser.String(v)
 	case _Uint16:
 		v, offset := d.decodeUint(size, offset)
@@ -581,7 +590,10 @@ func (d *Decoder) unmarshalSlice(
 }
 
 func (d *Decoder) unmarshalString(size, offset uint, result reflect.Value) (uint, error) {
-	value, newOffset := d.decodeString(size, offset)
+	value, newOffset, err := d.decodeString(size, offset)
+	if err != nil {
+		return 0, err
+	}
 
 	switch result.Kind() {
 	case reflect.String:
@@ -838,9 +850,25 @@ func (d *Decoder) decodeSliceToDeserializer(
 	return offset, nil
 }
 
-func (d *Decoder) decodeString(size, offset uint) (string, uint) {
+func (d *Decoder) decodeString(size, offset uint) (string, uint, error) {
 	newOffset := offset + size
-	return string(d.buffer[offset:newOffset]), newOffset
+	if d.cache == nil {
+		return string(d.buffer[offset:newOffset]), newOffset, nil
+	}
+
+	v, ok := d.cache.Load(offset)
+	if ok {
+		if s, ok := v.(string); ok {
+			return s, newOffset, nil
+		}
+		return "", 0, mmdberrors.NewCacheTypeStrError(v, "string")
+	}
+
+	s := string(d.buffer[offset:newOffset])
+
+	d.cache.Store(offset, s)
+
+	return s, newOffset, nil
 }
 
 func (d *Decoder) decodeStruct(

internal/mmdberrors/errors.go

@@ -23,6 +23,22 @@ func (e InvalidDatabaseError) Error() string {
 	return e.message
 }
 
+type CacheTypeError struct {
+	Type  string
+	Value any
+}
+
+func NewCacheTypeStrError(value any, expType string) CacheTypeError {
+	return CacheTypeError{
+		Type:  expType,
+		Value: value,
+	}
+}
+
+func (e CacheTypeError) Error() string {
+	return fmt.Sprintf("maxminddb: expected %s type in cache but found %T", e.Type, e.Value)
+}
+
 // UnmarshalTypeError is returned when the value in the database cannot be
 // assigned to the specified data type.
 type UnmarshalTypeError struct {

reader.go

@@ -50,10 +50,18 @@ type Metadata struct {
 	RecordSize               uint              `maxminddb:"record_size"`
 }
 
-type readerOptions struct{}
+type readerOptions struct {
+	cache decoder.Cacher
+}
 
 type ReaderOption func(*readerOptions)
 
+func Cache(cache decoder.Cacher) ReaderOption {
+	return func(o *readerOptions) {
+		o.cache = cache
+	}
+}
+
 // Open takes a string path to a MaxMind DB file and any options. It returns a
 // Reader structure or an error. The database file is opened using a memory
 // map on supported platforms. On platforms without memory map support, such
@@ -143,7 +151,7 @@ func FromBytes(buffer []byte, options ...ReaderOption) (*Reader, error) {
 	}
 
 	metadataStart += len(metadataStartMarker)
-	metadataDecoder := decoder.New(buffer[metadataStart:])
+	metadataDecoder := decoder.New(buffer[metadataStart:], nil)
 
 	var metadata Metadata
 
@@ -159,7 +167,8 @@ func FromBytes(buffer []byte, options ...ReaderOption) (*Reader, error) {
 		return nil, mmdberrors.NewInvalidDatabaseError("the MaxMind DB contains invalid metadata")
 	}
 	d := decoder.New(
-		buffer[searchTreeSize+dataSectionSeparatorSize : metadataStart-len(metadataStartMarker)],
+		buffer[searchTreeSize+dataSectionSeparatorSize:metadataStart-len(metadataStartMarker)],
+		opts.cache,
 	)
 
 	nodeBuffer := buffer[:searchTreeSize]

reader_test.go

@@ -931,6 +931,44 @@ func BenchmarkCityLookup(b *testing.B) {
 	var result fullCity
 
 	s := make(net.IP, 4)
+
+	b.ResetTimer()
+
+	for range b.N {
+		ip := randomIPv4Address(r, s)
+		err = db.Lookup(ip).Decode(&result)
+		if err != nil {
+			b.Error(err)
+		}
+	}
+	require.NoError(b, db.Close(), "error on close")
+}
+
+type cache struct {
+	m map[uint]any
+}
+
+func (c *cache) Load(key uint) (any, bool) {
+	v, ok := c.m[key]
+	return v, ok
+}
+
+func (c *cache) Store(key uint, value any) {
+	c.m[key] = value
+}
+
+func BenchmarkCityLookupWithCache(b *testing.B) {
+	db, err := Open("GeoLite2-City.mmdb", Cache(&cache{m: map[uint]any{}}))
+	require.NoError(b, err)
+
+	//nolint:gosec // this is a test
+	r := rand.New(rand.NewSource(time.Now().UnixNano()))
+	var result fullCity
+
+	s := make(net.IP, 4)
+
+	b.ResetTimer()
+
 	for range b.N {
 		ip := randomIPv4Address(r, s)
 		err = db.Lookup(ip).Decode(&result)
@@ -949,6 +987,9 @@ func BenchmarkCityLookupOnly(b *testing.B) {
 	r := rand.New(rand.NewSource(time.Now().UnixNano()))
 
 	s := make(net.IP, 4)
+
+	b.ResetTimer()
+
 	for range b.N {
 		ip := randomIPv4Address(r, s)
 		result := db.Lookup(ip)