feat: add on timer & watermarkmain

8 files changed, 288 insertions, 99 deletions

diff --git a/core/filter.go b/core/operator/filter.go
index 6c95fbc..183148b 100644
--- a/core/filter.go
+++ b/core/operator/filter.go
@@ -1,6 +1,9 @@
-package core
+package operator
 
-import "context"
+import (
+	"context"
+	"goflink/core"
+)
 
 type FilterFunc[IN any] func(in IN) bool
 
@@ -16,8 +19,8 @@ func (f *FilterOperator[IN]) Open(ctx context.Context) error {
 	return nil
 }
 
-func (f *FilterOperator[IN]) Process(ctx context.Context, record StreamRecord[IN], out Emitter[IN]) error {
-	if record.Type != RecordData {
+func (f *FilterOperator[IN]) Process(ctx context.Context, record core.StreamRecord[IN], out Emitter[IN]) error {
+	if record.Type != core.RecordData {
 		return out.Emit(ctx, record)
 	}
 
diff --git a/core/operator/keyed_process.go b/core/operator/keyed_process.go
new file mode 100644
index 0000000..b507d87
--- /dev/null
+++ b/core/operator/keyed_process.go
@@ -0,0 +1,81 @@
+package operator
+
+import (
+	"context"
+	"fmt"
+	"goflink/core"
+)
+
+type KeyedProcesser[IN, OUT any] interface {
+	Open(ctx context.Context) error
+	ProcessElement(ctx context.Context, record core.StreamRecord[IN], out Emitter[OUT]) error
+	Close(ctx context.Context) error
+	OnTimer(ctx context.Context, timestamp int64, out Emitter[OUT]) error
+}
+
+type KeyProcessOperator[IN, OUT any] struct {
+	userFunc    KeyedProcesser[IN, OUT]
+	timeService core.InternalTimerService 
+	backend     core.StateBackend
+}
+
+func NewKeyProcessOperator[IN, OUT any](userFunc KeyedProcesser[IN, OUT]) *KeyProcessOperator[IN, OUT] {
+	return &KeyProcessOperator[IN, OUT]{
+		userFunc:    userFunc,
+		timeService: core.NewTimerService(),
+	}
+}
+
+func (k *KeyProcessOperator[IN, OUT]) Open(ctx context.Context) error {
+	// init backend
+	be, ok := core.ExtractStateBackend(ctx)
+	if !ok {
+		return fmt.Errorf("state backend not found")
+	}
+
+	k.backend = be
+	// inject the timer to context for user use
+	ctx = core.InjectTimerService(ctx, k.timeService)
+	return k.userFunc.Open(ctx)
+
+}
+
+func (k *KeyProcessOperator[IN, OUT]) Process(ctx context.Context, record core.StreamRecord[IN], out Emitter[OUT]) error {
+
+	switch record.Type {
+	case core.RecordData:
+		if err := k.backend.SetCurrentKey(record.Key); err != nil {
+			return err
+		}
+
+		return k.userFunc.ProcessElement(ctx, record, out)
+	case core.RecordWatermark:
+		// advance the timer
+		triggerTimerQueue := k.timeService.AdvanceWatermark(record.Timestamp)
+
+		for _, t := range triggerTimerQueue {
+
+			if err := k.backend.SetCurrentKey(t.Key); err != nil {
+				return err
+			}
+			if err := k.userFunc.OnTimer(ctx, t.Timestamp, out); err != nil {
+				return err
+			}
+		}
+		return out.Emit(ctx, core.StreamRecord[OUT]{
+			Type:      core.RecordWatermark,
+			Timestamp: record.Timestamp,
+		})
+	default:
+		return out.Emit(ctx, core.StreamRecord[OUT]{
+			Timestamp: record.Timestamp,
+			BarrierID: record.BarrierID,
+			Key:       record.Key,
+		})
+	}
+
+}
+
+func (k *KeyProcessOperator[IN, OUT]) Close(ctx context.Context) error {
+	return k.userFunc.Close(ctx)
+}
diff --git a/core/operator/watermark.go b/core/operator/watermark.go
new file mode 100644
index 0000000..75390a8
--- /dev/null
+++ b/core/operator/watermark.go
@@ -0,0 +1,58 @@
+package operator
+
+import (
+	"context"
+	"goflink/core"
+)
+
+// TimestampAssigner is responsible for event time defined by user
+type TimestampAssigner[IN any] func(in IN) int64
+
+type WaterGeneratorOperator[IN any] struct {
+	assigner             TimestampAssigner[IN]
+	maxOutOfOrder        int64
+	currentMaxTimestamp  int64
+	lastEmittedWatermark int64
+}
+
+func NewWatermarkGeneratorOperator[IN any](maxOutOfOrder int64, assigner TimestampAssigner[IN]) *WaterGeneratorOperator[IN] {
+	return &WaterGeneratorOperator[IN]{
+		maxOutOfOrder:        maxOutOfOrder,
+		assigner:             assigner,
+		lastEmittedWatermark: -1,
+	}
+}
+
+func (w *WaterGeneratorOperator[IN]) Open(ctx context.Context) error {
+	return nil
+}
+
+func (w *WaterGeneratorOperator[IN]) Process(ctx context.Context, record core.StreamRecord[IN], out Emitter[IN]) error {
+	if record.Type == core.RecordData {
+
+		// query the time
+		eventTime := w.assigner(record.Value)
+		record.Timestamp = eventTime
+
+		// update the max timestamp
+		if w.currentMaxTimestamp < eventTime {
+			w.currentMaxTimestamp = eventTime
+		}
+
+		// calc the watermark
+
+		watermark := w.currentMaxTimestamp - w.maxOutOfOrder
+
+		// check if current watermark later than previous watermark
+		if watermark > w.lastEmittedWatermark {
+			w.lastEmittedWatermark = watermark
+			return out.Emit(ctx, core.NewWatermarkRecord[IN](watermark))
+		}
+		return nil
+	}
+	return out.Emit(ctx, record)
+}
+
+func (w *WaterGeneratorOperator[IN]) Close(ctx context.Context) error {
+	return nil
+}
diff --git a/core/rich_map.go b/core/rich_map.go
deleted file mode 100644
index 9a8bc95..0000000
--- a/core/rich_map.go
+++ /dev/null
@@ -1 +0,0 @@
-package core
diff --git a/core/state_backend.go b/core/state_backend.go
index f8b83fb..729f0fe 100644
--- a/core/state_backend.go
+++ b/core/state_backend.go
@@ -3,11 +3,15 @@ package core
 import "context"
 
 type StateBackend interface {
+	// For Framework ===
 	// SetCurrentKey will perform in Memory
 	SetCurrentKey(key string) error
 	// GetCurrentKey will perform in Memory
 	GetCurrentKey() string
+	SetCurrentNamespace(namespace string) error
+	GetCurrentNamespace() string
 
+	// For User ===
 	// GetValueState will perform in disk
 	GetValueState(ctx context.Context, key string) (ValueState[any], error)
 }
@@ -16,6 +20,36 @@ type stateBackendKeyType struct{}
 
 var stateBackendKey stateBackendKeyType
 
+type typedValueState[T any] struct {
+	inner ValueState[any]
+}
+
+func CastValueType[T any](inner ValueState[any]) ValueState[T] {
+	return &typedValueState[T]{inner: inner}
+}
+
+func (s *typedValueState[T]) Value(ctx context.Context) (T, error) {
+	val, err := s.inner.Value(ctx)
+	if err != nil {
+		var zero T
+		return zero, err
+	}
+	if val == nil {
+		var zero T
+		return zero, nil
+	}
+	return val.(T), nil
+}
+
+func (s *typedValueState[T]) Update(ctx context.Context, newVal T) (T, error) {
+	_, err := s.inner.Update(ctx, newVal)
+	return newVal, err
+}
+
+func (s *typedValueState[T]) Clear(ctx context.Context) error {
+	return s.inner.Clear(ctx)
+}
+
 func InjectStateBackend(ctx context.Context, state StateBackend) context.Context {
 	return context.WithValue(ctx, stateBackendKey, state)
 }
diff --git a/core/timer.go b/core/timer.go
new file mode 100644
index 0000000..f744382
--- /dev/null
+++ b/core/timer.go
@@ -0,0 +1,108 @@
+package core
+
+import (
+	"container/heap"
+	"context"
+)
+
+// If using rockDB no need priority queue anymore
+// implement priority queue from golang
+// ref: https://pkg.go.dev/container/heap#example-package-PriorityQueue
+
+type timerServiceKey struct{}
+
+// InjectTimerService inject the timer into Context
+func InjectTimerService(ctx context.Context, ts TimerService) context.Context {
+	return context.WithValue(ctx, timerServiceKey{}, ts)
+}
+
+// ExtractTimerService get the timer for user from context
+func ExtractTimerService(ctx context.Context) (TimerService, bool) {
+	ts, ok := ctx.Value(timerServiceKey{}).(TimerService)
+	return ts, ok
+}
+
+type Timer struct {
+	// Timestamp is the time when Timer is triggered
+	Timestamp int64
+	// Key is for who User involve
+	Key string
+}
+
+// Priority Queue (Min Heap). Smallest timestamp will be on top
+
+type timerHeap []Timer
+
+func (h *timerHeap) Len() int           { return len(*h) }
+func (h *timerHeap) Less(i, j int) bool { return (*h)[i].Timestamp < (*h)[j].Timestamp }
+func (h *timerHeap) Swap(i, j int) {
+	(*h)[i], (*h)[j] = (*h)[j], (*h)[i]
+}
+
+func (h *timerHeap) Push(x any) {
+	*h = append(*h, x.(Timer))
+}
+
+func (h *timerHeap) Pop() any {
+	old := *h
+	n := len(old)
+	last := old[n-1]
+	*h = old[:n-1]
+	return last
+}
+
+type TimerService interface {
+	Register(time int64, key string)
+	CurrentWatermark() int64
+}
+
+type InternalTimerService interface {
+	TimerService
+	AdvanceWatermark(watermark int64) []Timer
+}
+
+type timerService struct {
+	queue            timerHeap
+	currentWatermark int64
+}
+
+func NewTimerService() InternalTimerService {
+	return &timerService{
+		queue:            make(timerHeap, 0),
+		currentWatermark: -1,
+	}
+}
+
+func (t *timerService) Register(time int64, key string) {
+	heap.Push(&t.queue, Timer{
+		Timestamp: time,
+		Key:       key,
+	})
+}
+
+func (t *timerService) CurrentWatermark() int64 {
+	return t.currentWatermark
+}
+
+func (t *timerService) AdvanceWatermark(watermark int64) []Timer {
+	if watermark < t.currentWatermark {
+		// no need to advance when get smaller
+		return nil
+	}
+
+	t.currentWatermark = watermark
+
+	var timers []Timer
+	for t.queue.Len() > 0 {
+		earliestTimer := t.queue[0]
+
+		if earliestTimer.Timestamp > watermark {
+			break
+		}
+
+		pop := heap.Pop(&t.queue).(Timer)
+		timers = append(timers, pop)
+	}
+
+	return timers
+}
diff --git a/goflink b/goflink
new file mode 100755
index 0000000..9a670ba
--- /dev/null
+++ b/goflink
diff --git a/main.go b/main.go
index 1d86d66..da29a2c 100644
--- a/main.go
+++ b/main.go
@@ -1,98 +1,4 @@
 package main
 
-import (
-	"context"
-	"errors"
-	"fmt"
-
-	"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()
-
-	// 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])
-
-	// keyBy(name) -> count occurrences per name
-	keyed := runtime.RunTask(ctx,
-		transport.NewLocalReceiveChannel(source),
-		transport.NewLocalEmitChannel(mid),
-		operator.NewKeyByOperator(func(in string) string { return in }),
-	)
-	counted := runtime.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)
-		}
-	}()
-
-	for record := range sink {
-		fmt.Println(record.Key, record.Value)
-	}
-
-	for _, errCh := range []<-chan error{keyed, counted} {
-		if err := <-errCh; err != nil {
-			fmt.Println("task failed:", err)
-		}
-	}
 }