From 0958e46c3651f2c8b6d76bcffbd8820063647d59 Mon Sep 17 00:00:00 2001
From: SangTran-127 <tranquangsang12.7@gmail.com>
Date: Sat, 8 Aug 2026 17:24:01 +0700
Subject: refactor: add context support to transport operators and runtime task
 management

---
 runtime/task.go | 66 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 66 insertions(+)
 create mode 100644 runtime/task.go

(limited to 'runtime/task.go')

diff --git a/runtime/task.go b/runtime/task.go
new file mode 100644
index 0000000..3c887bb
--- /dev/null
+++ b/runtime/task.go
@@ -0,0 +1,66 @@
+package runtime
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"goflink/core"
+	"goflink/transport"
+)
+
+// RunTask runs a stream task in the background. The returned channel yields the
+// first error that stopped it (context.Canceled when ctx ends early) and is
+// closed when the task is done, so callers can join on it.
+func RunTask[IN, OUT any](ctx context.Context, input transport.Receiver[IN], output transport.Emitter[OUT], operator core.Operator[IN, OUT]) <-chan error {
+	done := make(chan error, 1)
+
+	go func() {
+		defer close(done)
+		if err := runTask(ctx, input, output, operator); err != nil {
+			done <- err
+		}
+	}()
+
+	return done
+}
+
+func runTask[IN, OUT any](ctx context.Context, input transport.Receiver[IN], output transport.Emitter[OUT], operator core.Operator[IN, OUT]) (err error) {
+	// ponytail: assumes one writer per output channel. Fan-in needs a refcount
+	// or a dedicated closer, otherwise the second Close panics.
+	defer func() {
+		if cerr := output.Close(); cerr != nil && err == nil {
+			err = fmt.Errorf("close output: %w", cerr)
+		}
+	}()
+
+	if err := operator.Open(ctx); err != nil {
+		return fmt.Errorf("open operator: %w", err)
+	}
+	defer func() {
+		// ctx is usually already cancelled by the time we get here — that is
+		// why the task stopped. Close still needs a live window to flush, so
+		// give it one that keeps ctx's values but drops the cancellation.
+		closeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
+		defer cancel()
+
+		if cerr := operator.Close(closeCtx); cerr != nil && err == nil {
+			err = fmt.Errorf("close operator: %w", cerr)
+		}
+	}()
+
+	for {
+		record, ok := input.Receive(ctx)
+		if !ok {
+			break
+		}
+
+		if err := operator.Process(ctx, record, output); err != nil {
+			return fmt.Errorf("process record: %w", err)
+		}
+	}
+
+	// Receive also returns !ok on cancellation, so report that as an error
+	// rather than a clean drain.
+	return ctx.Err()
+}
-- 
cgit v1.2.3