feat: lru store (#43)

Author: scorix

go.mod

@@ -10,6 +10,7 @@ require (
 	github.com/scorix/go-eccodes v0.1.5
 	github.com/scorix/walg v0.4.0
 	github.com/stretchr/testify v1.9.0
+	go.opentelemetry.io/otel v1.32.0
 	go.opentelemetry.io/otel/trace v1.32.0
 	golang.org/x/exp v0.0.0-20240904232852-e7e105dedf7e
 	golang.org/x/sync v0.9.0
@@ -19,7 +20,6 @@ require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/kr/text v0.2.0 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	go.opentelemetry.io/otel v1.32.0 // indirect
 	golang.org/x/time v0.7.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

pkg/grib2/cache/custom_test.go

@@ -0,0 +1,37 @@
+package cache_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/scorix/grib-go/pkg/grib2/cache"
+)
+
+// 实现一个简单的 GridDataSource
+type testDataSource struct{}
+
+func (t *testDataSource) ReadGridAt(ctx context.Context, index int) (float32, error) {
+	return 1.0, nil
+}
+
+func BenchmarkCustom_ReadGridAt_Parallel(b *testing.B) {
+	store := cache.NewMapStore()
+
+	c := cache.NewCustom(
+		func(lat, lon float32) bool {
+			return true
+		},
+		&testDataSource{},
+		store,
+	)
+
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			_, err := c.ReadGridAt(context.Background(), 1, 1, 1)
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+}

pkg/grib2/cache/store.go

@@ -1,6 +1,7 @@
 package cache
 
 import (
+	"container/list"
 	"context"
 	"sync"
 
@@ -13,6 +14,68 @@ type Store interface {
 	Set(ctx context.Context, key int, value float32)
 }
 
+// LRU Cache 实现
+type lruStore struct {
+	mu       sync.RWMutex
+	capacity int
+	cache    map[int]*list.Element
+	lru      *list.List
+}
+
+type entry struct {
+	key   int
+	value float32
+}
+
+func NewLRUStore(capacity int) Store {
+	return &lruStore{
+		capacity: capacity,
+		cache:    make(map[int]*list.Element),
+		lru:      list.New(),
+	}
+}
+
+func (l *lruStore) Get(ctx context.Context, key int) (float32, bool) {
+	l.mu.RLock()
+	defer l.mu.RUnlock()
+
+	if elem, ok := l.cache[key]; ok {
+		l.lru.MoveToFront(elem)
+		sp := trace.SpanFromContext(ctx)
+		sp.SetAttributes(attribute.Int("cache.grid", key), attribute.Bool("cache.hit", true))
+		return elem.Value.(*entry).value, true
+	}
+
+	sp := trace.SpanFromContext(ctx)
+	sp.SetAttributes(attribute.Int("cache.grid", key), attribute.Bool("cache.hit", false))
+	return 0, false
+}
+
+func (l *lruStore) Set(ctx context.Context, key int, value float32) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	if elem, ok := l.cache[key]; ok {
+		l.lru.MoveToFront(elem)
+		elem.Value.(*entry).value = value
+		return
+	}
+
+	// 如果缓存已满，删除最久未使用的条目
+	if l.lru.Len() >= l.capacity {
+		oldest := l.lru.Back()
+		if oldest != nil {
+			delete(l.cache, oldest.Value.(*entry).key)
+			l.lru.Remove(oldest)
+		}
+	}
+
+	// 添加新条目
+	elem := l.lru.PushFront(&entry{key: key, value: value})
+	l.cache[key] = elem
+}
+
+// 简单的 map 实现
 type mapStore struct {
 	mu    sync.RWMutex
 	cache map[int]float32

pkg/grib2/cache/store_test.go

@@ -0,0 +1,53 @@
+package cache
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func BenchmarkStore_Parallel(b *testing.B) {
+	b.Run("LRU", func(b *testing.B) {
+		store := NewLRUStore(1000)
+		ctx := context.Background()
+
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				store.Set(ctx, 1, 1.0)
+				_, _ = store.Get(ctx, 1)
+			}
+		})
+	})
+
+	b.Run("Map", func(b *testing.B) {
+		store := NewMapStore()
+		ctx := context.Background()
+
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				store.Set(ctx, 1, 1.0)
+				_, _ = store.Get(ctx, 1)
+			}
+		})
+	})
+}
+
+func TestLRUStore(t *testing.T) {
+	store := NewLRUStore(1)
+	ctx := context.Background()
+
+	store.Set(ctx, 1, 1.0)
+	v, ok := store.Get(ctx, 1)
+	assert.True(t, ok)
+	assert.Equal(t, float32(1.0), v)
+
+	store.Set(ctx, 2, 2.0)
+	v, ok = store.Get(ctx, 2)
+	assert.True(t, ok)
+	assert.Equal(t, float32(2.0), v)
+
+	v, ok = store.Get(ctx, 3)
+	assert.False(t, ok)
+	assert.Equal(t, float32(0.0), v)
+}