fix: Additional tests for the in-flight key tracking

Makefile

@@ -18,7 +18,7 @@ test:
  @echo 'Removing test cache...'
  go clean -testcache
  @echo 'Running tests...'
- go test -race -vet=off -timeout 30s ./...
+ go test -race -vet=off -timeout 15s ./...
 
 ## bench: run all benchmarks
 .PHONY: bench

README.md

@@ -1,4 +1,4 @@
-# `sturdyc`: a caching library for building sturdy systemsreadme
+# `sturdyc`: a caching library for building sturdy systems
 
 [![Go Reference](https://pkg.go.dev/badge/github.com/creativecreature/sturdyc.svg)](https://pkg.go.dev/github.com/creativecreature/sturdyc)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://github.com/creativecreature/sturdyc/blob/master/LICENSE)

cache.go

@@ -204,3 +204,12 @@ func (c *Client[T]) Delete(key string) {
  shard := c.getShard(key)
  shard.delete(key)
 }
+
+// NumKeysInflight returns the number of keys that are currently being fetched.
+func (c *Client[T]) NumKeysInflight() int {
+ c.inFlightMutex.Lock()
+ defer c.inFlightMutex.Unlock()
+ c.inFlightBatchMutex.Lock()
+ defer c.inFlightBatchMutex.Unlock()
+ return len(c.inFlightMap) + len(c.inFlightBatchMap)
+}

inflight.go

@@ -3,6 +3,7 @@ package sturdyc
 import (
  "context"
  "errors"
+ "fmt"
  "sync"
 )
 
@@ -20,6 +21,54 @@ func (c *Client[T]) newFlight(key string) *inFlightCall[T] {
  return call
 }
 
+func makeCall[T, V any](ctx context.Context, c *Client[T], key string, fn FetchFn[V], call *inFlightCall[T]) {
+ defer func() {
+ if err := recover(); err != nil {
+ call.err = fmt.Errorf("panic recovered: %v", err)
+ }
+ call.Done()
+ c.inFlightMutex.Lock()
+ delete(c.inFlightMap, key)
+ c.inFlightMutex.Unlock()
+ }()
+
+ response, err := fn(ctx)
+ if err != nil && c.storeMissingRecords && errors.Is(err, ErrStoreMissingRecord) {
+ c.SetMissing(key, *new(T), true)
+ call.err = ErrMissingRecord
+ return
+ }
+
+ if err != nil {
+ call.err = err
+ return
+ }
+
+ res, ok := any(response).(T)
+ if !ok {
+ call.err = ErrInvalidType
+ return
+ }
+
+ call.err = nil
+ call.val = res
+ c.SetMissing(key, res, false)
+}
+
+func callAndCache[V, T any](ctx context.Context, c *Client[T], key string, fn FetchFn[V]) (V, error) {
+ c.inFlightMutex.Lock()
+ if call, ok := c.inFlightMap[key]; ok {
+ c.inFlightMutex.Unlock()
+ call.Wait()
+ return unwrap[V, T](call.val, call.err)
+ }
+
+ call := c.newFlight(key)
+ c.inFlightMutex.Unlock()
+ makeCall(ctx, c, key, fn, call)
+ return unwrap[V, T](call.val, call.err)
+}
+
 // newBatchFlight should be called with a lock.
 func (c *Client[T]) newBatchFlight(ids []string, keyFn KeyFn) *inFlightCall[map[string]T] {
  call := new(inFlightCall[map[string]T])
@@ -31,13 +80,6 @@ func (c *Client[T]) newBatchFlight(ids []string, keyFn KeyFn) *inFlightCall[map[
  return call
 }
 
-func (c *Client[T]) endFlight(call *inFlightCall[T], key string) {
- call.Done()
- c.inFlightMutex.Lock()
- delete(c.inFlightMap, key)
- c.inFlightMutex.Unlock()
-}
-
 func (c *Client[T]) endBatchFlight(ids []string, keyFn KeyFn, call *inFlightCall[map[string]T]) {
  call.Done()
  c.inFlightBatchMutex.Lock()
@@ -47,12 +89,6 @@ func (c *Client[T]) endBatchFlight(ids []string, keyFn KeyFn, call *inFlightCall
  c.inFlightBatchMutex.Unlock()
 }
 
-func (c *Client[T]) endErrorFlight(call *inFlightCall[T], key string, err error) error {
- call.err = err
- c.endFlight(call, key)
- return err
-}
-
 type makeBatchCallOpts[T, V any] struct {
  ids []string
  fn BatchFetchFn[V]
@@ -64,16 +100,14 @@ func makeBatchCall[T, V any](ctx context.Context, c *Client[T], opts makeBatchCa
  response, err := opts.fn(ctx, opts.ids)
  if err != nil {
  opts.call.err = err
- c.endBatchFlight(opts.ids, opts.keyFn, opts.call)
  return
  }
 
  // Check if we should store any of these IDs as a missing record.
  if c.storeMissingRecords && len(response) < len(opts.ids) {
  for _, id := range opts.ids {
  if _, ok := response[id]; !ok {
- var zero T
- c.SetMissing(opts.keyFn(id), zero, true)
+ c.SetMissing(opts.keyFn(id), *new(T), true)
  }
  }
  }
@@ -87,40 +121,6 @@ func makeBatchCall[T, V any](ctx context.Context, c *Client[T], opts makeBatchCa
  c.SetMissing(opts.keyFn(id), v, false)
  opts.call.val[id] = v
  }
- c.endBatchFlight(opts.ids, opts.keyFn, opts.call)
-}
-
-func callAndCache[V, T any](ctx context.Context, c *Client[T], key string, fn FetchFn[V]) (V, error) {
- c.inFlightMutex.Lock()
- if call, ok := c.inFlightMap[key]; ok {
- c.inFlightMutex.Unlock()
- call.Wait()
- return unwrap[V, T](call.val, call.err)
- }
-
- call := c.newFlight(key)
- c.inFlightMutex.Unlock()
-
- response, err := fn(ctx)
- if err != nil && c.storeMissingRecords && errors.Is(err, ErrStoreMissingRecord) {
- c.SetMissing(key, *new(T), true)
- return response, c.endErrorFlight(call, key, ErrMissingRecord)
- }
-
- if err != nil {
- return response, c.endErrorFlight(call, key, err)
- }
-
- res, ok := any(response).(T)
- if !ok {
- return response, c.endErrorFlight(call, key, ErrInvalidType)
- }
-
- c.SetMissing(key, res, false)
- call.val = res
- call.err = nil
- c.endFlight(call, key)
- return response, nil
 }
 
 type callBatchOpts[T, V any] struct {
@@ -132,7 +132,6 @@ type callBatchOpts[T, V any] struct {
 func callAndCacheBatch[V, T any](ctx context.Context, c *Client[T], opts callBatchOpts[T, V]) (map[string]V, error) {
  c.inFlightBatchMutex.Lock()
 
- // We need to keep track of the specific IDs we're after for a particular call.
  callIDs := make(map[*inFlightCall[map[string]T]][]string)
  uniqueIDs := make([]string, 0, len(opts.ids))
  for _, id := range opts.ids {
@@ -145,20 +144,22 @@ func callAndCacheBatch[V, T any](ctx context.Context, c *Client[T], opts callBat
 
  if len(uniqueIDs) > 0 {
  call := c.newBatchFlight(uniqueIDs, opts.keyFn)
- for _, id := range uniqueIDs {
- c.inFlightBatchMap[opts.keyFn(id)] = call
- callIDs[call] = append(callIDs[call], id)
- }
-
- safeGo(func() {
+ callIDs[call] = append(callIDs[call], uniqueIDs...)
+ go func() {
+ defer func() {
+ if err := recover(); err != nil {
+ call.err = fmt.Errorf("panic recovered: %v", err)
+ }
+ c.endBatchFlight(uniqueIDs, opts.keyFn, call)
+ }()
  batchCallOpts := makeBatchCallOpts[T, V]{
  ids: uniqueIDs,
  fn: opts.fn,
  keyFn: opts.keyFn,
  call: call,
  }
  makeBatchCall(ctx, c, batchCallOpts)
- })
+ }()
  }
  c.inFlightBatchMutex.Unlock()