Prevent panic when batchloader returns nil values by replacing them with Result errors

Author: jpastoor

dataloader.go

@@ -460,6 +460,13 @@ func (b *batcher) batch(originalContext context.Context) {
 		return
 	}
 
+	// When a batchFunc returns a nil in it's items, we replace those by a Result struct containing an error
+	for key, value := range items {
+		if value == nil {
+			items[key] = &Result{Error: fmt.Errorf("no value for key")}
+		}
+	}
+
 	for i, req := range reqs {
 		req.channel <- items[i]
 		close(req.channel)

dataloader_test.go

@@ -211,6 +211,32 @@ func TestLoader(t *testing.T) {
 		// TODO: expect to get some kind of warning
 	})
 
+	t.Run("first result is a nil", func(t *testing.T) {
+		t.Parallel()
+		faultyLoader, _ := LoaderNilInsteadOfResult()
+		ctx := context.Background()
+
+		n := 10
+		reqs := []Thunk{}
+		keys := Keys{}
+		for i := 0; i < n; i++ {
+			key := StringKey(strconv.Itoa(i))
+			reqs = append(reqs, faultyLoader.Load(ctx, key))
+			keys = append(keys, key)
+		}
+
+		for i, future := range reqs {
+			_, err := future()
+			if i == 0 && err == nil {
+				t.Error("expected first result to contain an error")
+			}
+
+			if i != 0 && err != nil {
+				t.Error("expected rest of results not to contain an error")
+			}
+		}
+	})
+
 	t.Run("responds to max batch size", func(t *testing.T) {
 		t.Parallel()
 		identityLoader, loadCalls := IDLoader(2)
@@ -586,6 +612,30 @@ func FaultyLoader() (*Loader, *[][]string) {
 	return loader, &loadCalls
 }
 
+// LoaderNilInsteadOfResult gives a nil result back for the first key.
+func LoaderNilInsteadOfResult() (*Loader, *[][]string) {
+	var mu sync.Mutex
+	var loadCalls [][]string
+
+	loader := NewBatchedLoader(func(_ context.Context, keys Keys) []*Result {
+		var results []*Result
+		mu.Lock()
+		loadCalls = append(loadCalls, keys.Keys())
+		mu.Unlock()
+
+		for i, key := range keys {
+			if i == 0 {
+				results = append(results, nil)
+			} else {
+				results = append(results, &Result{key, nil})
+			}
+		}
+		return results
+	})
+
+	return loader, &loadCalls
+}
+
 ///////////////////////////////////////////////////
 // Benchmarks
 ///////////////////////////////////////////////////