feat: deduplicate identical requests

Author: sysulq

README.md

@@ -9,8 +9,9 @@ This is a implementation of a dataloader in Go.
 - 200+ lines of code, easy to understand and maintain.
 - 100% test coverage, bug free and reliable.
 - Based on generics and can be used with any type of data.
-- Use a lru cache to store the loaded values.
+- Use a LRU cache to store the loaded values.
 - Can be used to batch and cache multiple requests.
+- Deduplicate identical requests, reducing the number of requests.
 
 Installation
 ---

benchmark_test.go

@@ -23,7 +23,7 @@ func BenchmarkDataLoader(b *testing.B) {
 		}
 	})
 
-	b.Run("dataloader.AsyncLoad", func(b *testing.B) {
+	b.Run("dataloader.Go", func(b *testing.B) {
 		for i := 0; i < b.N; i++ {
 			results := make([]<-chan Result[string], 10)
 			for j := 0; j < 10; j++ {

dataloader.go

@@ -33,7 +33,7 @@ type dataLoader[K comparable, V any] struct {
 	config config
 	mu     sync.Mutex
 	batch  []K
-	chs    []chan Result[V]
+	chs    map[K][]chan Result[V]
 }
 
 // Interface is a `DataLoader` Interface which defines a public API for loading data from a particular
@@ -93,10 +93,10 @@ func New[K comparable, V any](loader Loader[K, V], options ...Option) Interface[
 	dl := &dataLoader[K, V]{
 		loader: loader,
 		config: config,
-		batch:  make([]K, 0, config.BatchSize),
-		chs:    make([]chan Result[V], 0, config.BatchSize),
 	}
 
+	dl.reset()
+
 	// Create a cache if the cache size is greater than 0
 	if config.CacheSize > 0 {
 		dl.cache = expirable.NewLRU[K, V](config.CacheSize, nil, config.CacheExpire)
@@ -150,18 +150,24 @@ func (d *dataLoader[K, V]) Go(ctx context.Context, key K) <-chan Result[V] {
 		go d.scheduleBatch(ctx, ch)
 	}
 
+	// Check if the key is in flight
+	if chs, ok := d.chs[key]; ok {
+		d.chs[key] = append(chs, ch)
+		d.mu.Unlock()
+		return ch
+	}
+
 	// If the current batch is full, start processing it
 	if len(d.batch) >= d.config.BatchSize {
 		// spawn a new goroutine to process the batch
 		go d.processBatch(ctx, d.batch, d.chs)
 		// Create a new batch, and a new set of channels
-		d.batch = make([]K, 0, d.config.BatchSize)
-		d.chs = make([]chan Result[V], 0, d.config.BatchSize)
+		d.reset()
 	}
 
 	// Add the key and channel to the current batch
 	d.batch = append(d.batch, key)
-	d.chs = append(d.chs, ch)
+	d.chs[key] = []chan Result[V]{ch}
 
 	// Unlock the DataLoader
 	d.mu.Unlock()
@@ -204,15 +210,20 @@ func (d *dataLoader[K, V]) LoadMap(ctx context.Context, keys []K) map[K]Result[V
 	return results
 }
 
+// reset resets the DataLoader
+func (d *dataLoader[K, V]) reset() {
+	d.batch = make([]K, 0, d.config.BatchSize)
+	d.chs = make(map[K][]chan Result[V], d.config.BatchSize)
+}
+
 // scheduleBatch schedules a batch to be processed
 func (d *dataLoader[K, V]) scheduleBatch(ctx context.Context, ch chan Result[V]) {
 	select {
 	case <-time.After(d.config.Wait):
 		d.mu.Lock()
 		if len(d.batch) > 0 {
 			go d.processBatch(ctx, d.batch, d.chs)
-			d.batch = make([]K, 0, d.config.BatchSize)
-			d.chs = make([]chan Result[V], 0, d.config.BatchSize)
+			d.reset()
 		}
 		d.mu.Unlock()
 	case <-ctx.Done():
@@ -221,17 +232,19 @@ func (d *dataLoader[K, V]) scheduleBatch(ctx context.Context, ch chan Result[V])
 }
 
 // processBatch processes a batch of keys
-func (d *dataLoader[K, V]) processBatch(ctx context.Context, keys []K, chs []chan Result[V]) {
+func (d *dataLoader[K, V]) processBatch(ctx context.Context, keys []K, chs map[K][]chan Result[V]) {
 	defer func() {
 		if r := recover(); r != nil {
 			const size = 64 << 10
 			buf := make([]byte, size)
 			buf = buf[:runtime.Stack(buf, false)]
 			fmt.Fprintf(os.Stderr, "Dataloader: Panic received in loader function: %v\n%s", r, buf)
 
-			for _, ch := range chs {
-				ch <- Result[V]{err: fmt.Errorf("panic received in loader function: %v", r)}
-				close(ch)
+			for _, chs := range chs {
+				for _, ch := range chs {
+					ch <- Result[V]{err: fmt.Errorf("panic received in loader function: %v", r)}
+					close(ch)
+				}
 			}
 			return
 		}
@@ -242,8 +255,11 @@ func (d *dataLoader[K, V]) processBatch(ctx context.Context, keys []K, chs []cha
 		if results[i].err == nil && d.cache != nil {
 			d.cache.Add(key, results[i].data)
 		}
-		chs[i] <- results[i]
-		close(chs[i])
+
+		for _, ch := range chs[key] {
+			ch <- results[i]
+			close(ch)
+		}
 	}
 }
 

dataloader_test.go

@@ -21,6 +21,38 @@ func TestDataLoader(t *testing.T) {
 	t.Run("LoadMap", testLoadMap)
 	t.Run("Panic recovered", testPanicRecovered)
 	t.Run("Prime", testPrime)
+	t.Run("Inflight", testInflight)
+}
+
+func testInflight(t *testing.T) {
+	loader := New(func(ctx context.Context, keys []int) []Result[string] {
+		if len(keys) != 5 {
+			t.Errorf("Expected 5 keys, got %d", keys)
+		}
+
+		results := make([]Result[string], len(keys))
+		for i, key := range keys {
+			results[i] = Result[string]{data: fmt.Sprintf("Result for %d", key)}
+		}
+		return results
+	}, WithBatchSize(5))
+
+	chs := make([]<-chan Result[string], 0)
+	for i := 0; i < 10; i++ {
+		chs = append(chs, loader.Go(context.Background(), i/2))
+	}
+
+	for idx, ch := range chs {
+		result := <-ch
+		data, err := result.Unwrap()
+		if err != nil {
+			t.Errorf("Unexpected error: %v", err)
+		}
+
+		if data != fmt.Sprintf("Result for %d", idx/2) {
+			t.Errorf("Unexpected result: %v", data)
+		}
+	}
 }
 
 func testBasicFunctionality(t *testing.T) {