1 files changed, 81 insertions, 0 deletions

diff --git a/core/operator/keyed_process.go b/core/operator/keyed_process.go
new file mode 100644
index 0000000..1d3c86c
--- /dev/null
+++ b/core/operator/keyed_process.go
@@ -0,0 +1,81 @@
+package operator
+
+import (
+        "context"
+        "flux/core"
+        "fmt"
+)
+
+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)
+}