From f3b182e5bff8d4791f02f16a6ac1143d7a3ce97d Mon Sep 17 00:00:00 2001
From: Andrew Regner <andrew@aregner.com>
Date: Sun, 23 Jun 2019 22:27:07 -0700
Subject: [PATCH] Add OnMissing callback

---
 CONTRIBUTORS  |  1 +
 cache.go      | 19 +++++++++++++++++++
 cache_test.go | 36 ++++++++++++++++++++++++++++++++++++
 3 files changed, 56 insertions(+)

diff --git a/CONTRIBUTORS b/CONTRIBUTORS
index 2b16e99..0a30517 100644
--- a/CONTRIBUTORS
+++ b/CONTRIBUTORS
@@ -7,3 +7,4 @@ Dustin Sallings <dustin@spy.net>
 Jason Mooberry <jasonmoo@me.com>
 Sergey Shepelev <temotor@gmail.com>
 Alex Edwards <ajmedwards@gmail.com>
+Andrew Regner <andrew@aregner.com>
diff --git a/cache.go b/cache.go
index db88d2f..25273b0 100644
--- a/cache.go
+++ b/cache.go
@@ -42,6 +42,7 @@ type cache struct {
 	items             map[string]Item
 	mu                sync.RWMutex
 	onEvicted         func(string, interface{})
+	onMissing         func(string) (*Item, error)
 	janitor           *janitor
 }
 
@@ -123,6 +124,16 @@ func (c *cache) Get(k string) (interface{}, bool) {
 	item, found := c.items[k]
 	if !found {
 		c.mu.RUnlock()
+		// try to generate the missing value
+		if c.onMissing != nil {
+			item, err := c.onMissing(k)
+			if err == nil {
+				c.mu.Lock()
+				c.items[k] = *item
+				c.mu.Unlock()
+				return item.Object, true
+			}
+		}
 		return nil, false
 	}
 	if item.Expiration > 0 {
@@ -956,6 +967,14 @@ func (c *cache) OnEvicted(f func(string, interface{})) {
 	c.mu.Unlock()
 }
 
+// Sets an (optional) function that is called to optionally generate a value for
+// a key that was not found when calling Get.
+func (c *cache) OnMissing(f func(string) (*Item, error)) {
+	c.mu.Lock()
+	c.onMissing = f
+	c.mu.Unlock()
+}
+
 // Write the cache's items (using Gob) to an io.Writer.
 //
 // NOTE: This method is deprecated in favor of c.Items() and NewFrom() (see the
diff --git a/cache_test.go b/cache_test.go
index cb80b38..3aa347f 100644
--- a/cache_test.go
+++ b/cache_test.go
@@ -2,6 +2,7 @@ package cache
 
 import (
 	"bytes"
+	"fmt"
 	"io/ioutil"
 	"runtime"
 	"strconv"
@@ -1247,6 +1248,41 @@ func TestOnEvicted(t *testing.T) {
 	}
 }
 
+func TestOnMissing(t *testing.T) {
+	tc := New(DefaultExpiration, 0)
+	tc.Set("foo", 3, DefaultExpiration)
+	if tc.onMissing != nil {
+		t.Fatal("tc.onMissing is not nil")
+	}
+	tc.OnMissing(func(k string) (*Item, error) {
+		return &Item{
+			Object:     42,
+			Expiration: 0,
+		}, nil
+	})
+	x, _ := tc.Get("foo")
+	if x != 3 {
+		t.Error("getting a set item with a generator set produced unexpected result:", x)
+	}
+	x, found := tc.Get("potato")
+	if !found {
+		t.Error("item was not generated for missing key")
+	} else if x.(int) != 42 {
+		t.Error("generated item was not expected; value:", x)
+	}
+	tc.OnMissing(func(k string) (*Item, error) {
+		return nil, fmt.Errorf("some error")
+	})
+	x, found = tc.Get("apples")
+	if found || x != nil {
+		t.Error("onMissing should not have generated this:", x)
+	}
+	x, found = tc.Get("potato")
+	if !found || x != 42 {
+		t.Error("prior generated value should still have been set")
+	}
+}
+
 func TestCacheSerialization(t *testing.T) {
 	tc := New(DefaultExpiration, 0)
 	testFillAndSerialize(t, tc)