Merge pull request #91 from nirrattner/expose-context

Expose context for auto evict lifecycle

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,

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")
+	}
+}