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

1 files changed, 27 insertions, 8 deletions

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
 }