v0.22.1 cli flag and behavior, and test output improvement

internal/pipeline/pipeline.go

@@ -7,12 +7,12 @@ import (
 	"sync/atomic"
 	"time"
 
-	"logwisp/internal/config"
-	"logwisp/internal/core"
-	"logwisp/internal/flow"
-	"logwisp/internal/session"
-	"logwisp/internal/sink"
-	"logwisp/internal/source"
+	"github.com/lixenwraith/logwisp/internal/config"
+	"github.com/lixenwraith/logwisp/internal/core"
+	"github.com/lixenwraith/logwisp/internal/flow"
+	"github.com/lixenwraith/logwisp/internal/session"
+	"github.com/lixenwraith/logwisp/internal/sink"
+	"github.com/lixenwraith/logwisp/internal/source"
 
 	"github.com/lixenwraith/log"
 )
@@ -32,11 +32,17 @@ type Pipeline struct {
 	Stats  *PipelineStats
 	logger *log.Logger
 
+	// Sinks that stall dispatch rather than drop (core.CapBackpressure)
+	backpressure map[string]bool
+
 	// Runtime
-	ctx     context.Context
-	cancel  context.CancelFunc
-	wg      sync.WaitGroup
-	running atomic.Bool
+	ctx          context.Context
+	cancel       context.CancelFunc
+	wg           sync.WaitGroup
+	running      atomic.Bool
+	halt         chan struct{} // closed by Stop: backpressure sinks drop again
+	finished     chan struct{} // closed when every source has ended
+	finishedOnce sync.Once
 }
 
 // PipelineStats contains runtime statistics for a pipeline
@@ -70,6 +76,9 @@ func NewPipeline(
 		logger:   logger,
 		ctx:      pipelineCtx,
 		cancel:   pipelineCancel,
+
+		backpressure: make(map[string]bool),
+		finished:     make(chan struct{}),
 	}
 
 	// Create flow processor
