@@ -22,9 +22,9 @@ import (
2222type Interface interface {
2323 Load (context.Context , string ) Thunk
2424 LoadMany (context.Context , []string ) ThunkMany
25- Clear (string ) Interface
25+ Clear (context. Context , string ) Interface
2626 ClearAll () Interface
27- Prime (key string , value interface {}) Interface
27+ Prime (ctx context. Context , key string , value interface {}) Interface
2828}
2929
3030// BatchFunc is a function, which when given a slice of keys (string), returns an slice of `results`.
@@ -193,6 +193,7 @@ func NewBatchedLoader(batchFn BatchFunc, opts ...Option) *Loader {
193193// Load load/resolves the given key, returning a channel that will contain the value and error
194194func (l * Loader ) Load (originalContext context.Context , key string ) Thunk {
195195 ctx , finish := l .tracer .TraceLoad (originalContext , key )
196+
196197 c := make (chan * Result , 1 )
197198 var result struct {
198199 mu sync.RWMutex
@@ -201,7 +202,7 @@ func (l *Loader) Load(originalContext context.Context, key string) Thunk {
201202
202203 // lock to prevent duplicate keys coming in before item has been added to cache.
203204 l .cacheLock .Lock ()
204- if v , ok := l .cache .Get (key ); ok {
205+ if v , ok := l .cache .Get (ctx , key ); ok {
205206 defer finish (v )
206207 defer l .cacheLock .Unlock ()
207208 return v
@@ -223,8 +224,9 @@ func (l *Loader) Load(originalContext context.Context, key string) Thunk {
223224 defer result .mu .RUnlock ()
224225 return result .value .Data , result .value .Error
225226 }
227+ defer finish (thunk )
226228
227- l .cache .Set (key , thunk )
229+ l .cache .Set (ctx , key , thunk )
228230 l .cacheLock .Unlock ()
229231
230232 // this is sent to batch fn. It contains the key and the channel to return the
@@ -236,7 +238,7 @@ func (l *Loader) Load(originalContext context.Context, key string) Thunk {
236238 if l .curBatcher == nil {
237239 l .curBatcher = l .newBatcher (l .silent , l .tracer )
238240 // start the current batcher batch function
239- go l .curBatcher .batch (ctx )
241+ go l .curBatcher .batch (originalContext )
240242 // start a sleeper for the current batcher
241243 l .endSleeper = make (chan bool )
242244 go l .sleeper (l .curBatcher , l .endSleeper )
@@ -261,13 +263,13 @@ func (l *Loader) Load(originalContext context.Context, key string) Thunk {
261263 }
262264 l .batchLock .Unlock ()
263265
264- defer finish (thunk )
265266 return thunk
266267}
267268
268269// LoadMany loads mulitiple keys, returning a thunk (type: ThunkMany) that will resolve the keys passed in.
269270func (l * Loader ) LoadMany (originalContext context.Context , keys []string ) ThunkMany {
270271 ctx , finish := l .tracer .TraceLoadMany (originalContext , keys )
272+
271273 length := len (keys )
272274 data := make ([]interface {}, length )
273275 errors := make ([]error , length )
@@ -276,13 +278,13 @@ func (l *Loader) LoadMany(originalContext context.Context, keys []string) ThunkM
276278
277279 wg .Add (length )
278280 for i := range keys {
279- go func (i int ) {
281+ go func (ctx context. Context , i int ) {
280282 defer wg .Done ()
281283 thunk := l .Load (ctx , keys [i ])
282284 result , err := thunk ()
283285 data [i ] = result
284286 errors [i ] = err
285- }(i )
287+ }(ctx , i )
286288 }
287289
288290 go func () {
@@ -318,9 +320,9 @@ func (l *Loader) LoadMany(originalContext context.Context, keys []string) ThunkM
318320}
319321
320322// Clear clears the value at `key` from the cache, it it exsits. Returs self for method chaining
321- func (l * Loader ) Clear (key string ) Interface {
323+ func (l * Loader ) Clear (ctx context. Context , key string ) Interface {
322324 l .cacheLock .Lock ()
323- l .cache .Delete (key )
325+ l .cache .Delete (ctx , key )
324326 l .cacheLock .Unlock ()
325327 return l
326328}
@@ -336,12 +338,12 @@ func (l *Loader) ClearAll() Interface {
336338
337339// Prime adds the provided key and value to the cache. If the key already exists, no change is made.
338340// Returns self for method chaining
339- func (l * Loader ) Prime (key string , value interface {}) Interface {
340- if _ , ok := l .cache .Get (key ); ! ok {
341+ func (l * Loader ) Prime (ctx context. Context , key string , value interface {}) Interface {
342+ if _ , ok := l .cache .Get (ctx , key ); ! ok {
341343 thunk := func () (interface {}, error ) {
342344 return value , nil
343345 }
344- l .cache .Set (key , thunk )
346+ l .cache .Set (ctx , key , thunk )
345347 }
346348 return l
347349}
0 commit comments