Enhancements -

1) handling potential deadlock if MultiFetch function do not give response for all the keys.
2) Functionality to take cache dump.
3) Preload cache with boot up key id - value pairs.
4) Cache explicit clean up and clean up with context.

Author: saurav534

README.md

@@ -40,7 +40,7 @@ the keys are overlapping as well.
 
 ```go
 func main() {
-	cache, _ := txncache.NewCache(context.Background(), GetValue, GetMultiValue, 10)
+	cache, _ := txncache.NewCache(GetValue, GetMultiValue, 10)
 	var wg sync.WaitGroup
 
 	wg.Add(1)
@@ -91,4 +91,18 @@ func (k ZKey) Id() string {
 }
 ```
 
-For robust example checkout the test file.
+For robust example checkout the test file.
+
+## Dump & Pre Loading
+One might need to reuse the current cache store in the next transaction as well. So **GetAll()** can be used to take
+the dump of current cache storage. Client can store it somewhere and case use it later to preload.
+**Preload(map[string]Value)** preload a cache with given key id and value map and do not call fetch function for those keys
+
+## Default Value
+Ideally the *Fetch* & *MultiFetch* function should return value for all the provided keys. However if the implementation 
+does not guarantee that, one should provide a default value to be used as return value of such keys. This is compulsory to
+avoid deadlock situation while waiting for these keys.
+
+## Cache cleanup
+Cache should be closed post usage to avoid leaks, Close up can be done explicitly by calling **Close()** function or
+can be closed along with ctx cancellation by using **CloseWithCtx(context.Context)** function.

cache.go

@@ -7,13 +7,16 @@ import (
 	"sync"
 )
 