@@ -199,6 +208,9 @@ func (p *Pipeline) initSinkCapabilities(s sink.Sink, cfg config.PluginSinkConfig
 	var hasTLS, hasAuth bool
 	for _, c := range s.Capabilities() {
 		switch c {
+		case core.CapBackpressure:
+			p.backpressure[cfg.ID] = true
+
 		// Network capabilities
 		case core.CapNetLimit:
 			continue // No-op for now, placeholder
@@ -224,31 +236,31 @@ func (p *Pipeline) initSinkCapabilities(s sink.Sink, cfg config.PluginSinkConfig
 	return nil
 }
 
-// run is the central processing loop that connects sources, flow, and sinks
-func (p *Pipeline) run() {
+// run connects the subscribed sources through the flow to the sinks. It
+// returns once every source channel has closed, at Stop or at end of input.
+func (p *Pipeline) run(subscriptions []<-chan core.LogEntry) {
 	defer p.wg.Done()
+	defer p.finish()
 	defer p.logger.Info("msg", "Pipeline processing loop stopped", "pipeline", p.Config.Name)
 
 	var componentWg sync.WaitGroup
-	// Start a goroutine for each source to fan-in data
-	for _, src := range p.Sources {
+	for _, ch := range subscriptions {
 		componentWg.Add(1)
-		go func(s source.Source) {
+		go func() {
 			defer componentWg.Done()
-			ch := s.Subscribe()
-			// Range allows in-flight data to drain cleanly once Source.Stop() closes the channel
+			// Range drains in-flight entries once the source closes the channel
 			for entry := range ch {
 				if event, passed := p.Flow.Process(entry); passed {
-					// Use non-blocking dispatcher
 					p.dispatch(event)
 				}
 			}
-		}(src)
+		}()
 	}
 
+	// Sinks keep p.ctx until Stop, so the heartbeat ends on its own context
+	hbCtx, hbCancel := context.WithCancel(p.ctx)
 	var hbWg sync.WaitGroup
-	// Start heartbeat generator if enabled
-	if heartbeatCh := p.Flow.StartHeartbeat(p.ctx); heartbeatCh != nil {
+	if heartbeatCh := p.Flow.StartHeartbeat(hbCtx); heartbeatCh != nil {
 		hbWg.Add(1)
 		go func() {
 			defer hbWg.Done()
@@ -258,9 +270,8 @@ func (p *Pipeline) run() {
 					if !ok {
 						return
 					}
-					// Use non-blocking dispatcher
 					p.dispatch(event)
-				case <-p.ctx.Done():
+				case <-hbCtx.Done():
 					return
 				}
 			}
@@ -268,25 +279,36 @@ func (p *Pipeline) run() {
 	}
 
 	componentWg.Wait()
-
-	// Terminate internal contexts (heartbeat) once flow is complete
-	p.cancel()
+	hbCancel()
 	hbWg.Wait()
 }
 
-// dispatch performs a non-blocking send to all sinks.
-// A full/stalled sink must never block the run loop or starve sibling sinks.
+// dispatch sends to every sink. A full sink drops the event so it cannot
+// stall its siblings, except a backpressure sink, which waits until Stop.
 func (p *Pipeline) dispatch(event core.TransportEvent) {
-	for _, snk := range p.Sinks {
+	for id, snk := range p.Sinks {
+		if p.backpressure[id] {
+			select {
+			case snk.Input() <- event:
+			case <-p.halt:
+				p.Stats.TotalEntriesDroppedBySink.Add(1)
+			}
+			continue
+		}
 		select {
 		case snk.Input() <- event:
 		default:
-			// Buffer full - drop to avoid deadlocking the pipeline
 			p.Stats.TotalEntriesDroppedBySink.Add(1)
 		}
 	}
 }
 
+func (p *Pipeline) finish() { p.finishedOnce.Do(func() { close(p.finished) }) }
+
+// Finished is closed when every source has ended, or the pipeline shut down.
+// Only a finite source (console at end of input) ends without Stop.
+func (p *Pipeline) Finished() <-chan struct{} { return p.finished }
+
 // Start starts the pipeline operation and all its components including flow, sources, and sinks
 func (p *Pipeline) Start() error {
 	if !p.running.CompareAndSwap(false, true) {
@@ -295,6 +317,7 @@ func (p *Pipeline) Start() error {
 
 	p.logger.Info("msg", "Starting pipeline", "pipeline", p.Config.Name)
 	p.ctx, p.cancel = context.WithCancel(context.Background())
+	p.halt = make(chan struct{})
 
 	// Start all sinks
 	for id, s := range p.Sinks {
@@ -303,7 +326,11 @@ func (p *Pipeline) Start() error {
 		}
 	}
 
-	// Start all sources
+	// Subscribed before Start: a source publishes from its first entry on
+	subscriptions := make([]<-chan core.LogEntry, 0, len(p.Sources))
+	for _, src := range p.Sources {
+		subscriptions = append(subscriptions, src.Subscribe())
+	}
 	for id, src := range p.Sources {
 		if err := src.Start(); err != nil {
 			return fmt.Errorf("failed to start source %s: %w", id, err)
@@ -313,7 +340,7 @@ func (p *Pipeline) Start() error {
 	// Start the central processing loop
 	p.Stats.StartTime = time.Now()
 	p.wg.Add(1)
-	go p.run()
+	go p.run(subscriptions)
 
 	return nil
 }
@@ -325,6 +352,7 @@ func (p *Pipeline) Stop() error {
 	}
 
 	p.logger.Info("msg", "Stopping pipeline", "pipeline", p.Config.Name)
+	close(p.halt)
 
 	// 1. Stop all sources concurrently to halt new data ingress and close their channels
 	var sourceWg sync.WaitGroup
@@ -338,7 +366,6 @@ func (p *Pipeline) Stop() error {
 	sourceWg.Wait()
 
 	// 2. Wait for the run loop to finish processing and sending all in-flight data
-	// run() inherently calls p.cancel() when the source channels are empty
 	p.wg.Wait()
 
 	// 3. Stop all sinks concurrently now that no new data will be sent
@@ -351,6 +378,7 @@ func (p *Pipeline) Stop() error {
 		}(s)
 	}
 	sinkWg.Wait()
+	p.cancel()
 
 	p.logger.Info("msg", "Pipeline stopped", "pipeline", p.Config.Name)
 	return nil
@@ -373,6 +401,7 @@ func (p *Pipeline) Shutdown() {
 	if p.Sessions != nil {
 		p.Sessions.Stop()
 	}
+	p.finish() // a pipeline that never started has no run loop to close it
 
 	p.logger.Info("msg", "Pipeline shutdown complete",
 		"component", "pipeline",

internal/pipeline/pipeline_test.go

@@ -0,0 +1,132 @@
+package pipeline
+
+import (
+	"context"
+	"fmt"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/lixenwraith/logwisp/internal/config"
+	"github.com/lixenwraith/logwisp/internal/core"
+	"github.com/lixenwraith/logwisp/internal/plugin"
+	"github.com/lixenwraith/logwisp/internal/session"
+	"github.com/lixenwraith/logwisp/internal/sink"
+	"github.com/lixenwraith/logwisp/internal/source"
+
+	"github.com/lixenwraith/log"
+)
+
+// finiteSource publishes n entries as soon as it starts, then ends
+type finiteSource struct {
+	n    int
+	subs []chan core.LogEntry
+}
+
+func (s *finiteSource) Capabilities() []core.Capability { return nil }
+func (s *finiteSource) Subscribe() <-chan core.LogEntry {
+	ch := make(chan core.LogEntry, 1)
+	s.subs = append(s.subs, ch)
+	return ch
+}
+func (s *finiteSource) Start() error {
+	go func() {
+		for i := range s.n {
+			for _, ch := range s.subs {
+				ch <- core.LogEntry{Time: time.Now(), Message: strconv.Itoa(i)}
+			}
+		}
+		for _, ch := range s.subs {
+			close(ch)
+		}
+	}()
+	return nil
+}
+func (s *finiteSource) Stop()                        {}
+func (s *finiteSource) GetStats() source.SourceStats { return source.SourceStats{} }
+
+// slowSink applies backpressure; a nil out never reads at all
+type slowSink struct {
+	in  chan core.TransportEvent
+	out chan string
+}
+
+func (s *slowSink) Capabilities() []core.Capability   { return []core.Capability{core.CapBackpressure} }
+func (s *slowSink) Input() chan<- core.TransportEvent { return s.in }
+func (s *slowSink) Start(context.Context) error {
+	if s.out != nil {
+		go func() {
+			for e := range s.in {
+				time.Sleep(10 * time.Microsecond)
+				s.out <- e.Entry.Message
+			}
+		}()
+	}
+	return nil
+}
+func (s *slowSink) Stop()                    {}
+func (s *slowSink) GetStats() sink.SinkStats { return sink.SinkStats{} }
+
+var testSinks = map[string]*slowSink{}
+
+func init() {
+	plugin.RegisterSource("test_finite", func(id string, cfg map[string]any, _ *log.Logger, _ *session.Proxy) (source.Source, error) {
+		n, _ := strconv.Atoi(fmt.Sprint(cfg["n"]))
+		return &finiteSource{n: n}, nil
+	})
+	plugin.RegisterSink("test_slow", func(id string, _ map[string]any, _ *log.Logger, _ *session.Proxy) (sink.Sink, error) {
+		return testSinks[id], nil
+	})
+}
+
+func newTestPipeline(t *testing.T, n int, s *slowSink) *Pipeline {
+	t.Helper()
+	testSinks[t.Name()] = s
+	p, err := NewPipeline(&config.PipelineConfig{
+		Name:          t.Name(),
+		PluginSources: []config.PluginSourceConfig{{ID: "src", Type: "test_finite", Config: map[string]any{"n": n}}},
+		PluginSinks:   []config.PluginSinkConfig{{ID: t.Name(), Type: "test_slow"}},
+	}, log.NewLogger())
+	if err != nil {
+		t.Fatal(err)
+	}
+	return p
+}
+
+// A source that publishes from its first moment loses nothing to a slow sink,
+// and the pipeline reports Finished when its input ends.
+func TestFiniteInputReachesBackpressureSinkWhole(t *testing.T) {
+	const n = 2000
+	s := &slowSink{in: make(chan core.TransportEvent, 1), out: make(chan string, n)}
+	p := newTestPipeline(t, n, s)
+	if err := p.Start(); err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case <-p.Finished():
+	case <-time.After(10 * time.Second):
+		t.Fatal("pipeline did not finish at the end of its input")
+	}
+	for i := range n {
+		if got := <-s.out; got != strconv.Itoa(i) {
+			t.Fatalf("entry %d = %s", i, got)
+		}
+	}
+	p.Shutdown()
+}
+
+// Stop releases a dispatch blocked on a sink that never reads
+func TestStopReleasesStalledBackpressureSink(t *testing.T) {
+	p := newTestPipeline(t, 100, &slowSink{in: make(chan core.TransportEvent, 1)})
+	if err := p.Start(); err != nil {
+		t.Fatal(err)
+	}
+	time.Sleep(50 * time.Millisecond)
+	stopped := make(chan struct{})
+	go func() { p.Shutdown(); close(stopped) }()
+	select {
+	case <-stopped:
+	case <-time.After(5 * time.Second):
+		t.Fatal("Stop blocked behind a stalled backpressure sink")
+	}
+}

internal/pipeline/registry.go

@@ -4,10 +4,10 @@ import (
 	"fmt"
 	"sync"
 
-	"logwisp/internal/plugin"
-	"logwisp/internal/session"
-	"logwisp/internal/sink"
-	"logwisp/internal/source"
+	"github.com/lixenwraith/logwisp/internal/plugin"
+	"github.com/lixenwraith/logwisp/internal/session"
+	"github.com/lixenwraith/logwisp/internal/sink"
+	"github.com/lixenwraith/logwisp/internal/source"
 
 	"github.com/lixenwraith/log"
 )