10 files changed, 324 insertions, 44 deletions

diff --git a/core/operator.go b/core/operator.go
deleted file mode 100644
index 7785d44..0000000
--- a/core/operator.go
+++ /dev/null
@@ -1,13 +0,0 @@
-package core
-
-import "context"
-
-type Emitter[OUT any] interface {
-	Emit(ctx context.Context, record StreamRecord[OUT]) error
-}
-
-type Operator[IN, OUT any] interface {
-	Open(ctx context.Context) error
-	Process(ctx context.Context, record StreamRecord[IN], out Emitter[OUT]) error
-	Close(ctx context.Context) error
-}
diff --git a/core/operator/keyby.go b/core/operator/keyby.go
new file mode 100644
index 0000000..1296df0
--- /dev/null
+++ b/core/operator/keyby.go
@@ -0,0 +1,32 @@
+package operator
+
+import (
+	"context"
+	"goflink/core"
+)
+
+type KeySelector[IN any] func(in IN) string
+type KeyByOperator[IN any] struct {
+	selector KeySelector[IN]
+}
+
+func NewKeyByOperator[IN any](selector KeySelector[IN]) *KeyByOperator[IN] {
+	return &KeyByOperator[IN]{selector: selector}
+}
+
+func (k *KeyByOperator[IN]) Open(ctx context.Context) error {
+	return nil
+}
+
+func (k *KeyByOperator[IN]) Process(ctx context.Context, record core.StreamRecord[IN], out Emitter[IN]) error {
+	if record.Type != core.RecordData {
+		return out.Emit(ctx, record)
+	}
+
+	record.Key = k.selector(record.Value)
+	return out.Emit(ctx, record)
+}
+
+func (k *KeyByOperator[IN]) Close(ctx context.Context) error {
+	return nil
+}
diff --git a/core/map.go b/core/operator/map.go
index bbb8dcd..ef6912e 100644
--- a/core/map.go
+++ b/core/operator/map.go
@@ -1,8 +1,9 @@
-package core
+package operator
 
 import (
 	"context"
 	"fmt"
+	"goflink/core"
 )
 
 // MapFunc is business logic for user implement
@@ -22,10 +23,10 @@ func (fo *MapOperator[IN, OUT]) Open(ctx context.Context) error {
 	return nil
 }
 
