diff --git a/flow_test.go b/flow_test.go index 8d6238a..43852e4 100644 --- a/flow_test.go +++ b/flow_test.go @@ -6,6 +6,7 @@ import ( "fmt" "log" "math/rand" + "sync/atomic" "testing" "time" @@ -146,47 +147,91 @@ func TestFlowWithInputAndSecondFlow(t *testing.T) { } func TestFlowWithTimeout(t *testing.T) { - ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) defer cancel() + const workers = 40 + var processed int64 + proceed := make(chan struct{}, workers) // releases as many inputs as there are workers and no more + slowHandler := func(ctx context.Context, ch chan interface{}) (chan interface{}, func() error) { - st := time.Now() resCh := make(chan interface{}) resFn := func() error { defer close(resCh) - processed := 0 for inp := range ch { - // slow operation + // slow operation, released by the test rather than by the clock select { case <-ctx.Done(): - log.Printf("error %s, [%02d] - %d", ctx.Err(), CID(ctx), processed) + log.Printf("error %s, [%02d]", ctx.Err(), CID(ctx)) return ctx.Err() - case <-time.After(time.Second): + case <-proceed: } - resCh <- inp - processed++ - + atomic.AddInt64(&processed, 1) + if err := Send(ctx, resCh, inp); err != nil { + return err + } } - log.Println("slow handler completed", CID(ctx), time.Since(st), processed) return nil } return resCh, resFn } + for i := 0; i < workers; i++ { + proceed <- struct{}{} + } + f := New(Context(ctx)) f.Add( seedHandler, - f.Parallel(40, slowHandler), + f.Parallel(workers, slowHandler), collectorHandler, ) - st := time.Now() - err := f.Go().Wait() - t.Logf("%s", time.Since(st)) - assert.EqualError(t, err, "context deadline exceeded") - assert.True(t, time.Since(st) < 3010*time.Millisecond) + assert.EqualError(t, f.Go().Wait(), "context deadline exceeded") + // the deadline aborts the flow, the 100 seeded inputs are never drained + assert.True(t, atomic.LoadInt64(&processed) <= workers, "processed %d", atomic.LoadInt64(&processed)) +} + +// merge blocks sending to the next stage until the flow canceled, the stage below never reads +func TestFlowParallelCanceledOnMerge(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sent := make(chan struct{}) + emitter := func(ctx context.Context, _ chan interface{}) (chan interface{}, func() error) { + outCh := make(chan interface{}) + return outCh, func() error { + defer close(outCh) + if CID(ctx) != 0 { // the second worker closes its channel without emitting anything + return nil + } + if err := Send(ctx, outCh, 1); err != nil { + return err + } + close(sent) + return nil + } + } + + // consumes nothing from the input channel, so the merged records have nowhere to go + stalled := func(ctx context.Context, _ chan interface{}) (chan interface{}, func() error) { + outCh := make(chan interface{}) + return outCh, func() error { + defer close(outCh) + <-ctx.Done() + return ctx.Err() + } + } + + f := New(Context(ctx)) + f.Add(f.Parallel(2, emitter), stalled) + + res := f.Go() + <-sent + cancel() + assert.EqualError(t, res.Wait(), "context canceled") } func seedHandler(_ context.Context, _ chan interface{}) (chan interface{}, func() error) {