5 files changed, 86 insertions, 13 deletions

diff --git a/core/operator/harness_test.go b/core/operator/harness_test.go
new file mode 100644
index 0000000..383ab8c
--- /dev/null
+++ b/core/operator/harness_test.go
@@ -0,0 +1,23 @@
+package operator_test
+
+import (
+        "context"
+
+        "flux/core"
+)
+
+type captureEmitter[T any] struct {
+        records []core.StreamRecord[T]
+}
+
+func (e *captureEmitter[T]) Emit(_ context.Context, record core.StreamRecord[T]) error {
+        e.records = append(e.records, record)
+        return nil
+}
+
+func (e *captureEmitter[T]) Records() []core.StreamRecord[T] {
+        result := make([]core.StreamRecord[T], len(e.records))
+        copy(result, e.records)
+
+        return result
+}
diff --git a/core/operator/watermark.go b/core/operator/watermark.go
index 3d591e5..43d63fa 100644
--- a/core/operator/watermark.go
+++ b/core/operator/watermark.go
@@ -6,20 +6,15 @@ import (
         "flux/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] {
+func NewWatermarkGeneratorOperator[IN any](maxOutOfOrder int64) *WaterGeneratorOperator[IN] {
         return &WaterGeneratorOperator[IN]{
                 maxOutOfOrder:        maxOutOfOrder,
-                assigner:             assigner,
                 lastEmittedWatermark: -1,
         }
 }
@@ -34,8 +29,7 @@ func (w *WaterGeneratorOperator[IN]) Process(ctx context.Context, record core.St
                         return err
                 }
                 // query the time
-                eventTime := w.assigner(record.Value)
-                record.Timestamp = eventTime
+                eventTime := record.Timestamp
 
                 // update the max timestamp
                 if w.currentMaxTimestamp < eventTime {
diff --git a/core/operator/watermark_test.go b/core/operator/watermark_test.go
new file mode 100644
index 0000000..0a0495e
--- /dev/null
+++ b/core/operator/watermark_test.go
@@ -0,0 +1,58 @@
+package operator_test
+
+import (
+        "context"
+        "flux/core"
+        "flux/core/operator"
+        "testing"
+)
+
+type timestampEvent struct {
+        EventTime int64
+}
+
+func TestWatermarkGeneratorPreservesAssignedEventTimestamp(t *testing.T) {
+        ctx := context.Background()
+
+        op := operator.NewWatermarkGeneratorOperator[timestampEvent](0)
+
+        if err := op.Open(ctx); err != nil {
+                t.Fatalf("open operator failed: %v", err)
+        }
+
+        t.Cleanup(func() {
+                if err := op.Close(ctx); err != nil {
+                        t.Errorf("close operator: %v", err)
+                }
+        })
+
+        out := &captureEmitter[timestampEvent]{}
+        input := core.NewDataRecord("user-1", timestampEvent{EventTime: 999}, 40)
+
+        if err := op.Process(ctx, input, out); err != nil {
+                t.Fatalf("process operator failed: %v", err)
+        }
+
+        records := out.Records()
+        if len(records) == 0 {
+                t.Fatal("operator emitted no records")
+        }
+
+        got := records[0]
+
+        if got.Type != core.RecordData {
+                t.Fatalf("first output type: got %v, want RecordData", got.Type)
+        }
+
+        if got.Key != "user-1" {
+                t.Errorf("data key: got %q, want %q", got.Key, "user-1")
+        }
+
+        if got.Value != input.Value {
+                t.Errorf("data value: got %+v, want %+v", got.Value, input.Value)
+        }
+
+        if got.Timestamp != 40 {
+                t.Errorf("data timestamp: got %d, want %d", got.Timestamp, 40)
+        }
+}
diff --git a/main.go b/main.go
index 2399ceb..f829de6 100644
--- a/main.go
+++ b/main.go
@@ -86,9 +86,7 @@ func main() {
         ch3 := make(chan core.StreamRecord[Transaction], len(txList))
         ch4 := make(chan core.StreamRecord[string], len(txList))
 
-        wo := operator.NewWatermarkGeneratorOperator[Transaction](0, func(in Transaction) int64 {
-                return in.Timestamp
-        })
+        wo := operator.NewWatermarkGeneratorOperator[Transaction](0)
 
         ko := operator.NewKeyByOperator[Transaction](func(in Transaction) string {
                 return in.User
diff --git a/state/pebble_backend.go b/state/pebble_backend.go
index f1aae0c..e8efc3a 100644
--- a/state/pebble_backend.go
+++ b/state/pebble_backend.go
@@ -140,7 +140,7 @@ func (p *PebbleStateBackend) Snapshot(checkpointID uint64, targetDir string) err
         // get parent dir
         parentDir := filepath.Dir(targetDir)
 
-        if err := os.MkdirAll(parentDir, 0o755); err != nil {
+        if err := os.MkdirAll(parentDir, 00755); err != nil {
                 return err
         }
 
@@ -156,7 +156,7 @@ func (p *PebbleStateBackend) Snapshot(checkpointID uint64, targetDir string) err
                 }
         }()
 
-        if err := p.db.Checkpoint(tmpDir); err != nil {
+        if err := p.db.Checkpoint(tmpDir, pebble.WithFlushedWAL()); err != nil {
                 return fmt.Errorf("pebble checkpoint failed for ID %d: %w", checkpointID, err)
         }