-func (fo *MapOperator[IN, OUT]) Process(ctx context.Context, record StreamRecord[IN], out Emitter[OUT]) error {
-	if record.Type != RecordData {
+func (fo *MapOperator[IN, OUT]) Process(ctx context.Context, record core.StreamRecord[IN], out Emitter[OUT]) error {
+	if record.Type != core.RecordData {
 		//	 watermark and barrier
-		return out.Emit(ctx, StreamRecord[OUT]{
+		return out.Emit(ctx, core.StreamRecord[OUT]{
 			Type:      record.Type,
 			Key:       record.Key,
 			Timestamp: record.Timestamp,
@@ -39,8 +40,8 @@ func (fo *MapOperator[IN, OUT]) Process(ctx context.Context, record StreamRecord
 		return fmt.Errorf("cannot process value %v: %w", record.Value, err)
 	}
 
-	return out.Emit(ctx, StreamRecord[OUT]{
-		Type:      RecordData,
+	return out.Emit(ctx, core.StreamRecord[OUT]{
+		Type:      core.RecordData,
 		Key:       record.Key,
 		Timestamp: record.Timestamp,
 		Value:     outValue,
diff --git a/core/operator/operator.go b/core/operator/operator.go
new file mode 100644
index 0000000..d4bf21b
--- /dev/null
+++ b/core/operator/operator.go
@@ -0,0 +1,16 @@
+package operator
+
+import (
+	"context"
+	"goflink/core"
+)
+
+type Emitter[OUT any] interface {
+	Emit(ctx context.Context, record core.StreamRecord[OUT]) error
+}
+
+type Operator[IN, OUT any] interface {
+	Open(ctx context.Context) error
+	Process(ctx context.Context, record core.StreamRecord[IN], out Emitter[OUT]) error
+	Close(ctx context.Context) error
+}
diff --git a/core/operator/stateful_map.go b/core/operator/stateful_map.go
new file mode 100644
index 0000000..7263170
--- /dev/null
+++ b/core/operator/stateful_map.go
@@ -0,0 +1,75 @@
+package operator
+
+import (
+	"context"
+	"fmt"
+
+	"goflink/core"
+)
+
+// RichMapFunc is a stateful map on the user side: Open grabs the state handles
+// off ctx, Map runs per record with the backend already pointed at that
+// record's key.
+type RichMapFunc[IN, OUT any] interface {
+	Open(ctx context.Context) error
+	Map(ctx context.Context, in IN) (OUT, error)
+	Close(ctx context.Context) error
+}
+
+type StatefulMapOperator[IN, OUT any] struct {
+	userFunc RichMapFunc[IN, OUT]
+	backend  core.StateBackend
+}
+
+func NewStatefulMapOperator[IN, OUT any](fn RichMapFunc[IN, OUT]) *StatefulMapOperator[IN, OUT] {
+	return &StatefulMapOperator[IN, OUT]{userFunc: fn}
+}
+
+func (s *StatefulMapOperator[IN, OUT]) Open(ctx context.Context) error {
+	// Get backend state from context via using context KV
+	be, ok := core.ExtractStateBackend(ctx)
+	if !ok || be == nil {
+		return fmt.Errorf("stateful map: no state backend in context")
+	}
+
+	s.backend = be
+	// The user func opens its ValueState off the same ctx.
+	return s.userFunc.Open(ctx)
+}
+
+func (s *StatefulMapOperator[IN, OUT]) Process(ctx context.Context, record core.StreamRecord[IN], out Emitter[OUT]) error {
+	if record.Type != core.RecordData {
+		return out.Emit(ctx, core.StreamRecord[OUT]{
+			Type:      record.Type,
+			Timestamp: record.Timestamp,
+			Key:       record.Key,
+			BarrierID: record.BarrierID,
+		})
+	}
+
+	// Without a key every record would silently share the "" state slot.
+	if record.Key == "" {
+		return fmt.Errorf("stateful map: record %v has no key, keyBy it first", record.Value)
+	}
+
+	if err := s.backend.SetCurrentKey(record.Key); err != nil {
+		return fmt.Errorf("set current key %q: %w", record.Key, err)
+	}
+
+	// run the user transform
+	res, err := s.userFunc.Map(ctx, record.Value)
+	if err != nil {
+		return fmt.Errorf("cannot process value %v: %w", record.Value, err)
+	}
+
+	return out.Emit(ctx, core.StreamRecord[OUT]{
+		Type:      core.RecordData,
+		Timestamp: record.Timestamp,
+		Key:       record.Key,
+		Value:     res,
+	})
+}
+
+func (s *StatefulMapOperator[IN, OUT]) Close(ctx context.Context) error {
+	return s.userFunc.Close(ctx)
+}
diff --git a/core/rich_map.go b/core/rich_map.go
new file mode 100644
index 0000000..9a8bc95
--- /dev/null
+++ b/core/rich_map.go
@@ -0,0 +1 @@
+package core
diff --git a/core/state_backend.go b/core/state_backend.go
index 219bb63..f8b83fb 100644
--- a/core/state_backend.go
+++ b/core/state_backend.go
@@ -11,3 +11,16 @@ type StateBackend interface {
 	// GetValueState will perform in disk
 	GetValueState(ctx context.Context, key string) (ValueState[any], error)
 }
+
+type stateBackendKeyType struct{}
+
+var stateBackendKey stateBackendKeyType
+
+func InjectStateBackend(ctx context.Context, state StateBackend) context.Context {
+	return context.WithValue(ctx, stateBackendKey, state)
+}
+
+func ExtractStateBackend(ctx context.Context) (StateBackend, bool) {
+	backend, ok := ctx.Value(stateBackendKey).(StateBackend)
+	return backend, ok
+}
diff --git a/main.go b/main.go
index 806c1ed..1d86d66 100644
--- a/main.go
+++ b/main.go
@@ -2,46 +2,95 @@ package main
 
 import (
 	"context"
+	"errors"
 	"fmt"
-	"strconv"
 
 	"goflink/core"
+	"goflink/core/operator"
 	"goflink/runtime"
+	"goflink/state"
 	"goflink/transport"
 )
 
+// wordCounter counts how many times each key has been seen, in keyed state.
+type wordCounter struct {
+	count core.ValueState[any]
+}
+
+func (w *wordCounter) Open(ctx context.Context) error {
+	be, ok := core.ExtractStateBackend(ctx)
+	if !ok {
+		return errors.New("no state backend in context")
+	}
+	s, err := be.GetValueState(ctx, "word-count")
+	if err != nil {
+		return err
+	}
+	w.count = s
+	return nil
+}
+
+func (w *wordCounter) Map(ctx context.Context, in string) (int, error) {
+	v, err := w.count.Value(ctx)
+	if err != nil {
+		return 0, err
+	}
+	n, _ := v.(int) // unseen key -> nil -> 0
+	n++
+	_, err = w.count.Update(ctx, n)
+	return n, err
+}
+
+func (w *wordCounter) Close(ctx context.Context) error { return nil }
+
 func main() {
+
+	/*
+			Bây giờ bạn đã có đầy đủ đồ chơi: MemoryStateBackend, KeyByOperator, StatefulMapOperator.
+
+		Trong file main.go, bạn hãy thiết lập một Pipeline có dạng: Source ➡️ KeyBy (Trích xuất Tên) ➡️ StatefulMap (Đếm số lần Tên xuất hiện) ➡️ Sink
+
+		Yêu cầu Logic của hàm StatefulMap:
+
+		Trong hàm Open, bạn hãy khởi tạo một ValueState[int] tên là "word-count".
+		Trong hàm xử lý (User Func), bạn hãy đọc giá trị cũ lên, cộng thêm 1, lưu lại vào State, và trả về con số mới đó.
+		Hãy chạy thử xem nếu bạn truyền vào ["Sang", "Tran", "Sang", "Sang"], nó có in ra 1, 1, 2, 3 không nhé. Nếu thành công, bạn đã tự tay chế tạo xong một Streaming Engine Stateful đích thực!
+	*/
 	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[string])
+	// ponytail: one backend for the whole job. Each stateful task needs its own
+	// once tasks run on separate key-group shards.
+	ctx = core.InjectStateBackend(ctx, state.NewMemoryStateBackend(core.DefaultMaxParallelism))
+
+	source := make(chan core.StreamRecord[string])
+	mid := make(chan core.StreamRecord[string])
+	sink := make(chan core.StreamRecord[int])
 
-	// x*10, then render as text
-	mapped := runtime.RunTask(ctx,
+	// keyBy(name) -> count occurrences per name
+	keyed := runtime.RunTask(ctx,
 		transport.NewLocalReceiveChannel(source),
 		transport.NewLocalEmitChannel(mid),
-		core.NewMapOperator(func(in int) (int, error) { return in * 10, nil }),
+		operator.NewKeyByOperator(func(in string) string { return in }),
 	)
-	printed := runtime.RunTask(ctx,
+	counted := runtime.RunTask(ctx,
 		transport.NewLocalReceiveChannel(mid),
 		transport.NewLocalEmitChannel(sink),
-		core.NewMapOperator(func(in int) (string, error) { return "v=" + strconv.Itoa(in), nil }),
+		operator.NewStatefulMapOperator[string, int](&wordCounter{}),
 	)
 
 	go func() {
 		defer close(source)
-		for i := 1; i <= 5; i++ {
-			source <- core.NewDataRecord("k", i, 0)
+		for _, name := range []string{"Sang", "Tran", "Sang", "Sang"} {
+			source <- core.NewDataRecord("", name, 0)
 		}
 	}()
 
 	for record := range sink {
-		fmt.Println(record.Value)
+		fmt.Println(record.Key, record.Value)
 	}
 
-	for _, errCh := range []<-chan error{mapped, printed} {
+	for _, errCh := range []<-chan error{keyed, counted} {
 		if err := <-errCh; err != nil {
 			fmt.Println("task failed:", err)
 		}
diff --git a/runtime/task.go b/runtime/task.go
index 3c887bb..b3d1caa 100644
--- a/runtime/task.go
+++ b/runtime/task.go
@@ -5,19 +5,19 @@ import (
 	"fmt"
 	"time"
 
-	"goflink/core"
+	"goflink/core/operator"
 	"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 {
+func RunTask[IN, OUT any](ctx context.Context, input transport.Receiver[IN], output transport.Emitter[OUT], op operator.Operator[IN, OUT]) <-chan error {
 	done := make(chan error, 1)
 
 	go func() {
 		defer close(done)
-		if err := runTask(ctx, input, output, operator); err != nil {
+		if err := runTask(ctx, input, output, op); err != nil {
 			done <- err
 		}
 	}()
@@ -25,7 +25,7 @@ func RunTask[IN, OUT any](ctx context.Context, input transport.Receiver[IN], out
 	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) {
+func runTask[IN, OUT any](ctx context.Context, input transport.Receiver[IN], output transport.Emitter[OUT], op operator.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() {
@@ -34,7 +34,7 @@ func runTask[IN, OUT any](ctx context.Context, input transport.Receiver[IN], out
 		}
 	}()
 
-	if err := operator.Open(ctx); err != nil {
+	if err := op.Open(ctx); err != nil {
 		return fmt.Errorf("open operator: %w", err)
 	}
 	defer func() {
@@ -44,7 +44,7 @@ func runTask[IN, OUT any](ctx context.Context, input transport.Receiver[IN], out
 		closeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
 		defer cancel()
 
-		if cerr := operator.Close(closeCtx); cerr != nil && err == nil {
+		if cerr := op.Close(closeCtx); cerr != nil && err == nil {
 			err = fmt.Errorf("close operator: %w", cerr)
 		}
 	}()
@@ -55,7 +55,7 @@ func runTask[IN, OUT any](ctx context.Context, input transport.Receiver[IN], out
 			break
 		}
 
-		if err := operator.Process(ctx, record, output); err != nil {
+		if err := op.Process(ctx, record, output); err != nil {
 			return fmt.Errorf("process record: %w", err)
 		}
 	}
diff --git a/runtime/task_test.go b/runtime/task_test.go
index f27a5eb..ae72c14 100644
--- a/runtime/task_test.go
+++ b/runtime/task_test.go
@@ -6,6 +6,8 @@ import (
 	"testing"
 
 	"goflink/core"
+	"goflink/core/operator"
+	"goflink/state"
 	"goflink/transport"
 )
 
@@ -21,12 +23,12 @@ func TestPipeline(t *testing.T) {
 	mapped := RunTask(ctx,
 		transport.NewLocalReceiveChannel(source),
 		transport.NewLocalEmitChannel(mid),
-		core.NewMapOperator(func(in int) (int, error) { return in * 10, nil }),
+		operator.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 }),
+		operator.NewFilterOperator(func(in int) bool { return in > 20 }),
 	)
 
 	go func() {
@@ -58,6 +60,110 @@ func TestPipeline(t *testing.T) {
 	}
 }
 
+// wordCounter is the canonical stateful user func: read the old count for the
+// current key, +1, store it back.
+type wordCounter struct {
+	count core.ValueState[any]
+}
+
+func (w *wordCounter) Open(ctx context.Context) error {
+	be, ok := core.ExtractStateBackend(ctx)
+	if !ok {
+		return errors.New("no state backend in context")
+	}
+	s, err := be.GetValueState(ctx, "word-count")
+	if err != nil {
+		return err
+	}
+	w.count = s
+	return nil
+}
+
+func (w *wordCounter) Map(ctx context.Context, in string) (int, error) {
+	v, err := w.count.Value(ctx)
+	if err != nil {
+		return 0, err
+	}
+	n, _ := v.(int) // unseen key -> nil -> 0
+	n++
+	_, err = w.count.Update(ctx, n)
+	return n, err
+}
+
+func (w *wordCounter) Close(ctx context.Context) error { return nil }
+
+// keyBy -> statefulMap: state must be per key and survive across records.
+func TestStatefulMapCountsPerKey(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	ctx = core.InjectStateBackend(ctx, state.NewMemoryStateBackend(core.DefaultMaxParallelism))
+
+	source := make(chan core.StreamRecord[string])
+	mid := make(chan core.StreamRecord[string])
+	sink := make(chan core.StreamRecord[int])
+
+	keyed := RunTask(ctx,
+		transport.NewLocalReceiveChannel(source),
+		transport.NewLocalEmitChannel(mid),
+		operator.NewKeyByOperator(func(in string) string { return in }),
+	)
+	counted := RunTask(ctx,
+		transport.NewLocalReceiveChannel(mid),
+		transport.NewLocalEmitChannel(sink),
+		operator.NewStatefulMapOperator[string, int](&wordCounter{}),
+	)
+
+	go func() {
+		defer close(source)
+		for _, name := range []string{"Sang", "Tran", "Sang", "Sang"} {
+			source <- core.NewDataRecord("", name, 0)
+		}
+	}()
+
+	var got []int
+	for record := range sink {
+		got = append(got, record.Value)
+	}
+
+	want := []int{1, 1, 2, 3}
+	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{keyed, counted} {
+		if err := <-errCh; err != nil {
+			t.Fatalf("task failed: %v", err)
+		}
+	}
+}
+
+// An unkeyed record must fail loudly instead of silently sharing one state slot.
+func TestStatefulMapRejectsUnkeyed(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	ctx = core.InjectStateBackend(ctx, state.NewMemoryStateBackend(core.DefaultMaxParallelism))
+
+	source := make(chan core.StreamRecord[string], 1)
+	sink := make(chan core.StreamRecord[int], 1)
+
+	done := RunTask(ctx,
+		transport.NewLocalReceiveChannel(source),
+		transport.NewLocalEmitChannel(sink),
+		operator.NewStatefulMapOperator[string, int](&wordCounter{}),
+	)
+
+	source <- core.NewDataRecord("", "Sang", 0)
+
+	if err := <-done; err == nil {
+		t.Fatal("unkeyed record was accepted, want an error")
+	}
+}
+
 // A cancelled ctx must unblock a task parked in Receive.
 func TestCancelStopsTask(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
@@ -68,7 +174,7 @@ func TestCancelStopsTask(t *testing.T) {
 	done := RunTask(ctx,
 		transport.NewLocalReceiveChannel(source),
 		transport.NewLocalEmitChannel(sink),
-		core.NewMapOperator(func(in int) (int, error) { return in, nil }),
+		operator.NewMapOperator(func(in int) (int, error) { return in, nil }),
 	)
 
 	cancel()
@@ -90,7 +196,7 @@ func TestProcessErrorPropagates(t *testing.T) {
 	done := RunTask(ctx,
 		transport.NewLocalReceiveChannel(source),
 		transport.NewLocalEmitChannel(sink),
-		core.NewMapOperator(func(in int) (int, error) { return 0, boom }),
+		operator.NewMapOperator(func(in int) (int, error) { return 0, boom }),
 	)
 
 	source <- core.NewDataRecord("k", 1, 0)
@@ -107,7 +213,7 @@ type closeProbe struct {
 
 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 {
+func (o *closeProbe) Process(ctx context.Context, r core.StreamRecord[int], out operator.Emitter[int]) error {
 	return out.Emit(ctx, r)
 }