From f7237ae7a6ebee22a5f4643985878ba072e0ae85 Mon Sep 17 00:00:00 2001 From: AhmedGoudaa Date: Wed, 2 Sep 2026 14:54:46 +0400 Subject: [PATCH] fix: release flights after compute panic --- .agents/skills/go-memoize-package/SKILL.md | 1 + cache.go | 44 +++++++++++++++++----- cache_engine_test.go | 42 +++++++++++++++++++++ cache_flight_test.go | 22 +++++++++++ docs/API.md | 2 + 5 files changed, 102 insertions(+), 9 deletions(-) create mode 100644 cache_flight_test.go diff --git a/.agents/skills/go-memoize-package/SKILL.md b/.agents/skills/go-memoize-package/SKILL.md index 86f3749..90f3961 100644 --- a/.agents/skills/go-memoize-package/SKILL.md +++ b/.agents/skills/go-memoize-package/SKILL.md @@ -48,6 +48,7 @@ Choose the API and store before editing: - Default clocks and clocks created by `WithTickerClock` are cache-owned and stopped by `Cache.Stop`; clocks injected with `WithClock` are caller-owned and may be shared across caches. - Metrics use one public method: `RecordMetric(memoize.MetricEvent)`. - `background.Keep` and `loader.New` share internal periodic refresh-loop infrastructure. Cache stale refresh uses cache flight machinery, not the shared refresh loop. +- A foreground `GetOrCompute` panic is propagated to the leader and same-flight followers. The flight must be removed without caching a result so later calls can retry. - Source compatibility may change for performance or clarity; users can pin module versions. - Keep core code standard-library-only unless the user explicitly approves a dependency. diff --git a/cache.go b/cache.go index 8aa955a..a7d85ba 100644 --- a/cache.go +++ b/cache.go @@ -8,9 +8,19 @@ import ( ) type flight[V any] struct { - wg sync.WaitGroup - value V - err error + wg sync.WaitGroup + value V + err error + panicValue any + panicked bool +} + +func (f *flight[V]) wait() (V, error) { + f.wg.Wait() + if f.panicked { + panic(f.panicValue) + } + return f.value, f.err } // peekingStore lets GetOrCompute inspect stored entry state without applying @@ -112,8 +122,8 @@ func (c *Cache[K, V]) waitForFlight(key K) (V, bool, error) { var zero V return zero, false, nil } - existing.wg.Wait() - return existing.value, true, existing.err + value, err := existing.wait() + return value, true, err } func (c *Cache[K, V]) startFlight(key K) (*flight[V], bool) { @@ -132,21 +142,37 @@ func (c *Cache[K, V]) startFlight(key K) (*flight[V], bool) { func (c *Cache[K, V]) finishFlight(key K, f *flight[V], value V, err error) { f.value = value f.err = err + + c.flightMu.Lock() + delete(c.flights, key) + c.flightMu.Unlock() f.wg.Done() +} + +func (c *Cache[K, V]) finishPanickedFlight(key K, f *flight[V], panicValue any) { + f.panicValue = panicValue + f.panicked = true c.flightMu.Lock() delete(c.flights, key) c.flightMu.Unlock() + f.wg.Done() } -func (c *Cache[K, V]) do(key K, fn func() (V, error)) (V, error) { +func (c *Cache[K, V]) do(key K, fn func() (V, error)) (value V, err error) { f, leader := c.startFlight(key) if !leader { - f.wg.Wait() - return f.value, f.err + return f.wait() } - value, err := fn() + defer func() { + if panicValue := recover(); panicValue != nil { + c.finishPanickedFlight(key, f, panicValue) + panic(panicValue) + } + }() + + value, err = fn() c.finishFlight(key, f, value, err) return value, err } diff --git a/cache_engine_test.go b/cache_engine_test.go index 61c75b0..7c6025a 100644 --- a/cache_engine_test.go +++ b/cache_engine_test.go @@ -424,3 +424,45 @@ func TestConcurrentMissComputesOnce(t *testing.T) { t.Fatalf("expected 1 compute call, got %d", calls) } } + +func TestGetOrComputePanicDoesNotPoisonKey(t *testing.T) { + cache, err := memoize.New[string, int]( + memoize.Opts().WithStore(memory.New[string, int](8)).WithTTL(time.Minute), + ) + if err != nil { + t.Fatalf("new cache failed: %v", err) + } + t.Cleanup(cache.Stop) + + func() { + defer func() { + if got := recover(); got != "boom" { + t.Fatalf("recovered panic = %v, want boom", got) + } + }() + _, _ = cache.GetOrCompute(t.Context(), "key", func(context.Context) (int, error) { + panic("boom") + }) + }() + + type result struct { + value int + err error + } + done := make(chan result, 1) + go func() { + value, err := cache.GetOrCompute(t.Context(), "key", func(context.Context) (int, error) { + return 42, nil + }) + done <- result{value: value, err: err} + }() + + select { + case got := <-done: + if got.err != nil || got.value != 42 { + t.Fatalf("retry = (%d, %v), want (42, nil)", got.value, got.err) + } + case <-time.After(time.Second): + t.Fatal("retry blocked on a poisoned single-flight") + } +} diff --git a/cache_flight_test.go b/cache_flight_test.go new file mode 100644 index 0000000..a400102 --- /dev/null +++ b/cache_flight_test.go @@ -0,0 +1,22 @@ +package memoize + +import ( + "testing" +) + +func TestFlightWaitRepanicsWithOriginalValue(t *testing.T) { + sentinel := &struct{ message string }{message: "boom"} + f := &flight[int]{ + panicValue: sentinel, + panicked: true, + } + f.wg.Add(1) + f.wg.Done() + + defer func() { + if got := recover(); got != sentinel { + t.Fatalf("recovered panic = %v, want identical sentinel", got) + } + }() + _, _ = f.wait() +} diff --git a/docs/API.md b/docs/API.md index b74d962..fa9aedb 100644 --- a/docs/API.md +++ b/docs/API.md @@ -124,6 +124,8 @@ defer cache.Stop() The cache engine owns freshness decisions. Stores persist `memoize.Stored[V]` envelopes and return entries even when they might be stale or expired; `Cache[K,V]` decides whether to serve, refresh, or miss. +If a foreground `GetOrCompute` function panics, the leader and all callers waiting on that same-key computation panic with the same value. The failed flight is removed without caching a result, so a later call can retry normally. + `GetOrCompute` example: ```go