diff --git a/batch.go b/batch.go index 6e74f97..f418698 100644 --- a/batch.go +++ b/batch.go @@ -2,8 +2,6 @@ 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. @@ -11,8 +9,12 @@ import ( // 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,10 +23,97 @@ 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. @@ -32,7 +121,25 @@ func Batch[A any](in <-chan Try[A], size int, timeout time.Duration) <-chan Try[ // 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 } diff --git a/batch_test.go b/batch_test.go index 4ab030a..4feeb18 100644 --- a/batch_test.go +++ b/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"}) }) } diff --git a/internal/core/batch.go b/internal/core/batch.go deleted file mode 100644 index 9bde0e3..0000000 --- a/internal/core/batch.go +++ /dev/null @@ -1,116 +0,0 @@ -package core - -import ( - "fmt" - "time" -) - -// Batch groups items from an input channel into batches based on a maximum size and a timeout. -// A batch is emitted when it reaches the maximum size, the timeout expires, or the input channel closes. -// This function never emits empty batches. The timeout countdown starts when the first item is added to a new batch. -// To emit batches only when full, set the timeout to -1. Zero timeout is not supported and will panic. -func Batch[A any](in <-chan A, size int, timeout time.Duration) <-chan []A { - if in == nil { - return nil - } - - out := make(chan []A) - - switch { - case timeout == 0: - panic(fmt.Errorf("zero timeout is not supported yet")) - - case timeout < 0: - // infinite timeout - go func() { - defer close(out) - var batch []A - for a := range in { - batch = append(batch, a) - if len(batch) >= size { - out <- batch - batch = make([]A, 0, size) - } - } - if len(batch) > 0 { - out <- batch - } - }() - - default: - // finite timeout - go func() { - batch := make([]A, 0, size) - t := time.NewTicker(1 * time.Hour) - t.Stop() - - flush := func() { - if len(batch) > 0 { - out <- batch - batch = make([]A, 0, size) - } - - t.Stop() - // consume a tick that might have been sent while we were flushing - select { - case <-t.C: - default: - } - } - - for { - select { - case <-t.C: - // timeout - flush() - - case a, ok := <-in: - if !ok { - // end of input - flush() - close(out) - return - } - - // got new item - batch = append(batch, a) - - if len(batch) == 1 { - // we've just started collecting a new batch. - // start the timer to flush the batch after the timeout. - t.Reset(timeout) - } - - if len(batch) >= size { - // batch is full - flush() - } - } - - } - }() - - } - - return out -} - -// Unbatch is the inverse of Batch. It takes a channel of batches and emits individual items. -func Unbatch[A any](in <-chan []A) <-chan A { - if in == nil { - return nil - } - - out := make(chan A) - - go func() { - defer close(out) - for batch := range in { - for _, a := range batch { - out <- a - } - } - }() - - return out -} diff --git a/internal/core/batch_test.go b/internal/core/batch_test.go deleted file mode 100644 index 8359361..0000000 --- a/internal/core/batch_test.go +++ /dev/null @@ -1,97 +0,0 @@ -package core - -import ( - "testing" - "time" - - "github.com/destel/rill/internal/th" -) - -func TestBatch(t *testing.T) { - addDelayAfter := func(in <-chan int, value int, delay time.Duration) <-chan int { - out := make(chan int) - go func() { - defer close(out) - for x := range in { - out <- x - if x == value { - time.Sleep(delay) - } - } - }() - return out - } - - t.Run("nil", func(t *testing.T) { - th.ExpectValue(t, Batch[string](nil, 10, 10*time.Second), nil) - }) - - th.RunSynctest(t, "empty", func(t *testing.T) { - in := make(chan int) - close(in) - - out := Batch(in, 10, 5*time.Second) - - outSlice := th.ToSlice(out) - th.ExpectValue(t, len(outSlice), 0) - }) - - th.RunSynctest(t, "slow input", func(t *testing.T) { - in := th.FromRange(0, 10) - in = addDelayAfter(in, 2, 1*time.Second) - in = addDelayAfter(in, 5, 10*time.Second) - in = addDelayAfter(in, 8, 10*time.Second) - - out := Batch(in, 4, 5*time.Second) - - outSlice := th.ToSlice(out) - - th.ExpectValue(t, len(outSlice), 4) - th.ExpectSlice(t, outSlice[0], []int{0, 1, 2, 3}) - th.ExpectSlice(t, outSlice[1], []int{4, 5}) - th.ExpectSlice(t, outSlice[2], []int{6, 7, 8}) - th.ExpectSlice(t, outSlice[3], []int{9}) - }) - - th.RunSynctest(t, "fast input", func(t *testing.T) { - in := th.FromRange(0, 10) - in = addDelayAfter(in, 2, 1*time.Second) - in = addDelayAfter(in, 5, 1*time.Second) - in = addDelayAfter(in, 8, 1*time.Second) - - out := Batch(in, 4, 5*time.Second) - outSlice := th.ToSlice(out) - th.ExpectValue(t, len(outSlice), 3) - th.ExpectSlice(t, outSlice[0], []int{0, 1, 2, 3}) - th.ExpectSlice(t, outSlice[1], []int{4, 5, 6, 7}) - th.ExpectSlice(t, outSlice[2], []int{8, 9}) - }) - - th.RunSynctest(t, "no timeout", func(t *testing.T) { - in := th.FromRange(0, 10) - in = addDelayAfter(in, 1, 30*time.Second) - in = addDelayAfter(in, 5, 30*time.Second) - - out := Batch(in, 4, -1) - outSlice := th.ToSlice(out) - th.ExpectValue(t, len(outSlice), 3) - th.ExpectSlice(t, outSlice[0], []int{0, 1, 2, 3}) - th.ExpectSlice(t, outSlice[1], []int{4, 5, 6, 7}) - th.ExpectSlice(t, outSlice[2], []int{8, 9}) - }) -} - -func TestUnbatch(t *testing.T) { - t.Run("nil", func(t *testing.T) { - th.ExpectValue(t, Unbatch[string](nil), nil) - }) - - th.RunSynctest(t, "normal", func(t *testing.T) { - in := th.FromSlice([][]int{{1, 2, 3}, {4}, {5, 6, 7, 8}, {9, 10}}) - - out := Unbatch(in) - outSlice := th.ToSlice(out) - - th.ExpectSlice(t, outSlice, []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}) - }) -}