diff --git a/bicache.go b/bicache.go index 1c55db5..b9f686c 100644 --- a/bicache.go +++ b/bicache.go @@ -69,6 +69,7 @@ type Config struct { EvictLog bool ShardCount int NoOverflow bool + Context context.Context } // Entry is a container type for scored @@ -141,7 +142,10 @@ func New(c *Config) (*Bicache, error) { } } - ctx, cf := context.WithCancel(context.Background()) + if c.Context == nil { + c.Context = context.Background() + } + ctx, cf := context.WithCancel(c.Context) cache := &Bicache{ shards: shards, diff --git a/bicache_test.go b/bicache_test.go index 994cdf2..c1d7e88 100644 --- a/bicache_test.go +++ b/bicache_test.go @@ -1,6 +1,8 @@ package bicache_test import ( + "context" + "fmt" "log" "strconv" "testing" @@ -222,3 +224,43 @@ func TestPromoteEvict(t *testing.T) { } } } + +func TestCanceledContext(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + c, _ := bicache.New(&bicache.Config{ + MRUSize: 30, + ShardCount: 2, + AutoEvict: 1000, + Context: ctx, + }) + + for i := 0; i < 40; i++ { + c.Set(fmt.Sprintf("initial-%d", i), "value") + } + + log.Printf("Sleeping for 2 seconds to allow evictions") + time.Sleep(2 * time.Second) + + stats := c.Stats() + + // Check that auto evict is maintaining capacity + if stats.MRUSize != 30 { + t.Error("Unexpected initial MRU count") + } + + // Cancel context to close auto evict + cancel() + + for i := 0; i < 40; i++ { + c.Set(fmt.Sprintf("after-cancel-%d", i), "value") + } + log.Printf("Sleeping for 2 seconds to allow evictions") + time.Sleep(2 * time.Second) + + stats = c.Stats() + + // Check that auto evict is not active and not maintaining capacity + if stats.MRUSize != 70 { + t.Error("Unexpected MRU count after cancel") + } +}