-type TxnCache struct {
+type txnCache struct {
 	store          *sync.Map
 	lock           *sync.Map
 	fetch          Fetch
 	multiFetch     MultiFetch
+	defVal         Value
+	batchSize      int
 	multiKeyChan   chan *keyValChan
 	execMultiFetch chan bool
+	closeChan      chan bool
 }
 
 type keyWrap struct {
@@ -26,11 +29,11 @@ type keyValChan struct {
 	locked bool
 }
 
-// NewCache gives an instance of TxnCache for given fetch & multiFetch functions
+// NewCache gives an instance of txnCache for given fetch & multiFetch functions
 // Fetch is the function to be used for value fetch of a single key
 // MultiFetch is the batched version of fetch function
 // batchSize should be passed greater than 0 if MultiFetch function is expected to be called in batch
-func NewCache(ctx context.Context, fetch Fetch, multiFetch MultiFetch, batchSize int) (*TxnCache, error) {
+func NewCache(fetch Fetch, multiFetch MultiFetch, batchSize int) (Cache, error) {
 	if fetch == nil && multiFetch == nil {
 		return nil, fmt.Errorf("both fetch and multifetch can not be nil")
 	}
@@ -66,22 +69,49 @@ func NewCache(ctx context.Context, fetch Fetch, multiFetch MultiFetch, batchSize
 
 	multiKeyChan := make(chan *keyValChan)
 	execMultiFetch := make(chan bool)
+	closeChan := make(chan bool)
 
 	store := &sync.Map{}
 	lock := &sync.Map{}
 
+	cache := &txnCache{
+		store:          store,
+		lock:           lock,
+		fetch:          fetch,
+		multiFetch:     multiFetch,
+		multiKeyChan:   multiKeyChan,
+		defVal:         struct{}{},
+		batchSize:      batchSize,
+		execMultiFetch: execMultiFetch,
+		closeChan:      closeChan,
+	}
+	cache.setup()
+	return cache, nil
+}
+
+func (tc *txnCache) setup() {
 	uniqueKeys := make([]Key, 0)
 	keysMap := make(map[string][]*keyWrap)
 	resMap := make(map[*keyWrap]chan Value)
 	execute := func() {
-		for k, v := range multiFetch(uniqueKeys) {
-			store.Store(k.Id(), v)
-			if kws, ok := keysMap[k.Id()]; ok {
+		mf := tc.multiFetch(uniqueKeys)
+		for _, k := range uniqueKeys {
+			v := mf[k]
+			// if fetch function do not return any value for a key
+			// ideally fetch function should handle this and return a
+			// suitable value of such keys
+			// default values are not cached
+			if v == nil {
+				v = tc.defVal
+			} else {
+				tc.store.Store(k.ID(), v)
+			}
+			if kws, ok := keysMap[k.ID()]; ok {
 				for _, kw := range kws {
 					resMap[kw] <- v
 					delete(resMap, kw)
 				}
-				delete(keysMap, k.Id())
+				delete(keysMap, k.ID())
 			}
 		}
 		uniqueKeys = nil
@@ -90,80 +120,75 @@ func NewCache(ctx context.Context, fetch Fetch, multiFetch MultiFetch, batchSize
 	go func() {
 		for {
 			select {
-			case kr := <-multiKeyChan:
+			case kr := <-tc.multiKeyChan:
 				{
 					mutex.Lock()
-					if value, ok := store.Load(kr.key.Id()); ok {
+					if value, ok := tc.store.Load(kr.key.ID()); ok {
 						kr.val <- value
 						mutex.Unlock()
 						continue
 					}
 					kw := &keyWrap{kr.key}
-					if keys, ok := keysMap[kr.key.Id()]; ok {
-						keysMap[kr.key.Id()] = append(keys, kw)
+					if keys, ok := keysMap[kr.key.ID()]; ok {
+						keysMap[kr.key.ID()] = append(keys, kw)
 					} else {
 						uniqueKeys = append(uniqueKeys, kr.key)
-						keysMap[kr.key.Id()] = []*keyWrap{kw}
+						keysMap[kr.key.ID()] = []*keyWrap{kw}
 					}
 					resMap[kw] = kr.val
-					if batchSize > 0 && batchSize == len(uniqueKeys) {
+					if tc.batchSize > 0 && tc.batchSize == len(uniqueKeys) {
 						execute()
 					}
 					mutex.Unlock()
 				}
-			case <-execMultiFetch:
+			case <-tc.execMultiFetch:
 				{
 					mutex.Lock()
 					execute()
 					mutex.Unlock()
 				}
-			case <-ctx.Done():
+			case <-tc.closeChan:
 				{
-					close(multiKeyChan)
-					close(execMultiFetch)
+					close(tc.multiKeyChan)
+					close(tc.execMultiFetch)
 					return
 				}
 			}
 		}
 	}()
-	return &TxnCache{
-		store:          store,
-		lock:           lock,
-		fetch:          fetch,
-		multiFetch:     multiFetch,
-		multiKeyChan:   multiKeyChan,
-		execMultiFetch: execMultiFetch,
-	}, nil
 }
 
-func (tc *TxnCache) Get(key Key) Value {
-	if value, ok := tc.store.Load(key.Id()); ok {
+func (tc *txnCache) Get(key Key) Value {
+	if value, ok := tc.store.Load(key.ID()); ok {
 		return value
 	}
-	lock, _ := tc.lock.LoadOrStore(key.Id(), &sync.Mutex{})
+	lock, _ := tc.lock.LoadOrStore(key.ID(), &sync.Mutex{})
 	mutex := lock.(*sync.Mutex)
 	mutex.Lock()
 	defer mutex.Unlock()
-	if value, ok := tc.store.Load(key.Id()); ok {
+	if value, ok := tc.store.Load(key.ID()); ok {
 		return value
 	} else {
 		value := tc.fetch(key)
-		tc.store.Store(key.Id(), value)
+		if value == nil {
+			return tc.defVal
+		}
+		tc.store.Store(key.ID(), value)
 		return value
 	}
 }
 
-func (tc *TxnCache) MultiGet(keys []Key) map[Key]Value {
+func (tc *txnCache) MultiGet(keys []Key) map[Key]Value {
 	var data sync.Map
 	var wg sync.WaitGroup
 	var keyPush sync.WaitGroup
 	for _, key := range keys {
-		if value, ok := tc.store.Load(key.Id()); ok {
+		if value, ok := tc.store.Load(key.ID()); ok {
 			data.Store(key, value)
 			continue
 		}
 		wg.Add(1)
-		lock, _ := tc.lock.LoadOrStore(key.Id(), internal.NewMutex())
+		lock, _ := tc.lock.LoadOrStore(key.ID(), internal.NewMutex())
 		mutex := lock.(*internal.Mutex)
 		locked := mutex.TryLock()
 		keyPush.Add(1)
@@ -185,8 +210,38 @@ func (tc *TxnCache) MultiGet(keys []Key) map[Key]Value {
 	res := make(map[Key]Value)
 	wg.Wait()
 	data.Range(func(key, value interface{}) bool {
-		res[key.(Key)] = value.(Value)
+		res[key.(Key)] = value
 		return true
 	})
 	return res
 }
+
+func (tc *txnCache) GetAll() map[string]Value {
+	res := make(map[string]Value)
+	tc.store.Range(func(key, value interface{}) bool {
+		res[key.(string)] = value
+		return true
+	})
+	return res
+}
+
+func (tc *txnCache) Preload(preload map[string]Value) {
+	for k, v := range preload {
+		tc.store.Store(k, v)
+	}
+}
+
+func (tc *txnCache) DefaultValue(value Value) {
+	tc.defVal = value
+}
+
+func (tc *txnCache) CloseWithCtx(ctx context.Context) {
+	go func() {
+		<-ctx.Done()
+		tc.closeChan <- true
+	}()
+}
+
+func (tc *txnCache) Close() {
+	tc.closeChan <- true
+}

cache_api.go

@@ -1,14 +1,42 @@
 package txncache
 
+import "context"
+
+// Key construct of the cache
+// ID() is used to uniquely identify Keys
 type Key interface {
-	Id() string
+	ID() string
 }
 
+// Value construct of the cache
+// Can be of any type
 type Value interface{}
 
 type Cache interface {
-	Get(key Key) Value
-	MultiGet(keys []Key) map[Key]Value
+	// Get return value for a given key
+	Get(Key) Value
+
+	// MultiGet returns key value map for a given list of keys
+	MultiGet([]Key) map[Key]Value
+
+	// GetAll returns map of all the key id, value present in the cache
+	GetAll() map[string]Value
+
+	// Preload can be used to load cache with a map of key value upfront
+	Preload(map[string]Value)
+
+	// DefaultValue set the default value for the missing values while calling fetch functions
+	// This is mandatory if the provided MultiFetch function does not guarantee value for all
+	// the provided keys, default value is not cached.
+	DefaultValue(Value)
+
+	// CloseWithCtx clean up the cache once the provided context is canceled or done
+	// it's a non blocking function and can be used in place of explicit Close() call for guaranteed
+	// cleanup as soon as ctx is canceled
+	CloseWithCtx(context.Context)
+
+	// Close is necessary to call to cleanup cache and make it eligible for GC
+	Close()
 }
 
 type Fetch func(Key) Value