Skip to content

Commit baacb48

Browse files
author
Matt Quinn
committed
Pass correct results to TraceBatchFinishFunc
The slice of *Result objects passed to TraceBatchFinishFunc was always empty. By moving the `defer finish(items)` down below where items is assigned the results of the batch function, tracers are passed the actual results. Fixes #63.
1 parent c87fdce commit baacb48

File tree

2 files changed

+40
-1
lines changed

2 files changed

+40
-1
lines changed

dataloader.go

+2-1
Original file line numberDiff line numberDiff line change
@@ -414,7 +414,6 @@ func (b *batcher) batch(originalContext context.Context) {
414414
}
415415

416416
ctx, finish := b.tracer.TraceBatch(originalContext, keys)
417-
defer finish(items)
418417

419418
func() {
420419
defer func() {
@@ -432,6 +431,8 @@ func (b *batcher) batch(originalContext context.Context) {
432431
items = b.batchFn(ctx, keys)
433432
}()
434433

434+
defer finish(items)
435+
435436
if panicErr != nil {
436437
for _, req := range reqs {
437438
req.channel <- &Result{Error: fmt.Errorf("Panic received in batch function: %v", panicErr)}

dataloader_test.go

+38
Original file line numberDiff line numberDiff line change
@@ -457,6 +457,26 @@ func TestLoader(t *testing.T) {
457457
}
458458
})
459459

460+
t.Run("tracer's TraceBatch finish func is passed the Result slice", func(t *testing.T) {
461+
t.Parallel()
462+
identityLoader, _ := IDLoader(0)
463+
tracer := new(RecordingTracer)
464+
identityLoader.tracer = tracer
465+
ctx := context.Background()
466+
future := identityLoader.Load(ctx, StringKey("1"))
467+
_, err := future()
468+
if err != nil {
469+
t.Error(err.Error())
470+
}
471+
472+
calls := tracer.traceBatchFinishCalls
473+
inner := []*Result{{Data: "1"}}
474+
expected := [][]*Result{inner}
475+
if !reflect.DeepEqual(calls, expected) {
476+
t.Errorf("tracer did not receive expected results. Expected %#v, got %#v", expected, calls)
477+
}
478+
})
479+
460480
}
461481

462482
// test helpers
@@ -586,6 +606,24 @@ func FaultyLoader() (*Loader, *[][]string) {
586606
return loader, &loadCalls
587607
}
588608

609+
type RecordingTracer struct {
610+
traceBatchFinishCalls [][]*Result
611+
}
612+
613+
func (t *RecordingTracer) TraceLoad(ctx context.Context, key Key) (context.Context, TraceLoadFinishFunc) {
614+
return ctx, func(Thunk) {}
615+
}
616+
617+
func (t *RecordingTracer) TraceLoadMany(ctx context.Context, keys Keys) (context.Context, TraceLoadManyFinishFunc) {
618+
return ctx, func(ThunkMany) {}
619+
}
620+
621+
func (t *RecordingTracer) TraceBatch(ctx context.Context, keys Keys) (context.Context, TraceBatchFinishFunc) {
622+
return ctx, func(results []*Result) {
623+
t.traceBatchFinishCalls = append(t.traceBatchFinishCalls, results)
624+
}
625+
}
626+
589627
///////////////////////////////////////////////////
590628
// Benchmarks
591629
///////////////////////////////////////////////////

0 commit comments

Comments
 (0)