refactor: add context support to transport operators and runtime task management

1 files changed, 140 insertions, 0 deletions

diff --git a/runtime/task_test.go b/runtime/task_test.go
new file mode 100644
index 0000000..f27a5eb
--- /dev/null
+++ b/runtime/task_test.go
@@ -0,0 +1,140 @@
+package runtime
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"goflink/core"
+	"goflink/transport"
+)
+
+// map(x*10) -> filter(>20), end to end.
+func TestPipeline(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	source := make(chan core.StreamRecord[int])
+	mid := make(chan core.StreamRecord[int])
+	sink := make(chan core.StreamRecord[int])
+
+	mapped := RunTask(ctx,
+		transport.NewLocalReceiveChannel(source),
+		transport.NewLocalEmitChannel(mid),
+		core.NewMapOperator(func(in int) (int, error) { return in * 10, nil }),
+	)
+	filtered := RunTask(ctx,
+		transport.NewLocalReceiveChannel(mid),
+		transport.NewLocalEmitChannel(sink),
+		core.NewFilterOperator(func(in int) bool { return in > 20 }),
+	)
+
+	go func() {
+		defer close(source)
+		for i := 1; i <= 5; i++ {
+			source <- core.NewDataRecord("k", i, int64(i))
+		}
+	}()
+
+	var got []int
+	for record := range sink {
+		got = append(got, record.Value)
+	}
+
+	want := []int{30, 40, 50}
+	if len(got) != len(want) {
+		t.Fatalf("got %v, want %v", got, want)
+	}
+	for i := range want {
+		if got[i] != want[i] {
+			t.Fatalf("got %v, want %v", got, want)
+		}
+	}
+
+	for _, errCh := range []<-chan error{mapped, filtered} {
+		if err := <-errCh; err != nil {
+			t.Fatalf("task failed: %v", err)
+		}
+	}
+}
+
+// A cancelled ctx must unblock a task parked in Receive.
+func TestCancelStopsTask(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	source := make(chan core.StreamRecord[int])
+	sink := make(chan core.StreamRecord[int], 1)
+
+	done := RunTask(ctx,
+		transport.NewLocalReceiveChannel(source),
+		transport.NewLocalEmitChannel(sink),
+		core.NewMapOperator(func(in int) (int, error) { return in, nil }),
+	)
+
+	cancel()
+
+	if err := <-done; !errors.Is(err, context.Canceled) {
+		t.Fatalf("got %v, want context.Canceled", err)
+	}
+}
+
+// A user function error must reach the caller with its cause intact.
+func TestProcessErrorPropagates(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	boom := errors.New("boom")
+	source := make(chan core.StreamRecord[int], 1)
+	sink := make(chan core.StreamRecord[int], 1)
+
+	done := RunTask(ctx,
+		transport.NewLocalReceiveChannel(source),
+		transport.NewLocalEmitChannel(sink),
+		core.NewMapOperator(func(in int) (int, error) { return 0, boom }),
+	)
+
+	source <- core.NewDataRecord("k", 1, 0)
+
+	if err := <-done; !errors.Is(err, boom) {
+		t.Fatalf("got %v, want %v", err, boom)
+	}
+}
+
+// closeProbe records the ctx.Err() its Close saw.
+type closeProbe struct {
+	closeCtxErr error
+}
+
+func (o *closeProbe) Open(ctx context.Context) error { return nil }
+
+func (o *closeProbe) Process(ctx context.Context, r core.StreamRecord[int], out core.Emitter[int]) error {
+	return out.Emit(ctx, r)
+}
+
+func (o *closeProbe) Close(ctx context.Context) error {
+	o.closeCtxErr = ctx.Err()
+	return nil
+}
+
+// Close must get a live ctx even when the task was killed by cancellation,
+// otherwise an operator that flushes on shutdown loses whatever it buffered.
+func TestCloseGetsLiveContext(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	source := make(chan core.StreamRecord[int])
+	sink := make(chan core.StreamRecord[int], 1)
+
+	probe := &closeProbe{}
+	done := RunTask(ctx,
+		transport.NewLocalReceiveChannel(source),
+		transport.NewLocalEmitChannel(sink),
+		probe,
+	)
+
+	cancel()
+	<-done // Close has already run by the time this channel closes
+
+	if probe.closeCtxErr != nil {
+		t.Fatalf("Close got a dead ctx (%v), so it cannot flush", probe.closeCtxErr)
+	}
+}