Make batching stricter: errors become batch boundaries

batch.go

@@ -2,17 +2,19 @@ package rill
 
 import (
 	"time"
-
-	"github.com/destel/rill/internal/core"
 )
 
 // Batch takes a stream of items and returns a stream of batches based on a maximum size and a timeout.
 //
 // A batch is emitted when one of the following conditions is met:
 //   - The batch reaches the maximum size
 //   - The time since the first item was added to the batch exceeds the timeout
+//   - An error is encountered in the input stream
 //   - The input stream is closed
 //
+// Errors are never included in batches. Each error is forwarded to the output as a separate item,
+// preserving the relative order of values and errors.
+//
 // This function never emits empty batches. To disable the timeout and emit batches only based on the size,
 // set the timeout to -1. Setting the timeout to zero is not supported and will result in a panic
 //
@@ -21,18 +23,123 @@ import (
 // See the package documentation for more information on non-blocking ordered functions and error handling.
 func Batch[A any](in <-chan Try[A], size int, timeout time.Duration) <-chan Try[[]A] {
 	validateMinSize(size, 1)
+	if timeout == 0 {
+		// Zero timeout reads as "batch greedily until reading from the input blocks". With an unbuffered
+		// input channel, "reading blocks" is a scheduler accident, not an end-of-burst signal,
+		// so this degenerates into floods of 1-item batches. A small positive timeout is the
+		// reliable way to get the intended behavior.
+		panic("rill: zero timeout is not supported")
+	}
+
+	if in == nil {
+		return nil
+	}
+
+	out := make(chan Try[[]A])
+	t := time.NewTimer(1 * time.Hour)
+	t.Stop()
+
+	var batch []A
+
+	flush := func() {
+		t.Stop() // no need to drain t.C since Go 1.23
+
+		if len(batch) > 0 {
+			out <- Try[[]A]{Value: batch}
+			batch = make([]A, 0, size)
+		}
+	}
+
+	sendError := func(err error) {
+		out <- Try[[]A]{Error: err}
+	}
+
+	// infinite timeout
+	if timeout < 0 {
+		go func() {
+			defer close(out)
+			defer flush()
+
+			for x := range in {
+				if x.Error != nil {
+					flush()
+					sendError(x.Error)
+					continue
+				}
+
+				batch = append(batch, x.Value)
+				if len(batch) >= size {
+					flush()
+				}
+			}
+		}()
 
-	values, errs := ToChans(in)
-	batches := core.Batch(values, size, timeout)
-	return FromChans(batches, errs)
+		return out
+	}
+
+	// finite timeout
+	go func() {
+		defer close(out)
+		defer flush()
+
+		for {
+			select {
+			case <-t.C:
+				flush()
+
+			case x, ok := <-in:
+				if !ok {
+					return
+				}
+
+				if x.Error != nil {
+					flush()
+					sendError(x.Error)
+					continue
+				}
+
+				batch = append(batch, x.Value)
+				if len(batch) >= size {
+					flush()
+					continue
+				}
+
+				if len(batch) == 1 {
+					// First item was just added to the batch - start the timer.
+					// The invariant: no item sits in the batcher longer than timeout.
+					t.Reset(timeout)
+				}
+			}
+		}
+	}()
+
+	return out
 }
 
 // Unbatch is the inverse of [Batch]. It takes a stream of batches and returns a stream of individual items.
 //
 // This is a non-blocking ordered function that processes items sequentially.
 // See the package documentation for more information on non-blocking ordered functions and error handling.
 func Unbatch[A any](in <-chan Try[[]A]) <-chan Try[A] {
-	batches, errs := ToChans(in)
-	values := core.Unbatch(batches)
-	return FromChans(values, errs)
+	if in == nil {
+		return nil
+	}
+
+	out := make(chan Try[A])
+
+	go func() {
+		defer close(out)
+		for x := range in {
+			if x.Error != nil {
+				out <- Try[A]{Error: x.Error}
+				continue
+			}
+
+			for _, a := range x.Value {
+				out <- Try[A]{Value: a}
+			}
+		}
+	}()
+
+	return out
 }

batch_test.go

@@ -3,34 +3,121 @@ package rill
 import (
 	"fmt"
 	"testing"
+	"time"
 
 	"github.com/destel/rill/internal/th"
 )
 
 func TestBatch(t *testing.T) {
-	// most logic is covered by the chans pkg tests
+	th.TestVariants(t, "timeout", []time.Duration{5 * time.Second, -1}, func(t *testing.T, timeout time.Duration) {
+		t.Run("nil", func(t *testing.T) {
+			th.ExpectValue(t, Batch[string](nil, 10, timeout), nil)
+		})
 
-	th.RunSynctest(t, "correctness", func(t *testing.T) {
-		in := FromChan(th.FromRange(0, 15), nil)
-		in = replaceWithError(in, 0, fmt.Errorf("err0"))
-		in = replaceWithError(in, 5, fmt.Errorf("err5"))
-		in = replaceWithError(in, 10, fmt.Errorf("err10"))
+		t.Run("empty", func(t *testing.T) {
+			in := make(chan Try[int])
+			close(in)
 
-		batches, errs := toSliceAndErrors(Batch(in, 3, -1))
+			out := Batch(in, 10, timeout)
+			outSlice := toUnifiedStringSlice(out, "%v")
+			th.ExpectSlice(t, outSlice, nil)
+		})
+	})
 
-		th.ExpectValue(t, len(batches), 4)
-		th.ExpectSlice(t, batches[0], []int{1, 2, 3})
-		th.ExpectSlice(t, batches[1], []int{4, 6, 7})
-		th.ExpectSlice(t, batches[2], []int{8, 9, 11})
-		th.ExpectSlice(t, batches[3], []int{12, 13, 14})
+	t.Run("zero timeout panics", func(t *testing.T) {
+		defer func() {
+			if r := recover(); r == nil {
+				t.Errorf("expected panic")
+			}
+		}()
+		Batch[string](nil, 10, 0)
+	})
 
-		th.ExpectSlice(t, errs, []string{"err0", "err5", "err10"})
+	th.RunSynctest(t, "with timeout", func(t *testing.T) {
+		in := Generate(func(send func(int), sendErr func(error)) {
+			// flush by size: the mid-batch pause is too short to trigger the timeout
+			send(11)
+			send(12)
+			send(13)
+			time.Sleep(4 * time.Second)
+			send(14)
+
+			// flush by timeout: the deadline is 5s after the first item - it hits during the second sleep
+			send(21)
+			time.Sleep(4 * time.Second)
+			send(22)
+			send(23)
+			time.Sleep(4 * time.Second)
+
+			// no batch is open - the idle gap flushes nothing
+			time.Sleep(10 * time.Second)
+
+			// flush by error: the partial batch is flushed first, then the error is forwarded
+			send(31)
+			time.Sleep(4 * time.Second)
+			send(32)
+			sendErr(fmt.Errorf("err1"))
+
+			// flush by timeout again: the error disarmed the previous deadline,
+			// otherwise [41] would have flushed alone before 42 arrived
+			send(41)
+			time.Sleep(4 * time.Second)
+			send(42)
+			time.Sleep(4 * time.Second)
+
+			// errors hitting an empty batch are forwarded alone
+			sendErr(fmt.Errorf("err2"))
+			sendErr(fmt.Errorf("err3"))
+
+			// trailing partial is flushed when the input closes
+			send(51)
+			send(52)
+		})
+
+		out := Batch(in, 4, 5*time.Second)
+
+		outSlice := toUnifiedStringSlice(out, "%v")
+		th.ExpectSlice(t, outSlice, []string{
+			"[11 12 13 14]", "[21 22 23]", "[31 32]", "err1", "[41 42]", "err2", "err3", "[51 52]",
+		})
 	})
 
+	th.RunSynctest(t, "no timeout", func(t *testing.T) {
+		in := Generate(func(send func(int), sendErr func(error)) {
+			// flush by size: sleep has no effect
+			send(11)
+			send(12)
+			send(13)
+			time.Sleep(30 * time.Second)
+			send(14)
+
+			// flush by error: the partial batch is flushed first, then the error is forwarded
+			send(21)
+			send(22)
+			sendErr(fmt.Errorf("err1"))
+
+			// errors hitting an empty batch are forwarded alone
+			sendErr(fmt.Errorf("err2"))
+			sendErr(fmt.Errorf("err3"))
+
+			// trailing partial is flushed when the input closes
+			send(31)
+			send(32)
+		})
+
+		out := Batch(in, 4, -1)
+
+		outSlice := toUnifiedStringSlice(out, "%v")
+		th.ExpectSlice(t, outSlice, []string{
+			"[11 12 13 14]", "[21 22]", "err1", "err2", "err3", "[31 32]",
+		})
+	})
 }
 
 func TestUnbatch(t *testing.T) {
-	// most logic is covered by the common package tests
+	t.Run("nil", func(t *testing.T) {
+		th.ExpectValue(t, Unbatch[string](nil), nil)
+	})
 
 	th.RunSynctest(t, "correctness", func(t *testing.T) {
 		in := Generate(func(send func([]int), sendErr func(error)) {
@@ -44,8 +131,29 @@ func TestUnbatch(t *testing.T) {
 
 		out := Unbatch(in)
 
-		outSlice, outErrs := toSliceAndErrors(out)
-		th.ExpectSlice(t, outSlice, []int{1, 2, 10, 11, 12, 13, 14, 20, 21})
-		th.ExpectSlice(t, outErrs, []string{"err1", "err2"})
+		outSlice := toUnifiedStringSlice(out, "%v")
+		th.ExpectSlice(t, outSlice, []string{
+			"1", "2", "err1", "10", "11", "12", "13", "14", "err2", "20", "21",
+		})
+	})
+
+	th.RunSynctest(t, "inverse of Batch", func(t *testing.T) {
+		in := Generate(func(send func(int), sendErr func(error)) {
+			send(0)
+			send(1)
+			send(2)
+			send(3)
+			send(4)
+			sendErr(fmt.Errorf("err5"))
+			send(6)
+			sendErr(fmt.Errorf("err7"))
+			send(8)
+			send(9)
+		})
+
+		out := Unbatch(Batch(in, 3, -1))
+
+		outSlice := toUnifiedStringSlice(out, "%v")
+		th.ExpectSlice(t, outSlice, []string{"0", "1", "2", "3", "4", "err5", "6", "err7", "8", "9"})
 	})
 }