Add context to fetch functions

close vikstrous#11

Author: vikstrous

README.md

@@ -9,7 +9,7 @@ It works as follows:
 * Each of many concurrently executing graphql resolver functions call Load() on the Loader object with different keys. Let's say K1, K2, K3
 * Each call to Load() with a new key is delayed slightly (a few milliseconds) so that the Loader can load them together.
 * The customizable fetch function of the loader takes a list of keys and loads data for all of them in a single batched request to the data storage layer. It might send `[K1,K2,K3]` and get back `[V1,V2,V3]`.
-* The Loader takes case of sending the right resight to the right caller and the result is cached for the duration of the graphql request.
+* The Loader takes case of sending the right result to the right caller and the result is cached for the duration of the graphql request.
 
 Usage:
 
@@ -31,7 +31,8 @@ import (
 )
 
 // fetchFn is shown as a function here, but it might work better as a method
-func fetchFn(keys []string) (ret []int, errs []error) {
+// ctx is the context from the first call to Load for the current batch
+func fetchFn(ctx context.Context, keys []string) (ret []int, errs []error) {
     for _, key := range keys {
         num, err := strconv.ParseInt(key, 10, 32)
         ret = append(ret, int(num))
@@ -61,18 +62,20 @@ func main() {
 Benchmarks show that this package is faster than both of the above and I find it easier to use.
 
 ```
-BenchmarkDataloader/caches-8                             4152324               270.3 ns/op           168 B/op          5 allocs/op
-BenchmarkDataloader/random_spread-8                      1000000              1281 ns/op             626 B/op         11 allocs/op
-BenchmarkDataloader/concurently-8                          33159             55575 ns/op           32649 B/op        160 allocs/op
-BenchmarkDataloader/all_in_one_request-8                   10000           7556166 ns/op         2574411 B/op      60032 allocs/op
-
-BenchmarkDataloaden/caches-8                            17960090                67.73 ns/op           24 B/op          1 allocs/op
-BenchmarkDataloaden/random_spread-8                      1223949               955.0 ns/op           279 B/op          5 allocs/op
-BenchmarkDataloaden/concurently-8                          27093             43594 ns/op            2867 B/op         76 allocs/op
-BenchmarkDataloaden/all_in_one_request-8                   10000           1410499 ns/op          487876 B/op      10007 allocs/op
-
-BenchmarkDataloadgen/caches-8                           22032517                53.61 ns/op            8 B/op          0 allocs/op
-BenchmarkDataloadgen/random_spread-8                     2558128               483.7 ns/op           287 B/op          4 allocs/op
-BenchmarkDataloadgen/concurently-8                         31900             34903 ns/op            2906 B/op         71 allocs/op
-BenchmarkDataloadgen/all_in_one_request-8                  10000           1032841 ns/op          573619 B/op          7 allocs/op
-```
+BenchmarkDataloader/caches-8                             4512438               267.7 ns/op           168 B/op          5 allocs/op
+BenchmarkDataloader/random_spread-8                      1000000              1267 ns/op             626 B/op         11 allocs/op
+BenchmarkDataloader/concurently-8                          16305             78332 ns/op           28619 B/op        154 allocs/op
+BenchmarkDataloader/all_in_one_request-8                   10000           7081367 ns/op         2576182 B/op      60031 allocs/op
+
+BenchmarkDataloaden/caches-8                            18537674                68.09 ns/op           24 B/op          1 allocs/op
+BenchmarkDataloaden/random_spread-8                      2885696               696.8 ns/op           295 B/op          5 allocs/op
+BenchmarkDataloaden/concurently-8                          23302             50965 ns/op            2799 B/op         75 allocs/op
+BenchmarkDataloaden/all_in_one_request-8                   10000           1279412 ns/op          487867 B/op      10007 allocs/op
+
+BenchmarkDataloadgen/caches-8                           23384140                53.00 ns/op            8 B/op          0 allocs/op
+BenchmarkDataloadgen/random_spread-8                     2394097               604.9 ns/op           292 B/op          4 allocs/op
+BenchmarkDataloadgen/concurently-8                         28359             39131 ns/op            2852 B/op         68 allocs/op
+BenchmarkDataloadgen/all_in_one_request-8                  10000           1041825 ns/op          573643 B/op          8 allocs/op
+```
+
+To run the benchmarks, run `go test -bench=. . -benchmem` from the benchmark directory.

benchmark/benchmarks_test.go

@@ -114,7 +114,7 @@ func (k IntKey) Raw() interface{} { return k }
 
 func BenchmarkDataloadgen(b *testing.B) {
 	ctx := context.Background()
-	dl := dataloadgen.NewLoader(func(keys []int) ([]benchmarkUser, []error) {
+	dl := dataloadgen.NewLoader(func(_ context.Context, keys []int) ([]benchmarkUser, []error) {
 		users := make([]benchmarkUser, len(keys))
 		errors := make([]error, len(keys))
 

dataloaden_test.go

@@ -20,7 +20,7 @@ func TestUserLoader(t *testing.T) {
 	ctx := context.Background()
 	var fetches [][]string
 	var mu sync.Mutex
-	dl := dataloadgen.NewLoader(func(keys []string) ([]*benchmarkUser, []error) {
+	dl := dataloadgen.NewLoader(func(_ context.Context, keys []string) ([]*benchmarkUser, []error) {
 		mu.Lock()
 		fetches = append(fetches, keys)
 		mu.Unlock()

dataloader_test.go

@@ -16,7 +16,7 @@ import (
 func BenchmarkLoaderFromDataloader(b *testing.B) {
 	a := &Avg{}
 	ctx := context.Background()
-	dl := dataloadgen.NewLoader(func(keys []string) (results []string, errs []error) {
+	dl := dataloadgen.NewLoader(func(_ context.Context, keys []string) (results []string, errs []error) {
 		a.Add(len(keys))
 		results = make([]string, 0, len(keys))
 		for _, key := range keys {
@@ -485,7 +485,7 @@ func TestLoader(t *testing.T) {
 func IDLoader(max int) (*dataloadgen.Loader[string, string], *[][]string) {
 	var mu sync.Mutex
 	var loadCalls [][]string
-	identityLoader := dataloadgen.NewLoader(func(keys []string) (results []string, errs []error) {
+	identityLoader := dataloadgen.NewLoader(func(_ context.Context, keys []string) (results []string, errs []error) {
 		mu.Lock()
 		loadCalls = append(loadCalls, keys)
 		mu.Unlock()
@@ -500,7 +500,7 @@ func IDLoader(max int) (*dataloadgen.Loader[string, string], *[][]string) {
 func BatchOnlyLoader(max int) (*dataloadgen.Loader[string, string], *[][]string) {
 	var mu sync.Mutex
 	var loadCalls [][]string
-	identityLoader := dataloadgen.NewLoader(func(keys []string) (results []string, errs []error) {
+	identityLoader := dataloadgen.NewLoader(func(_ context.Context, keys []string) (results []string, errs []error) {
 		mu.Lock()
 		loadCalls = append(loadCalls, keys)
 		mu.Unlock()
@@ -515,7 +515,7 @@ func BatchOnlyLoader(max int) (*dataloadgen.Loader[string, string], *[][]string)
 func ErrorLoader(max int) (*dataloadgen.Loader[string, string], *[][]string) {
 	var mu sync.Mutex
 	var loadCalls [][]string
-	identityLoader := dataloadgen.NewLoader(func(keys []string) (results []string, errs []error) {
+	identityLoader := dataloadgen.NewLoader(func(_ context.Context, keys []string) (results []string, errs []error) {
 		mu.Lock()
 		loadCalls = append(loadCalls, keys)
 		mu.Unlock()
@@ -531,7 +531,7 @@ func ErrorLoader(max int) (*dataloadgen.Loader[string, string], *[][]string) {
 func OneErrorLoader(max int) (*dataloadgen.Loader[string, string], *[][]string) {
 	var mu sync.Mutex
 	var loadCalls [][]string
-	identityLoader := dataloadgen.NewLoader(func(keys []string) (results []string, errs []error) {
+	identityLoader := dataloadgen.NewLoader(func(_ context.Context, keys []string) (results []string, errs []error) {
 		results = make([]string, max)
 		errs = make([]error, max)
 		mu.Lock()
@@ -552,7 +552,7 @@ func OneErrorLoader(max int) (*dataloadgen.Loader[string, string], *[][]string)
 
 func PanicLoader(max int) (*dataloadgen.Loader[string, string], *[][]string) {
 	var loadCalls [][]string
-	panicLoader := dataloadgen.NewLoader(func(keys []string) (results []string, errs []error) {
+	panicLoader := dataloadgen.NewLoader(func(_ context.Context, keys []string) (results []string, errs []error) {
 		panic("Programming error")
 	}, dataloadgen.WithBatchCapacity(max)) //, withSilentLogger())
 	return panicLoader, &loadCalls
@@ -561,7 +561,7 @@ func PanicLoader(max int) (*dataloadgen.Loader[string, string], *[][]string) {
 func BadLoader(max int) (*dataloadgen.Loader[string, string], *[][]string) {
 	var mu sync.Mutex
 	var loadCalls [][]string
-	identityLoader := dataloadgen.NewLoader(func(keys []string) (results []string, errs []error) {
+	identityLoader := dataloadgen.NewLoader(func(_ context.Context, keys []string) (results []string, errs []error) {
 		mu.Lock()
 		loadCalls = append(loadCalls, keys)
 		mu.Unlock()
@@ -575,7 +575,7 @@ func NoCacheLoader(max int) (*dataloadgen.Loader[string, string], *[][]string) {
 	var mu sync.Mutex
 	var loadCalls [][]string
 	// cache := &NoCache{}
-	identityLoader := dataloadgen.NewLoader(func(keys []string) (results []string, errs []error) {
+	identityLoader := dataloadgen.NewLoader(func(_ context.Context, keys []string) (results []string, errs []error) {
 		mu.Lock()
 		loadCalls = append(loadCalls, keys)
 		mu.Unlock()
@@ -592,7 +592,7 @@ func FaultyLoader() (*dataloadgen.Loader[string, string], *[][]string) {
 	var mu sync.Mutex
 	var loadCalls [][]string
 
-	loader := dataloadgen.NewLoader(func(keys []string) (results []string, errs []error) {
+	loader := dataloadgen.NewLoader(func(_ context.Context, keys []string) (results []string, errs []error) {
 		mu.Lock()
 		loadCalls = append(loadCalls, keys)
 		mu.Unlock()

dataloadgen.go

@@ -35,7 +35,7 @@ func WithTracer(tracer trace.Tracer) Option {
 }
 
 // NewLoader creates a new GenericLoader given a fetch, wait, and maxBatch
-func NewLoader[KeyT comparable, ValueT any](fetch func(keys []KeyT) ([]ValueT, []error), options ...Option) *Loader[KeyT, ValueT] {
+func NewLoader[KeyT comparable, ValueT any](fetch func(ctx context.Context, keys []KeyT) ([]ValueT, []error), options ...Option) *Loader[KeyT, ValueT] {
 	config := &loaderConfig{
 		wait:     16 * time.Millisecond,
 		maxBatch: 0, // unlimited
@@ -64,7 +64,7 @@ type loaderConfig struct {
 // Loader batches and caches requests
 type Loader[KeyT comparable, ValueT any] struct {
 	// this method provides the data for the loader
-	fetch func(keys []KeyT) ([]ValueT, []error)
+	fetch func(ctx context.Context, keys []KeyT) ([]ValueT, []error)
 
 	*loaderConfig
 
@@ -87,6 +87,7 @@ type loaderBatch[KeyT comparable, ValueT any] struct {
 	errors        []error
 	fetchExecuted bool
 	done          chan struct{}
+	firstContext  context.Context
 	contexts      []context.Context
 	spans         []trace.Span
 }
@@ -106,7 +107,7 @@ func (l *Loader[KeyT, ValueT]) LoadThunk(ctx context.Context, key KeyT) func() (
 		return it
 	}
 
-	l.startBatch()
+	l.startBatch(ctx)
 
 	if l.tracer != nil {
 		_, loadSpan := l.tracer.Start(ctx, "dataloadgen.load")
@@ -230,9 +231,12 @@ func (l *Loader[KeyT, ValueT]) Clear(key KeyT) {
 	l.mu.Unlock()
 }
 
-func (l *Loader[KeyT, ValueT]) startBatch() {
+func (l *Loader[KeyT, ValueT]) startBatch(ctx context.Context) {
 	if l.batch == nil {
-		batch := &loaderBatch[KeyT, ValueT]{done: make(chan struct{})}
+		batch := &loaderBatch[KeyT, ValueT]{
+			done:         make(chan struct{}),
+			firstContext: ctx,
+		}
 		l.batch = batch
 		go func(l *Loader[KeyT, ValueT]) {
 			time.Sleep(l.wait)
@@ -257,7 +261,7 @@ func (l *Loader[KeyT, ValueT]) startBatch() {
 				}
 			}
 
-			batch.results, batch.errors = l.fetch(batch.keys)
+			batch.results, batch.errors = l.fetch(batch.firstContext, batch.keys)
 
 			if l.tracer != nil {
 				for _, span := range spans {
@@ -289,7 +293,7 @@ func (l *Loader[KeyT, ValueT]) addKeyToBatch(b *loaderBatch[KeyT, ValueT], key K
 				}
 			}
 
-			b.results, b.errors = l.fetch(b.keys)
+			b.results, b.errors = l.fetch(b.firstContext, b.keys)
 
 			if l.tracer != nil {
 				for _, span := range spans {

dataloadgen_test.go

@@ -13,7 +13,7 @@ import (
 
 func ExampleLoader() {
 	ctx := context.Background()
-	loader := dataloadgen.NewLoader(func(keys []string) (ret []int, errs []error) {
+	loader := dataloadgen.NewLoader(func(ctx context.Context, keys []string) (ret []int, errs []error) {
 		for _, key := range keys {
 			num, err := strconv.ParseInt(key, 10, 32)
 			ret = append(ret, int(num))
@@ -36,7 +36,7 @@ func TestEdgeCases(t *testing.T) {
 	ctx := context.Background()
 	var fetches [][]int
 	var mu sync.Mutex
-	dl := dataloadgen.NewLoader(func(keys []int) ([]string, []error) {
+	dl := dataloadgen.NewLoader(func(_ context.Context, keys []int) ([]string, []error) {
 		mu.Lock()
 		fetches = append(fetches, keys)
 		mu.Unlock()