feat: adding lo.BufferWithContext (#580)

* feat: adding lo.BufferWithContext

* Add BufferWithContext to README table of contents

Author: samber

README.md

@@ -191,6 +191,7 @@ Supported helpers for channels:
 - [SliceToChannel](#slicetochannel)
 - [Generator](#generator)
 - [Buffer](#buffer)
+- [BufferWithContext](#bufferwithcontext)
 - [BufferWithTimeout](#bufferwithtimeout)
 - [FanIn](#fanin)
 - [FanOut](#fanout)
@@ -1928,6 +1929,32 @@ for {
 }
 ```
 
+### BufferWithContext
+
+Creates a slice of n elements from a channel, with timeout. Returns the slice, the slice length, the read time and the channel status (opened/closed).
+
+```go
+ctx, cancel := context.WithCancel(context.TODO())
+go func() {
+    ch <- 0
+    time.Sleep(10*time.Millisecond)
+    ch <- 1
+    time.Sleep(10*time.Millisecond)
+    ch <- 2
+    time.Sleep(10*time.Millisecond)
+    ch <- 3
+    time.Sleep(10*time.Millisecond)
+    ch <- 4
+    time.Sleep(10*time.Millisecond)
+    cancel()
+}()
+
+items1, length1, duration1, ok1 := lo.BufferWithContext(ctx, ch, 3)
+// []int{0, 1, 2}, 3, 20ms, true
+items2, length2, duration2, ok2 := lo.BufferWithContext(ctx, ch, 3)
+// []int{3, 4}, 2, 30ms, false
+```
+
 ### BufferWithTimeout
 
 Creates a slice of n elements from a channel, with timeout. Returns the slice, the slice length, the read time and the channel status (opened/closed).

channel.go

@@ -1,6 +1,7 @@
 package lo
 
 import (
+	"context"
 	"sync"
 	"time"
 
@@ -220,17 +221,13 @@ func Batch[T any](ch <-chan T, size int) (collection []T, length int, readTime t
 	return Buffer(ch, size)
 }
 
-// BufferWithTimeout creates a slice of n elements from a channel, with timeout. Returns the slice and the slice length.
+// BufferWithContext creates a slice of n elements from a channel, with context. Returns the slice and the slice length.
 // @TODO: we should probably provide an helper that reuse the same buffer.
-func BufferWithTimeout[T any](ch <-chan T, size int, timeout time.Duration) (collection []T, length int, readTime time.Duration, ok bool) {
-	expire := time.NewTimer(timeout)
-	defer expire.Stop()
-
+func BufferWithContext[T any](ctx context.Context, ch <-chan T, size int) (collection []T, length int, readTime time.Duration, ok bool) {
 	buffer := make([]T, 0, size)
-	index := 0
 	now := time.Now()
 
-	for ; index < size; index++ {
+	for index := 0; index < size; index++ {
 		select {
 		case item, ok := <-ch:
 			if !ok {
@@ -239,12 +236,19 @@ func BufferWithTimeout[T any](ch <-chan T, size int, timeout time.Duration) (col
 
 			buffer = append(buffer, item)
 
-		case <-expire.C:
+		case <-ctx.Done():
 			return buffer, index, time.Since(now), true
 		}
 	}
 
-	return buffer, index, time.Since(now), true
+	return buffer, size, time.Since(now), true
+}
+
+// BufferWithTimeout creates a slice of n elements from a channel, with timeout. Returns the slice and the slice length.
+func BufferWithTimeout[T any](ch <-chan T, size int, timeout time.Duration) (collection []T, length int, readTime time.Duration, ok bool) {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+	return BufferWithContext(ctx, ch, size)
 }
 
 // BatchWithTimeout creates a slice of n elements from a channel, with timeout. Returns the slice and the slice length.

channel_test.go

@@ -1,6 +1,7 @@
 package lo
 
 import (
+	"context"
 	"math/rand"
 	"testing"
 	"time"
@@ -279,6 +280,42 @@ func TestBuffer(t *testing.T) {
 	is.False(ok3)
 }
 
+func TestBufferWithContext(t *testing.T) {
+	t.Parallel()
+	testWithTimeout(t, 200*time.Millisecond)
+	is := assert.New(t)
+
+	ch1 := make(chan int, 10)
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		ch1 <- 0
+		ch1 <- 1
+		ch1 <- 2
+		time.Sleep(5 * time.Millisecond)
+		cancel()
+		ch1 <- 3
+		ch1 <- 4
+		ch1 <- 5
+		close(ch1)
+	}()
+	items1, length1, _, ok1 := BufferWithContext(ctx, ch1, 20)
+	is.Equal([]int{0, 1, 2}, items1)
+	is.Equal(3, length1)
+	is.True(ok1)
+
+	ch2 := make(chan int, 10)
+	ctx, cancel = context.WithCancel(context.Background())
+	defer cancel()
+	defer close(ch2)
+	for i := 0; i < 10; i++ {
+		ch2 <- i
+	}
+	items2, length2, _, ok2 := BufferWithContext(ctx, ch2, 5)
+	is.Equal([]int{0, 1, 2, 3, 4}, items2)
+	is.Equal(5, length2)
+	is.True(ok2)
+}
+
 func TestBufferWithTimeout(t *testing.T) {
 	t.Parallel()
 	testWithTimeout(t, 200*time.Millisecond)