feat: implement acknowledgment tracking for pipeline records

internal/pkg/pipeline/ack/ack.go

@@ -0,0 +1,176 @@
+// Package ack tracks completion of a single source record as it flows through
+// a pipeline, so the source task can defer acknowledging it until every
+// downstream branch produced from it has finished processing.
+package ack
+
+import (
+	"context"
+	"sync/atomic"
+)
+
+type catterpillarAckKey string
+
+const CATERPILLAR_ACK catterpillarAckKey = "CATERPILLAR_ACK"
+
+// Ack is created once per source record with exactly one pending branch: the
+// record itself. It rides downstream on the record's context.Context, so tasks
+// that only transform a record need no explicit wiring.
+//
+// Tasks that change a record's branch count must say so before sending: Fanout
+// for one-to-many, Joined for many-to-one, Drop or Reject for a record they
+// don't forward. A task with a nil output channel settles every record it
+// consumes.
+type Ack struct {
+	remaining atomic.Int32
+	failed    atomic.Bool
+	done      chan struct{}
+	children  []*Ack // settled with this Ack's own outcome once it completes; see Joined
+}
+
+// New returns an Ack with a single pending branch.
+func New() *Ack {
+	a := &Ack{done: make(chan struct{})}
+	a.remaining.Store(1)
+	return a
+}
+
+// AddBranch registers cnt additional branches that must call Done or Fail
+// before Wait's channel closes.
+func (a *Ack) AddBranch(cnt int32) {
+	a.remaining.Add(cnt)
+}
+
+// Done marks one branch as complete.
+func (a *Ack) Done() {
+	a.complete(false)
+}
+
+// Fail marks one branch as complete but unsuccessful, so Failed reports true
+// once every branch has finished. Use it when a branch didn't make it to where
+// it needed to go, so the source knows not to acknowledge the record.
+func (a *Ack) Fail() {
+	a.complete(true)
+}
+
+func (a *Ack) complete(failed bool) {
+
+	if failed {
+		a.failed.Store(true)
+	}
+
+	if a.remaining.Add(-1) != 0 {
+		return
+	}
+
+	// a joined record's single downstream outcome applies equally to every
+	// record that went into it, so children get this Ack's overall result
+	// rather than the outcome of this particular call.
+	anyFailed := a.failed.Load()
+	for _, c := range a.children {
+		if anyFailed {
+			c.Fail()
+		} else {
+			c.Done()
+		}
+	}
+
+	close(a.done)
+
+}
+
+// Failed reports whether any branch called Fail instead of Done. Only
+// meaningful after Wait's channel has closed.
+func (a *Ack) Failed() bool {
+	return a.failed.Load()
+}
+
+// Wait returns a channel that closes once every branch has called Done or
+// Fail.
+func (a *Ack) Wait() <-chan struct{} {
+	return a.done
+}
+
+// WithContext returns a copy of ctx carrying a, recoverable later via
+// FromContext. A nil ctx is treated as context.Background(), since an
+// aggregating task may never have received a record to inherit one from.
+func WithContext(ctx context.Context, a *Ack) context.Context {
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	return context.WithValue(ctx, CATERPILLAR_ACK, a)
+}
+
+// FromContext recovers the Ack embedded in ctx, if any.
+func FromContext(ctx context.Context) (*Ack, bool) {
+	a, ok := ctx.Value(CATERPILLAR_ACK).(*Ack)
+	return a, ok
+}
+
+// Fanout adjusts the Ack embedded in ctx, if any, so it represents n branches
+// derived from the single branch ctx currently carries. Call it once, before
+// sending any of the n outputs, so no downstream Done/Fail can race ahead of
+// the adjustment. n counts sends rather than distinct records; n == 0
+// completes the Ack immediately.
+func Fanout(ctx context.Context, n int) {
+
+	a, ok := FromContext(ctx)
+	if !ok {
+		return
+	}
+
+	switch {
+	case n == 0:
+		a.Done()
+	case n > 1:
+		a.AddBranch(int32(n - 1))
+	}
+
+}
+
+// Drop completes the Ack embedded in ctx, if any, for a record a task decided
+// not to forward downstream (e.g. it was filtered out). It is the
+// single-record equivalent of Fanout(ctx, 0).
+func Drop(ctx context.Context) {
+	if a, ok := FromContext(ctx); ok {
+		a.Done()
+	}
+}
+
+// Reject completes the Ack embedded in ctx, if any, as failed, for a record a
+// task could not process. Where Drop means the record is legitimately finished
+// with, Reject means it never made it: the source leaves it unacknowledged and
+// the broker redelivers it, rather than the pipeline waiting for a completion
+// that can't come.
+func Reject(ctx context.Context) {
+	if a, ok := FromContext(ctx); ok {
+		a.Fail()
+	}
+}
+
+// Rejected is Reject followed by err, for a task bailing out on the record it
+// is holding. Keeping the settle and the return on one line stops the two
+// drifting apart: a bare return strands the record, and the symptom is the
+// whole pipeline hanging at shutdown rather than anything pointing here.
+func Rejected(ctx context.Context, err error) error {
+	Reject(ctx)
+	return err
+}
+
+// Joined returns a new Ack with a single pending branch representing one output
+// record combined from several inputs. Attach it to that output via
+// WithContext: completing it transitively completes every Ack found among
+// ctxs, so the combined record's outcome is attributed back to each record
+// that went into it. ctxs with no Ack attached are ignored.
+func Joined(ctxs ...context.Context) *Ack {
+
+	a := New()
+
+	for _, c := range ctxs {
+		if child, ok := FromContext(c); ok {
+			a.children = append(a.children, child)
+		}
+	}
+
+	return a
+
+}

