Fix File Egress Stage

docs/pages/stages/egress/file.md

@@ -28,14 +28,18 @@ This egress stage produces no downstream output message.
 | --- | --- | --- |
 | `Path` | constructor arg | Output file path. Parent directories are created automatically. |
 | `BufferSize` | `4096` | `bufio.Writer` buffer size. |
-| `FlushThresholdPercentage` | `0.75` | Flush when pending bytes reach this fraction of `BufferSize`. |
-| `FlushDeadline` | `time.Second` | Periodic flush interval. |
+| `FlushThresholdPercentage` | `0.75` | Flush when buffered bytes reach this fraction of `BufferSize`. |
+| `FlushDeadline` | `time.Second` | Maximum idle time before flushing buffered bytes. |
 
-File egress always runs in single mode and does not expose the generic
-worker-pool config.
+File egress always runs with a custom single runner and does not expose the
+generic worker-pool config.
 
 ## Internals
 
 The stage uses `os.File` and `bufio.Writer` from the standard library. It opens
-the file append-only, writes each message byte slice, flushes on threshold or
-deadline, and syncs/closes the file during stage close.
+the file append-only and writes each message byte slice exactly as returned by
+`GetBytes()`.
+
+A single runner goroutine owns the writer. It flushes when the buffer reaches
+the configured threshold, after `FlushDeadline` of input idleness while data is
+buffered, and during stage close before the file is synced and closed.

egress/file.go

@@ -3,17 +3,16 @@ package egress
 import (
 	"bufio"
 	"context"
+	"errors"
 	"os"
 	"path/filepath"
-	"sync"
-	"sync/atomic"
 	"time"
 
+	"github.com/FerroO2000/goccia/connector"
 	"github.com/FerroO2000/goccia/egress/metrics"
 	"github.com/FerroO2000/goccia/internal/config"
 	"github.com/FerroO2000/goccia/internal/stage"
 	"github.com/FerroO2000/goccia/internal/stage/env"
-	"github.com/FerroO2000/goccia/internal/stage/worker"
 	"go.opentelemetry.io/otel/attribute"
 )
 
@@ -122,114 +121,138 @@ func (fe *fileEnv) Close(ctx context.Context) {
 	fe.BaseEnv.Close(ctx)
 }
 
-// ─── Worker ─────────────────────────────────────────────────────────────────|
+// ─── Runner ─────────────────────────────────────────────────────────────────|
 
