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

---
 transport/channel.go       | 16 +++++++++++++---
 transport/local_channel.go | 35 +++++++++++++++++++++++++++--------
 2 files changed, 40 insertions(+), 11 deletions(-)

(limited to 'transport')

diff --git a/transport/channel.go b/transport/channel.go
index 011f8ae..1db6ea8 100644
--- a/transport/channel.go
+++ b/transport/channel.go
@@ -1,13 +1,23 @@
 package transport
 
-import "goflink/core"
+import (
+	"context"
+
+	"goflink/core"
+)
 
 type Receiver[T any] interface {
-	Receive() (core.StreamRecord[T], bool)
+	// Receive returns the next record, or ok=false once the stream is done or
+	// ctx is cancelled.
+	Receive(ctx context.Context) (core.StreamRecord[T], bool)
 	Close() error
 }
 
+// Emitter mirrors core.Emitter plus Close. It is declared separately because
+// core cannot import transport without an import cycle.
 type Emitter[T any] interface {
-	Emit(core.StreamRecord[T]) error
+	// Emit blocks until the record is handed off, or returns ctx.Err() if ctx
+	// is cancelled first.
+	Emit(ctx context.Context, record core.StreamRecord[T]) error
 	Close() error
 }
diff --git a/transport/local_channel.go b/transport/local_channel.go
index 4ac2c83..e3e4cba 100644
--- a/transport/local_channel.go
+++ b/transport/local_channel.go
@@ -1,6 +1,10 @@
 package transport
 
-import "goflink/core"
+import (
+	"context"
+
+	"goflink/core"
+)
 
 type LocalReceiveChannel[T any] struct {
 	recordCh <-chan core.StreamRecord[T]
@@ -12,9 +16,14 @@ func NewLocalReceiveChannel[T any](ch <-chan core.StreamRecord[T]) *LocalReceive
 	}
 }
 
-func (c *LocalReceiveChannel[T]) Receive() (core.StreamRecord[T], bool) {
-	val, ok := <-c.recordCh
-	return val, ok
+func (c *LocalReceiveChannel[T]) Receive(ctx context.Context) (core.StreamRecord[T], bool) {
+	select {
+	case val, ok := <-c.recordCh:
+		return val, ok
+	case <-ctx.Done():
+		var zero core.StreamRecord[T]
+		return zero, false
+	}
 }
 
 func (c *LocalReceiveChannel[T]) Close() error {
@@ -26,12 +35,22 @@ type LocalEmitChannel[T any] struct {
 	recordCh chan<- core.StreamRecord[T]
 }
 
-func (l LocalEmitChannel[T]) Emit(c core.StreamRecord[T]) error {
-	l.recordCh <- c
-	return nil
+func NewLocalEmitChannel[T any](ch chan<- core.StreamRecord[T]) *LocalEmitChannel[T] {
+	return &LocalEmitChannel[T]{
+		recordCh: ch,
+	}
+}
+
+func (l *LocalEmitChannel[T]) Emit(ctx context.Context, c core.StreamRecord[T]) error {
+	select {
+	case l.recordCh <- c:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 }
 
-func (l LocalEmitChannel[T]) Close() error {
+func (l *LocalEmitChannel[T]) Close() error {
 	close(l.recordCh)
 	return nil
 }
-- 
cgit v1.2.3