internal/pkg/pipeline/ack/tracker.go

@@ -0,0 +1,76 @@
+package ack
+
+import "sync"
+
+// Acknowledger is the broker-specific half of deferred acknowledgement: how a
+// source settles a single message once the pipeline is done with it. Tracker
+// owns the bookkeeping common to every broker, an Acknowledger owns what isn't.
+type Acknowledger interface {
+	// Ack settles one message. failed reports whether any downstream branch
+	// signalled Fail, in which case implementations should normally leave the
+	// message unacknowledged so the broker redelivers it. Called at most once
+	// per message, from its own goroutine, and at most the Tracker's
+	// concurrency at a time.
+	Ack(failed bool)
+}
+
+// Tracker lets a source task defer acknowledging a message until every
+// downstream task has finished with the record produced from it, without each
+// source having to re-implement the bookkeeping.
+//
+// Nothing here gates the source's receive loop: a cap on unacknowledged
+// messages would deadlock against any fan-in task that must accumulate records
+// before it can emit, since freeing a slot depends on the very completion the
+// fan-in is waiting to produce. Messages in flight are bounded by the
+// pipeline's channel capacity, which applies backpressure already.
+//
+// A Tracker must be created with NewTracker; the zero value is not usable.
+type Tracker struct {
+	slots chan struct{} // bounds concurrent Ack calls; acquired only after settling
+	wg    sync.WaitGroup
+}
+
+// NewTracker returns a Tracker that runs at most concurrency Ack calls at a
+// time. A value below 1 is treated as 1.
+func NewTracker(concurrency int) *Tracker {
+
+	if concurrency < 1 {
+		concurrency = 1
+	}
+
+	return &Tracker{slots: make(chan struct{}, concurrency)}
+
+}
+
+// Track watches a in the background and calls target.Ack once every
+// downstream task has signalled Done or Fail for it, passing on whether any
+// of them failed. It does not block.
+//
+// Every Track for a given Tracker must happen before its Wait.
+func (t *Tracker) Track(a *Ack, target Acknowledger) {
+
+	t.wg.Add(1)
+
+	go func() {
+
+		defer t.wg.Done()
+
+		<-a.Wait()
+
+		// take the slot after the wait, not before: the record is already
+		// through the pipeline, so freeing it depends only on Ack returning
+		// and never on the pipeline making further progress.
+		t.slots <- struct{}{}
+		defer func() { <-t.slots }()
+
+		target.Ack(a.Failed())
+
+	}()
+
+}
+
+// Wait blocks until every tracked Ack has settled and its acknowledgement
+// has been carried out. It is idempotent.
+func (t *Tracker) Wait() {
+	t.wg.Wait()
+}

internal/pkg/pipeline/pipeline.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"sync"
 
+	"github.com/patterninc/caterpillar/internal/pkg/pipeline/ack"
 	"github.com/patterninc/caterpillar/internal/pkg/pipeline/record"
 	"github.com/patterninc/caterpillar/internal/pkg/pipeline/task"
 	"gopkg.in/yaml.v3"
@@ -189,7 +190,18 @@ func (p *Pipeline) distributeToChannels(input <-chan *record.Record, outputs []c
 		}
 	}()
 
+	branches := 0
+	for _, ch := range outputs {
+		if ch != nil {
+			branches++
+		}
+	}
+
 	for rec := range input {
+		// structural fan-out: the record is duplicated to every parallel DAG
+		// branch, so its ack must represent all of them before any branch can
+		// complete it.
+		ack.Fanout(rec.Context, branches)
 		for _, ch := range outputs {
 			if ch != nil {
 				ch <- rec
@@ -251,11 +263,39 @@ func (p *Pipeline) runTaskConcurrently(t task.Task, input <-chan *record.Record,
 		}(t, input, output)
 	}
 
-	go func(wg *sync.WaitGroup, out chan<- *record.Record) {
+	go func(t task.Task, wg *sync.WaitGroup, in <-chan *record.Record, out chan<- *record.Record) {
+
 		wg.Wait()
+
+		// a worker that bailed out early can leave records in this task's
+		// input with nobody left to consume them, blocking upstream writers.
+		// Reject them so a source deferring acknowledgement redelivers them
+		// rather than waiting forever.
+		if in != nil {
+			for r := range in {
+				ack.Reject(r.Context)
+			}
+		}
+
 		if out != nil {
 			close(out)
 		}
+
+		// the output channel is closed, so downstream tasks can now drain to
+		// completion: the only safe point at which a source can wait for its
+		// deferred acknowledgements.
+		if f, ok := t.(task.Finisher); ok {
+			if err := f.Finish(); err != nil {
+				fmt.Printf("error finishing %s: %s\n", t.GetName(), err)
+				if t.GetFailOnError() {
+					p.locker.Lock()
+					p.errors[t.GetName()] = err
+					p.locker.Unlock()
+				}
+			}
+		}
+
 		p.wg.Done()
-	}(&taskWg, output)
+
+	}(t, &taskWg, input, output)
 }