-type fileWorker[T msgSer] struct {
-	worker.BaseWorker[*fileEnv]
+var _ stage.Runner[*fileEnv] = (*fileRunner[msgSer])(nil)
 
-	ticker   *time.Ticker
-	tickerWg *sync.WaitGroup
-	flushMux *sync.Mutex
+type fileRunner[T msgSer] struct {
+	*fileEnv
 
-	notFlushedBytes atomic.Int64
+	inConnector msgConn[T]
+
+	runDone chan struct{}
 }
 
-func newFileWorkerMaker[T msgSer]() func() *fileWorker[T] {
-	return func() *fileWorker[T] {
-		return &fileWorker[T]{
-			tickerWg: &sync.WaitGroup{},
-			flushMux: &sync.Mutex{},
-		}
+func newFileRunner[T msgSer](inConnector msgConn[T]) *fileRunner[T] {
+	return &fileRunner[T]{
+		inConnector: inConnector,
+
+		runDone: make(chan struct{}),
 	}
 }
 
-func (fw *fileWorker[T]) Init(ctx context.Context) error {
-	// Create the ticker
-	fw.ticker = time.NewTicker(fw.Env.Config.FlushDeadline)
-	go fw.runTicker(ctx)
-
-	return fw.BaseWorker.Init(ctx)
+func (r *fileRunner[T]) SetEnvironment(env *fileEnv) {
+	r.fileEnv = env
 }
 
-func (fw *fileWorker[T]) runTicker(ctx context.Context) {
-	fw.tickerWg.Add(1)
-	defer fw.tickerWg.Done()
+func (r *fileRunner[T]) Init(_ context.Context) error {
+	return nil
+}
 
-	defer fw.ticker.Stop()
+// Run owns the writer loop, so writes and flushes are never concurrent.
+func (r *fileRunner[T]) Run(ctx context.Context) {
+	defer close(r.runDone)
 
 	for {
-		select {
-		case <-ctx.Done():
-			return
-
-		case <-fw.ticker.C:
-			if err := fw.flush(); err != nil {
-				fw.Tel.LogError("periodic flush failed", err, "path", fw.Env.Config.Path)
+		msgIn, err := r.read(ctx)
+		if err != nil {
+			if errors.Is(err, context.DeadlineExceeded) && ctx.Err() == nil {
+				if err := r.flush(); err != nil {
+					r.Tel.LogError("periodic flush failed", err, "path", r.Config.Path)
+				}
+
+				continue
 			}
+
+			return
 		}
+
+		r.deliver(ctx, msgIn)
+		msgIn.Destroy()
 	}
 }
 
-func (fw *fileWorker[T]) Deliver(ctx context.Context, msgIn *msg[T]) error {
-	ctx, span := fw.Tel.StartTrace(ctx, "writing file")
+// read uses a timeout only while data is buffered, turning idle time into a
+// flush deadline without a separate ticker goroutine.
+func (r *fileRunner[T]) read(ctx context.Context) (*msg[T], error) {
+	if r.writer.Buffered() == 0 {
+		return r.inConnector.Read(ctx)
+	}
+
+	readCtx, cancelReadCtx := context.WithTimeout(ctx, r.Config.FlushDeadline)
+	defer cancelReadCtx()
+
+	return r.inConnector.Read(readCtx)
+}
+
+func (r *fileRunner[T]) deliver(ctx context.Context, msgIn *msg[T]) {
+	if err := r.write(ctx, msgIn); err != nil {
+		r.GetEgressMetrics().IncrementDeliveringErrors()
+	}
+
+	r.GetEgressMetrics().IncrementDeliveredMessages()
+	r.GetEgressMetrics().RecordTotalMessageProcessingTime(
+		ctx, int(time.Since(msgIn.GetReceiveTime()).Milliseconds()),
+	)
+}
+
+func (r *fileRunner[T]) write(ctx context.Context, msgIn *msg[T]) error {
+	_, span := r.Tel.StartTrace(msgIn.LoadSpanContext(ctx), "writing file")
 	defer span.End()
 
-	// Write message bytes to file
 	chunk := msgIn.GetBody().GetBytes()
-	n, err := fw.Env.writer.Write(chunk)
+	n, err := r.writer.Write(chunk)
 	if err != nil {
-		fw.Tel.LogError("failed to write to file", err, "path", fw.Env.Config.Path)
-		fw.Env.Metrics.IncrementWriteErrors()
+		r.Tel.LogError("failed to write to file", err, "path", r.Config.Path)
+		r.Metrics.IncrementWriteErrors()
 
 		return err
 	}
 
 	writtenBytes := int64(n)
-	bytesUnflushed := fw.notFlushedBytes.Add(writtenBytes)
-
 	span.SetAttributes(attribute.Int64("chunk_size", writtenBytes))
 
-	// Check wether to flush the writer
-	if bytesUnflushed >= fw.Env.bufSizeThreshold {
-		if err := fw.flush(); err != nil {
+	if int64(r.writer.Buffered()) >= r.bufSizeThreshold {
+		if err := r.flush(); err != nil {
 			return err
 		}
 	}
 
-	// Update metrics
-	fw.Env.Metrics.AddWrittenBytes(uint(writtenBytes))
+	r.Metrics.AddWrittenBytes(uint(writtenBytes))
 
 	return nil
 }
 
-func (fw *fileWorker[T]) flush() error {
-	fw.flushMux.Lock()
-	defer fw.flushMux.Unlock()
-
-	// Check if there is anything to flush
-	if fw.notFlushedBytes.Load() == 0 {
+func (r *fileRunner[T]) flush() error {
+	if r.writer.Buffered() == 0 {
 		return nil
 	}
 
-	if err := fw.Env.writer.Flush(); err != nil {
-		fw.Tel.LogError("failed to flush writer", err, "path", fw.Env.Config.Path)
-		fw.Env.Metrics.IncrementFlushErrors()
+	if err := r.writer.Flush(); err != nil {
+		r.Tel.LogError("failed to flush writer", err, "path", r.Config.Path)
+		r.Metrics.IncrementFlushErrors()
 
 		return err
 	}
 
-	fw.notFlushedBytes.Store(0)
-
 	return nil
 }
 
-func (fw *fileWorker[T]) Close(ctx context.Context) error {
-	fw.tickerWg.Wait()
-	if err := fw.flush(); err != nil {
-		return err
+// Close waits for Run to drain before flushing the remaining buffered bytes.
+func (r *fileRunner[T]) Close(_ context.Context) {
+	<-r.runDone
+
+	if err := r.flush(); err != nil {
+		r.Tel.LogError("failed to flush writer", err, "path", r.Config.Path)
 	}
+}
 
-	return fw.BaseWorker.Close(ctx)
+func (r *fileRunner[T]) Inputs() []uintptr {
+	return []uintptr{connector.GetConnectorID(r.inConnector)}
+}
+
+func (r *fileRunner[T]) Outputs() []uintptr {
+	return []uintptr{}
 }
 
 // ─── Stage ──────────────────────────────────────────────────────────────────|
@@ -242,13 +265,9 @@ type FileStage[T msgSer] struct {
 
 // NewFileStage returns a new file egress stage.
 func NewFileStage[T msgSer](inConnector msgConn[T], cfg *FileConfig) *FileStage[T] {
-	env := newFileEnv(cfg)
-
 	return &FileStage[T]{
-		EgressStage: stage.NewEgressStage(
-			"file", inConnector, env, newFileWorkerMaker[T](), &config.Stage{
-				RunningMode: config.StageRunningModeSingle,
-			},
+		EgressStage: stage.NewEgressStageFromRunner[T](
+			"file", newFileEnv(cfg), newFileRunner(inConnector),
 		),
 	}
 }

egress/file_test.go

@@ -0,0 +1,49 @@
+package egress
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/FerroO2000/goccia/connector"
+	"github.com/FerroO2000/goccia/internal/message"
+	"github.com/stretchr/testify/require"
+)
+
+type fileTestMsg struct {
+	data []byte
+}
+
+func (m *fileTestMsg) Destroy() {}
+
+func (m *fileTestMsg) GetBytes() []byte {
+	return m.data
+}
+
+func Test_FileStage_FlushesBufferedDataAfterIdleDeadline(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "out.txt")
+
+	cfg := NewFileConfig(path)
+	cfg.BufferSize = 64
+	cfg.FlushThresholdPercentage = 1
+	cfg.FlushDeadline = 10 * time.Millisecond
+
+	conn := connector.NewRingBuffer[*fileTestMsg](1)
+	stage := NewFileStage(conn, cfg)
+	require.NoError(t, stage.Init(t.Context()))
+
+	go stage.Run(t.Context())
+
+	msg := message.NewMessage(&fileTestMsg{data: []byte("hello")})
+	msg.SetReceiveTime(time.Now())
+	require.NoError(t, conn.Write(msg))
+
+	require.Eventually(t, func() bool {
+		data, err := os.ReadFile(path)
+		return err == nil && string(data) == "hello"
+	}, time.Second, time.Millisecond)
+
+	conn.Close()
+	stage.Close(t.Context())
+}

egress/sink.go

@@ -2,6 +2,7 @@ package egress
 
 import (
 	"context"
+	"time"
 
 	"github.com/FerroO2000/goccia/connector"
 	"github.com/FerroO2000/goccia/internal/config"
@@ -59,6 +60,12 @@ func (sr *sinkRunner[T]) Run(ctx context.Context) {
 			return
 		}
 
+		metricsCtx := msgIn.LoadSpanContext(ctx)
+		sr.GetEgressMetrics().IncrementDeliveredMessages()
+		sr.GetEgressMetrics().RecordTotalMessageProcessingTime(
+			metricsCtx, int(time.Since(msgIn.GetReceiveTime()).Milliseconds()),
+		)
+
 		msgIn.Destroy()
 	}
 }

processor/merge.go

@@ -123,6 +123,7 @@ func (mr *mergeRunner[T]) readInput(ctx context.Context, inConnector msgConn[T])
 			return
 		}
 
+		mr.GetProcessorMetrics().IncrementProcessedMessages()
 		if err := mr.fanIn.Write(msgIn); err != nil {
 			msgIn.Destroy()
 		}

processor/rob.go

@@ -145,6 +145,8 @@ func (rr *robRunner[T]) Run(ctx context.Context) {
 			continue
 		}
 
+		rr.GetProcessorMetrics().IncrementProcessedMessages()
+
 		// Set the sequence number encoded in the message
 		// value into the main message struct
 		msgIn.SetSequenceNumber(msgIn.GetBody().GetSequenceNumber())

processor/tee.go

@@ -72,6 +72,7 @@ func (tr *teeRunner[T]) Run(ctx context.Context) {
 			return
 		}
 
+		tr.GetProcessorMetrics().IncrementProcessedMessages()
 		tr.clone(ctx, msgIn)
 	}
 }