diff --git a/go.mod b/go.mod index 440073b..97f3722 100644 --- a/go.mod +++ b/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 ) diff --git a/pkg/grib2/cache/custom_test.go b/pkg/grib2/cache/custom_test.go new file mode 100644 index 0000000..1cfbd71 --- /dev/null +++ b/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) + } + } + }) +} diff --git a/pkg/grib2/cache/store.go b/pkg/grib2/cache/store.go index 7f837cb..19b10e8 100644 --- a/pkg/grib2/cache/store.go +++ b/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 diff --git a/pkg/grib2/cache/store_test.go b/pkg/grib2/cache/store_test.go new file mode 100644 index 0000000..16f1503 --- /dev/null +++ b/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) +}