diff --git a/internal/pkg/pipeline/ack/ack.go b/internal/pkg/pipeline/ack/ack.go new file mode 100644 index 0000000..f2a34d3 --- /dev/null +++ b/internal/pkg/pipeline/ack/ack.go @@ -0,0 +1,181 @@ +// 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" + "sync/atomic" +) + +type catterpillarAckKey string + +const CATERPILLAR_ACK catterpillarAckKey = "CATERPILLAR_ACK" + +// Ack is created once per source record and rides downstream on its context. The +// counter is the live branches descending from that record, and each send registers +// its own — so a task declares no fan-out count, it settles what it consumes. +type Ack struct { + mu sync.Mutex + settled bool + 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 no branches yet. The send that puts its record on an +// output channel registers the first one, which is why a counter that starts at +// one would never reach zero. +func New() *Ack { + return &Ack{done: make(chan struct{})} +} + +// AddBranch registers cnt additional branches that must call Done or Fail +// before Wait's channel closes. Sending a record does this for the caller. +func (a *Ack) AddBranch(cnt int32) { + + a.mu.Lock() + defer a.mu.Unlock() + + // a settled record has already been reported to the source; re-opening its + // counter would let it settle a second time and close done twice. + if a.settled { + return + } + + a.remaining.Add(cnt) + +} + +// Done marks one branch as complete, settling the record if it was the last. +func (a *Ack) Done() { + + a.mu.Lock() + + if a.settled || a.remaining.Add(-1) != 0 { + a.mu.Unlock() + return + } + + a.settled = true + a.mu.Unlock() + + a.finish() + +} + +// Fail abandons the record so the broker redelivers it. It settles immediately +// rather than decrementing like Done, because a sibling branch that never +// completes would otherwise suppress the failure and hang the source. +func (a *Ack) Fail() { + + a.mu.Lock() + + // after settling, the outcome is already reported: a late Fail must not flip a + // record the source has been told to acknowledge into one it redelivers. + if a.settled { + a.mu.Unlock() + return + } + + a.settled = true + a.failed.Store(true) + a.mu.Unlock() + + a.finish() + +} + +func (a *Ack) finish() { + + // 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 +} + +// Release completes one branch: this task is finished with the record, whether it +// forwarded, fanned out, or filtered it. Exactly once per record consumed — +// twice acknowledges early, never leaves it for redelivery. +func Release(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 Release 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 an Ack for one record combined from several inputs; completing it +// transitively completes every Ack among ctxs. An aggregator therefore does not +// Release its inputs — this settles them, per contribution rather than per Ack. +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 + +} diff --git a/internal/pkg/pipeline/ack/ack_test.go b/internal/pkg/pipeline/ack/ack_test.go new file mode 100644 index 0000000..e16166a --- /dev/null +++ b/internal/pkg/pipeline/ack/ack_test.go @@ -0,0 +1,408 @@ +package ack + +import ( + "context" + "fmt" + "sync" + "testing" +) + +// settled reports whether a has finished, without blocking. +func settled(a *Ack) bool { + select { + case <-a.Wait(): + return true + default: + return false + } +} + +func mustSettle(t *testing.T, a *Ack, wantFailed bool, format string, args ...any) { + + t.Helper() + + label := fmt.Sprintf(format, args...) + + if !settled(a) { + t.Fatalf("%s: ack never settled", label) + } + if got := a.Failed(); got != wantFailed { + t.Fatalf("%s: Failed() = %v, want %v", label, got, wantFailed) + } + +} + +// send models what Base.SendRecord does on every emit: register one branch. +func send(ctx context.Context) { + if a, ok := FromContext(ctx); ok { + a.AddBranch(1) + } +} + +// TestShapes walks the accounting rule — every send registers a branch, every task +// releases each record it consumes exactly once — over the pipeline shapes that +// exist, asserting each source record settles exactly once. +func TestShapes(t *testing.T) { + + tests := []struct { + name string + run func(ctx context.Context) + wantFailed bool + }{ + { + name: "source emit then terminal sink releases", + run: func(ctx context.Context) { + send(ctx) // source emits + Release(ctx) // sink is done with it + }, + }, + { + name: "one-to-one transform", + run: func(ctx context.Context) { + send(ctx) // source emits + send(ctx) // transform emits its derived record + Release(ctx) // transform releases its input + Release(ctx) // sink releases + }, + }, + { + name: "fan-out to three, then all three sink", + run: func(ctx context.Context) { + send(ctx) + for range 3 { + send(ctx) + } + Release(ctx) // the fanning task releases its one input + for range 3 { + Release(ctx) // each branch reaches a sink + } + }, + }, + { + name: "fan-out to zero is a drop", + run: func(ctx context.Context) { + send(ctx) + Release(ctx) // filtered away without emitting + }, + }, + { + name: "chained fan-out: 2 then each to 2", + run: func(ctx context.Context) { + send(ctx) + send(ctx) + send(ctx) + Release(ctx) // first task + for range 2 { + send(ctx) + send(ctx) + Release(ctx) // second task, once per input + } + for range 4 { + Release(ctx) + } + }, + }, + { + name: "reject anywhere fails the record", + run: func(ctx context.Context) { + send(ctx) + send(ctx) + Release(ctx) + Reject(ctx) // the derived record never made it + }, + wantFailed: true, + }, + { + name: "one failed branch of three fails the record", + run: func(ctx context.Context) { + send(ctx) + for range 3 { + send(ctx) + } + Release(ctx) + Release(ctx) + Reject(ctx) + Release(ctx) + }, + wantFailed: true, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + + a := New() + ctx := WithContext(context.Background(), a) + + if settled(a) { + t.Fatal("a fresh ack must not be settled") + } + + tc.run(ctx) + + mustSettle(t, a, tc.wantFailed, "%s", tc.name) + + }) + } + +} + +// TestPrematureSettle is the property the whole design turns on: while any +// descendant is still outstanding, the source record must not settle. +func TestPrematureSettle(t *testing.T) { + + a := New() + ctx := WithContext(context.Background(), a) + + send(ctx) // source emits + + // a task fans out to three and releases its input. Three branches outstanding. + send(ctx) + send(ctx) + send(ctx) + Release(ctx) + + for i := range 3 { + if settled(a) { + t.Fatalf("settled with %d branch(es) still outstanding", 3-i) + } + Release(ctx) + } + + mustSettle(t, a, false, "after all three branches") + +} + +// TestSettleOnce covers the guard that keeps a miscount from closing done twice, +// which would panic and take the process down. +func TestSettleOnce(t *testing.T) { + + t.Run("release after settling is a no-op", func(t *testing.T) { + + a := New() + ctx := WithContext(context.Background(), a) + + send(ctx) + Release(ctx) + mustSettle(t, a, false, "first release") + + // a buggy task releasing twice must not re-settle + Release(ctx) + Release(ctx) + + mustSettle(t, a, false, "after extra releases") + + }) + + t.Run("branch registered after settling is ignored", func(t *testing.T) { + + a := New() + ctx := WithContext(context.Background(), a) + + send(ctx) + Release(ctx) + mustSettle(t, a, false, "first release") + + // without the guard this reopens the counter, and the Release below + // drives it to zero a second time -> close of closed channel + send(ctx) + Release(ctx) + + mustSettle(t, a, false, "after late branch") + + }) + + t.Run("reject after settling does not flip the outcome", func(t *testing.T) { + + a := New() + ctx := WithContext(context.Background(), a) + + send(ctx) + Release(ctx) + Reject(ctx) + + if a.Failed() { + t.Fatal("a record that already succeeded must not become failed") + } + + }) + + t.Run("reject settles immediately, before siblings finish", func(t *testing.T) { + + a := New() + ctx := WithContext(context.Background(), a) + + send(ctx) + send(ctx) + send(ctx) + Release(ctx) + + Reject(ctx) + + // one leaked sibling must not suppress the failure, or the source can + // never learn why the record needs redelivery + mustSettle(t, a, true, "on reject with siblings outstanding") + + }) + +} + +// TestJoined covers aggregation: an aggregator does not release its buffered +// inputs, the combined record's outcome settles them. +func TestJoined(t *testing.T) { + + t.Run("combined record settles every input", func(t *testing.T) { + + parents := make([]*Ack, 3) + ctxs := make([]context.Context, 3) + for i := range parents { + parents[i] = New() + ctxs[i] = WithContext(context.Background(), parents[i]) + send(ctxs[i]) // each source emitted its record + } + + joined := Joined(ctxs...) + joinedCtx := WithContext(context.Background(), joined) + send(joinedCtx) // the aggregator emits the combined record + + for i, p := range parents { + if settled(p) { + t.Fatalf("input %d settled before the archive carrying it was handled", i) + } + } + + Release(joinedCtx) // the sink handled the combined record + + for i, p := range parents { + mustSettle(t, p, false, "input %d", i) + } + + }) + + t.Run("failing the combined record fails every input", func(t *testing.T) { + + parents := make([]*Ack, 2) + ctxs := make([]context.Context, 2) + for i := range parents { + parents[i] = New() + ctxs[i] = WithContext(context.Background(), parents[i]) + send(ctxs[i]) + } + + joined := Joined(ctxs...) + joinedCtx := WithContext(context.Background(), joined) + send(joinedCtx) + Reject(joinedCtx) + + for i, p := range parents { + mustSettle(t, p, true, "input %d", i) + } + + }) + + t.Run("one source contributing twice is completed twice", func(t *testing.T) { + + // an upstream fan-out sent two records from one source message, and both + // landed in the same batch. Collapsing them would settle the source + // while the batch was still in flight. + a := New() + ctx := WithContext(context.Background(), a) + + send(ctx) // source emit + send(ctx) // fan-out: two records from it + send(ctx) + Release(ctx) // the fanning task released its input + + joined := Joined(ctx, ctx) // both contributed to one batch + joinedCtx := WithContext(context.Background(), joined) + send(joinedCtx) + + if settled(a) { + t.Fatal("settled before the batch was handled") + } + + Release(joinedCtx) + + mustSettle(t, a, false, "after the batch was handled") + + }) + +} + +// TestUntracked: a pipeline with no acking source must be completely unaffected. +func TestUntracked(t *testing.T) { + + ctx := context.Background() + + // none of these have an ack to find, and none may panic + send(ctx) + Release(ctx) + Reject(ctx) + + if err := Rejected(ctx, context.Canceled); err != context.Canceled { + t.Fatalf("Rejected must pass the error through, got %v", err) + } + + if _, ok := FromContext(ctx); ok { + t.Fatal("a bare context must not yield an ack") + } + + // Joined over untracked contexts yields an ack with no children + j := Joined(ctx, ctx) + jctx := WithContext(context.Background(), j) + send(jctx) + Release(jctx) + mustSettle(t, j, false, "joined over untracked inputs") + +} + +// TestWithContextNilParent covers the aggregator that has no input record to +// inherit a context from. +func TestWithContextNilParent(t *testing.T) { + + a := New() + ctx := WithContext(context.TODO(), a) + + if got, ok := FromContext(ctx); !ok || got != a { + t.Fatal("ack must be recoverable from a context built on a nil parent") + } + +} + +// TestConcurrent runs the emit/release cycle from many goroutines, so -race can +// find unsynchronised access to the counter and the settle guard. +func TestConcurrent(t *testing.T) { + + const ( + workers = 16 + perTask = 200 + ) + + a := New() + ctx := WithContext(context.Background(), a) + + send(ctx) // the source's own emit, so the record is live while workers run + + var wg sync.WaitGroup + for range workers { + wg.Add(1) + go func() { + defer wg.Done() + for range perTask { + send(ctx) // fan out + Release(ctx) + } + }() + } + + wg.Wait() + + if settled(a) { + t.Fatal("settled while the source's own branch was still outstanding") + } + + Release(ctx) + + mustSettle(t, a, false, "after the final release") + +} diff --git a/internal/pkg/pipeline/ack/tracker.go b/internal/pkg/pipeline/ack/tracker.go new file mode 100644 index 0000000..99c56ae --- /dev/null +++ b/internal/pkg/pipeline/ack/tracker.go @@ -0,0 +1,77 @@ +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. Pipeline occupancy is already bounded by +// channel capacity; acks that have settled but not yet been deleted wait only +// on the concurrency slots around Ack, so a slow broker can accumulate them. +// +// 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() +} diff --git a/internal/pkg/pipeline/pipeline.go b/internal/pkg/pipeline/pipeline.go index 8540082..5f173ce 100644 --- a/internal/pkg/pipeline/pipeline.go +++ b/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" @@ -190,11 +191,26 @@ func (p *Pipeline) distributeToChannels(input <-chan *record.Record, outputs []c }() for rec := range input { + + // this duplication doesn't go through Base.SendRecord, so it registers each + // branch itself. The branch the record arrived with is released only once + // every copy has been handed over, so a branch that completes early cannot + // settle the record while later ones are still being dispatched. + a, tracked := ack.FromContext(rec.Context) + for _, ch := range outputs { if ch != nil { + if tracked { + a.AddBranch(1) + } ch <- rec } } + + if tracked { + a.Done() + } + } } @@ -251,11 +267,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) } diff --git a/internal/pkg/pipeline/task/acking_test.go b/internal/pkg/pipeline/task/acking_test.go new file mode 100644 index 0000000..0e3d458 --- /dev/null +++ b/internal/pkg/pipeline/task/acking_test.go @@ -0,0 +1,436 @@ +package task_test + +import ( + "context" + "testing" + "time" + + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/task/compress" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/task/delay" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/task/echo" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/task/flatten" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/task/join" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/task/jq" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/task/replace" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/task/sample" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/task/split" + "gopkg.in/yaml.v3" +) + +const settleTimeout = 5 * time.Second + +// tracked builds n records, each carrying its own Ack the way an acking source +// emits them: the ack starts with no branches and the send registers the first. +func tracked(n int, data func(i int) []byte) ([]*record.Record, []*ack.Ack) { + + records := make([]*record.Record, n) + acks := make([]*ack.Ack, n) + + for i := range records { + a := ack.New() + a.AddBranch(1) // the source's own emit + acks[i] = a + records[i] = &record.Record{ + ID: i + 1, + Data: data(i), + Context: ack.WithContext(context.Background(), a), + } + } + + return records, acks + +} + +// newTask builds a task of the given type from YAML, exactly as a pipeline would. +func newTask[T any](t *testing.T, new func() (task.Task, error), spec string) task.Task { + + t.Helper() + + tk, err := new() + if err != nil { + t.Fatalf("constructing task: %v", err) + } + if err := yaml.Unmarshal([]byte(spec), tk); err != nil { + t.Fatalf("unmarshalling %q: %v", spec, err) + } + if err := tk.Init(); err != nil { + t.Fatalf("initialising task: %v", err) + } + + return tk + +} + +// runTask feeds records through a task and collects what it emits, without +// releasing any of it: the caller decides when a downstream sink "handles" each +// emitted record, which is what makes premature settling observable. +func runTask(t *testing.T, tk task.Task, records []*record.Record, withOutput bool) ([]*record.Record, error) { + + t.Helper() + + input := make(chan *record.Record, len(records)) + for _, r := range records { + input <- r + } + close(input) + + var output chan *record.Record + if withOutput { + output = make(chan *record.Record, 4096) + } + + var emitted []*record.Record + drained := make(chan struct{}) + go func() { + defer close(drained) + if output == nil { + return + } + for r := range output { + emitted = append(emitted, r) + } + }() + + err := tk.Run(input, chanOrNil(output)) + if output != nil { + close(output) + } + <-drained + + return emitted, err + +} + +// sinkAll plays the part of the downstream sinks: every emitted record has one +// outstanding branch, and handling it releases exactly that branch. +func sinkAll(emitted []*record.Record) { + for _, r := range emitted { + ack.Release(r.Context) + } +} + +// chanOrNil keeps a nil *chan* nil when widened to a send-only channel: a typed +// nil wrapped in a non-nil interface would defeat every `output == nil` check. +func chanOrNil(c chan *record.Record) chan<- *record.Record { + if c == nil { + return nil + } + return c +} + +// assertSettled waits for every ack to settle, then checks the outcome. An ack +// that never settles is the failure mode this whole suite exists to catch: it +// hangs a real pipeline at shutdown. +func assertSettled(t *testing.T, acks []*ack.Ack, wantFailed bool) { + + t.Helper() + + for i, a := range acks { + select { + case <-a.Wait(): + if got := a.Failed(); got != wantFailed { + t.Errorf("record %d: Failed() = %v, want %v", i, got, wantFailed) + } + case <-time.After(settleTimeout): + t.Fatalf("record %d never settled: a source would wait on it forever", i) + } + } + +} + +// TestStreamingTasksSettleTheirInput covers the shapes that make up most of the +// pipeline: a task consumes a record, emits zero or more, and must settle the +// input exactly once either way. +func TestStreamingTasksSettleTheirInput(t *testing.T) { + + tests := []struct { + name string + task func(t *testing.T) task.Task + data func(i int) []byte + records int + withOutput bool + wantEmitted int + }{ + { + name: "echo forwards", + task: func(t *testing.T) task.Task { return newTask[any](t, echo.New, "name: e\ntype: echo") }, + data: func(int) []byte { return []byte(`{"a":1}`) }, + records: 3, + withOutput: true, + wantEmitted: 3, + }, + { + name: "echo as a terminal sink", + task: func(t *testing.T) task.Task { return newTask[any](t, echo.New, "name: e\ntype: echo") }, + data: func(int) []byte { return []byte(`{"a":1}`) }, + records: 3, + withOutput: false, + wantEmitted: 0, + }, + { + name: "delay forwards", + task: func(t *testing.T) task.Task { return newTask[any](t, delay.New, "name: d\ntype: delay\nduration: 1ms") }, + data: func(int) []byte { return []byte(`x`) }, + records: 3, + withOutput: true, + wantEmitted: 3, + }, + { + name: "flatten transforms one to one", + task: func(t *testing.T) task.Task { return newTask[any](t, flatten.New, "name: f\ntype: flatten") }, + data: func(int) []byte { return []byte(`{"a":{"b":1}}`) }, + records: 3, + withOutput: true, + wantEmitted: 3, + }, + { + name: "replace transforms one to one", + task: func(t *testing.T) task.Task { + return newTask[any](t, replace.New, "name: r\ntype: replace\nexpression: a\nreplacement: b") + }, + data: func(int) []byte { return []byte(`aaa`) }, + records: 3, + withOutput: true, + wantEmitted: 3, + }, + { + name: "replace with no output still drains", + task: func(t *testing.T) task.Task { + return newTask[any](t, replace.New, "name: r\ntype: replace\nexpression: a\nreplacement: b") + }, + data: func(int) []byte { return []byte(`aaa`) }, + records: 3, + withOutput: false, + wantEmitted: 0, + }, + { + // the shape that used to need a declared count: four lines out of one record + name: "split fans out", + task: func(t *testing.T) task.Task { return newTask[any](t, split.New, "name: s\ntype: split") }, + data: func(int) []byte { return []byte("1\n2\n3\n4") }, + records: 3, + withOutput: true, + wantEmitted: 12, + }, + { + name: "jq explodes an array", + task: func(t *testing.T) task.Task { + return newTask[any](t, jq.New, "name: j\ntype: jq\npath: .items\nexplode: true") + }, + data: func(int) []byte { return []byte(`{"items":[1,2,3]}`) }, + records: 4, + withOutput: true, + wantEmitted: 12, + }, + { + name: "jq dropping every record", + task: func(t *testing.T) task.Task { return newTask[any](t, jq.New, "name: j\ntype: jq\npath: .missing") }, + data: func(int) []byte { return []byte(`{"items":[1,2,3]}`) }, + records: 3, + withOutput: true, + wantEmitted: 0, + }, + { + name: "compress round-trips", + task: func(t *testing.T) task.Task { + return newTask[any](t, compress.New, "name: c\ntype: compress\nformat: gzip") + }, + data: func(int) []byte { return []byte(`hello world`) }, + records: 3, + withOutput: true, + wantEmitted: 3, + }, + { + name: "compress drops empty records", + task: func(t *testing.T) task.Task { + return newTask[any](t, compress.New, "name: c\ntype: compress\nformat: gzip") + }, + data: func(int) []byte { return nil }, + records: 3, + withOutput: true, + wantEmitted: 0, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + + records, acks := tracked(tc.records, tc.data) + + emitted, err := runTask(t, tc.task(t), records, tc.withOutput) + if err != nil { + t.Fatalf("Run returned %v", err) + } + if len(emitted) != tc.wantEmitted { + t.Errorf("emitted %d records, want %d", len(emitted), tc.wantEmitted) + } + + // a task that emitted nothing has legitimately finished with its input, + // so it may settle right away. One that emitted must not have. + if tc.wantEmitted > 0 { + assertOutstanding(t, acks) + } + + sinkAll(emitted) + + assertSettled(t, acks, false) + + }) + } + +} + +// assertOutstanding checks no ack settled while its descendants are still in +// flight — the premature-acknowledgement bug, which is silent data loss. +func assertOutstanding(t *testing.T, acks []*ack.Ack) { + + t.Helper() + + for i, a := range acks { + select { + case <-a.Wait(): + t.Fatalf("record %d settled while its emitted records were still outstanding", i) + default: + } + } + +} + +// TestSamplersSettleDroppedRecords: a sampler that discards a record still has to +// settle it, and one that buffers must not settle until it drains. +func TestSamplersSettleDroppedRecords(t *testing.T) { + + tests := []struct { + name string + spec string + records int + wantEmitted int + }{ + {name: "head keeps the first two", spec: "name: s\ntype: sample\nfilter: head\nlimit: 2", records: 6, wantEmitted: 2}, + {name: "nth keeps every third", spec: "name: s\ntype: sample\nfilter: nth\ndivider: 3", records: 9, wantEmitted: 3}, + {name: "tail keeps the last two", spec: "name: s\ntype: sample\nfilter: tail\nlimit: 2", records: 6, wantEmitted: 2}, + {name: "percent keeps none", spec: "name: s\ntype: sample\nfilter: percent\npercent: 0", records: 6, wantEmitted: 0}, + {name: "percent keeps all", spec: "name: s\ntype: sample\nfilter: percent\npercent: 100", records: 6, wantEmitted: 6}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + + records, acks := tracked(tc.records, func(i int) []byte { return []byte{byte('a' + i)} }) + + emitted, err := runTask(t, newTask[any](t, sample.New, tc.spec), records, true) + if err != nil { + t.Fatalf("Run returned %v", err) + } + if len(emitted) != tc.wantEmitted { + t.Errorf("emitted %d, want %d", len(emitted), tc.wantEmitted) + } + + // the sinks handle the survivors; the records the sampler discarded were + // already settled by the sampler itself + sinkAll(emitted) + + assertSettled(t, acks, false) + + }) + } + +} + +// TestJoinHoldsInputsUntilFlush is the aggregation invariant: buffered records +// must stay unsettled until the joined record they went into is handled, or a +// source acknowledges data that has not been written anywhere. +func TestJoinHoldsInputsUntilFlush(t *testing.T) { + + const records = 5 + + srcs, acks := tracked(records, func(i int) []byte { return []byte{byte('a' + i)} }) + + input := make(chan *record.Record, records) + for _, r := range srcs { + input <- r + } + close(input) + + output := make(chan *record.Record, 16) + + tk := newTask[any](t, join.New, "name: j\ntype: join\nnumber: 5") + + if err := tk.Run(input, output); err != nil { + t.Fatalf("Run returned %v", err) + } + close(output) + + joined := make([]*record.Record, 0, 1) + for r := range output { + joined = append(joined, r) + } + if len(joined) != 1 { + t.Fatalf("expected 1 joined record, got %d", len(joined)) + } + + // the joined record exists but nothing has handled it yet + assertOutstanding(t, acks) + + // a sink handles it: that settles every record that went into it + ack.Release(joined[0].Context) + + assertSettled(t, acks, false) + +} + +// TestJoinFailurePropagates: if the joined record can't be written, every record +// that went into it must be left for redelivery. +func TestJoinFailurePropagates(t *testing.T) { + + const records = 3 + + srcs, acks := tracked(records, func(i int) []byte { return []byte{byte('a' + i)} }) + + input := make(chan *record.Record, records) + for _, r := range srcs { + input <- r + } + close(input) + + output := make(chan *record.Record, 16) + + tk := newTask[any](t, join.New, "name: j\ntype: join\nnumber: 3") + if err := tk.Run(input, output); err != nil { + t.Fatalf("Run returned %v", err) + } + close(output) + + for r := range output { + ack.Reject(r.Context) + } + + assertSettled(t, acks, true) + +} + +// TestUntrackedPipelineUnaffected: with no acking source, the mechanism is inert +// and tasks behave exactly as they did before it existed. +func TestUntrackedPipelineUnaffected(t *testing.T) { + + records := make([]*record.Record, 4) + for i := range records { + records[i] = &record.Record{ + ID: i + 1, + Data: []byte("1\n2\n3"), + Context: context.Background(), + } + } + + emitted, err := runTask(t, newTask[any](t, split.New, "name: s\ntype: split"), records, true) + if err != nil { + t.Fatalf("Run returned %v", err) + } + if len(emitted) != 12 { + t.Errorf("emitted %d records, want 12", len(emitted)) + } + +} diff --git a/internal/pkg/pipeline/task/archive/tar.go b/internal/pkg/pipeline/task/archive/tar.go index e2d1198..644f0c0 100644 --- a/internal/pkg/pipeline/task/archive/tar.go +++ b/internal/pkg/pipeline/task/archive/tar.go @@ -3,11 +3,13 @@ package archive import ( "archive/tar" "bytes" + "context" "io" "log" "path/filepath" "strings" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" "github.com/patterninc/caterpillar/internal/pkg/textutil" @@ -27,6 +29,7 @@ func (t *tarArchive) Read() { } if len(rc.Data) == 0 { + ack.Release(rc.Context) continue } @@ -44,16 +47,20 @@ func (t *tarArchive) Read() { } // check the file type is regular file - if header.Typeflag == tar.TypeReg { - buf := make([]byte, header.Size) - if _, err := io.ReadFull(r, buf); err != nil && err != io.EOF { - log.Fatal(err) - } - rc.SetContextValue(string(task.CtxKeyArchiveFileNameWrite), textutil.SlugifyFileName(filepath.Base(header.Name))) - t.SendData(rc.Context, buf, t.OutputChan) + if header.Typeflag != tar.TypeReg { + continue } + buf := make([]byte, header.Size) + if _, err := io.ReadFull(r, buf); err != nil && err != io.EOF { + log.Fatal(err) + } + + rc.SetContextValue(string(task.CtxKeyArchiveFileNameWrite), textutil.SlugifyFileName(filepath.Base(header.Name))) + t.SendData(rc.Context, buf, t.OutputChan) } + + ack.Release(rc.Context) } } @@ -62,12 +69,15 @@ func (t *tarArchive) Write() { var buf bytes.Buffer tw := tar.NewWriter(&buf) var rc record.Record + var ctxs []context.Context for { rec, ok := t.GetRecord(t.InputChan) if !ok { break } + ctxs = append(ctxs, rec.Context) + b := rec.Data if len(b) == 0 { @@ -105,5 +115,9 @@ func (t *tarArchive) Write() { log.Fatal(err) } - t.SendData(rc.Context, buf.Bytes(), t.OutputChan) + // fan-in: the archive record is produced from every input consumed above, + // so its ack must transitively complete all of theirs. + joinedAck := ack.Joined(ctxs...) + + t.SendData(ack.WithContext(rc.Context, joinedAck), buf.Bytes(), t.OutputChan) } diff --git a/internal/pkg/pipeline/task/archive/zip.go b/internal/pkg/pipeline/task/archive/zip.go index bf0b694..0238fb5 100644 --- a/internal/pkg/pipeline/task/archive/zip.go +++ b/internal/pkg/pipeline/task/archive/zip.go @@ -3,11 +3,13 @@ package archive import ( "archive/zip" "bytes" + "context" "io" "log" "path/filepath" "strings" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" "github.com/patterninc/caterpillar/internal/pkg/textutil" @@ -26,6 +28,7 @@ func (z *zipArchive) Read() { } if len(rc.Data) == 0 { + ack.Release(rc.Context) continue } @@ -35,6 +38,7 @@ func (z *zipArchive) Read() { if err != nil { log.Fatal(err) } + for _, f := range r.File { // check the file type is regular file @@ -58,6 +62,8 @@ func (z *zipArchive) Read() { z.SendData(rc.Context, buf, z.OutputChan) } } + + ack.Release(rc.Context) } } @@ -66,6 +72,7 @@ func (z *zipArchive) Write() { zipBuf := new(bytes.Buffer) zipWriter := zip.NewWriter(zipBuf) var rc record.Record + var ctxs []context.Context for { rec, ok := z.GetRecord(z.InputChan) @@ -94,13 +101,18 @@ func (z *zipArchive) Write() { } rc.Context = rec.Context + ctxs = append(ctxs, rec.Context) } if err := zipWriter.Close(); err != nil { log.Fatal(err) } + // fan-in: the archive record is produced from every input consumed above, + // so its ack must transitively complete all of theirs. + joinedAck := ack.Joined(ctxs...) + // Send the complete ZIP archive - z.SendData(rc.Context, zipBuf.Bytes(), z.OutputChan) + z.SendData(ack.WithContext(rc.Context, joinedAck), zipBuf.Bytes(), z.OutputChan) } diff --git a/internal/pkg/pipeline/task/aws/parameter_store/parameter_store.go b/internal/pkg/pipeline/task/aws/parameter_store/parameter_store.go index f7f6fab..95e2d4f 100644 --- a/internal/pkg/pipeline/task/aws/parameter_store/parameter_store.go +++ b/internal/pkg/pipeline/task/aws/parameter_store/parameter_store.go @@ -15,6 +15,7 @@ import ( "github.com/patterninc/caterpillar/internal/pkg/config" "github.com/patterninc/caterpillar/internal/pkg/duration" "github.com/patterninc/caterpillar/internal/pkg/jq" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" ) @@ -138,12 +139,14 @@ func (p *parameterStore) Run(input <-chan *record.Record, output chan<- *record. // so a single misconfigured tenant need not stop every other one if errors.Is(err, errParameterNotFound) && p.OnMissing == onMissingSkip { fmt.Printf("WARN: skipping record: %s\n", err) + ack.Release(r.Context) continue } - return err + return ack.Rejected(r.Context, err) } p.SendRecord(r, output) + ack.Release(r.Context) } return nil diff --git a/internal/pkg/pipeline/task/compress/compress.go b/internal/pkg/pipeline/task/compress/compress.go index be7e20a..fa947ca 100644 --- a/internal/pkg/pipeline/task/compress/compress.go +++ b/internal/pkg/pipeline/task/compress/compress.go @@ -5,6 +5,7 @@ import ( "fmt" "io" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" ) @@ -61,6 +62,7 @@ func (c *core) Run(input <-chan *record.Record, output chan<- *record.Record) (e // skip empty records if len(r.Data) == 0 { + ack.Release(r.Context) continue } @@ -68,22 +70,22 @@ func (c *core) Run(input <-chan *record.Record, output chan<- *record.Record) (e var err error if c.Action == defaultAction { if transformedData, err = c.compress(r); err != nil { - return err + return ack.Rejected(r.Context, err) } } else { if transformedData, err = c.decompress(r); err != nil { - return err + return ack.Rejected(r.Context, err) } } // skip empty transformed data if len(transformedData) == 0 { + ack.Release(r.Context) continue } - if output != nil { - c.SendData(r.Context, transformedData, output) - } + c.SendData(r.Context, transformedData, output) + ack.Release(r.Context) } return nil diff --git a/internal/pkg/pipeline/task/converter/converter.go b/internal/pkg/pipeline/task/converter/converter.go index 40827b1..c2dfb88 100644 --- a/internal/pkg/pipeline/task/converter/converter.go +++ b/internal/pkg/pipeline/task/converter/converter.go @@ -3,6 +3,7 @@ package converter import ( "fmt" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" ) @@ -79,7 +80,7 @@ func (c *core) Run(input <-chan *record.Record, output chan<- *record.Record) er outputs, err := c.convert(r.Data, c.Delimiter) if err != nil { - return err + return ack.Rejected(r.Context, err) } for _, out := range outputs { @@ -92,6 +93,8 @@ func (c *core) Run(input <-chan *record.Record, output chan<- *record.Record) er c.SendData(r.Context, out.Data, output) } } + + ack.Release(r.Context) } return nil diff --git a/internal/pkg/pipeline/task/delay/delay.go b/internal/pkg/pipeline/task/delay/delay.go index 0bf7927..4e1caf2 100644 --- a/internal/pkg/pipeline/task/delay/delay.go +++ b/internal/pkg/pipeline/task/delay/delay.go @@ -4,6 +4,7 @@ import ( "time" "github.com/patterninc/caterpillar/internal/pkg/duration" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" ) @@ -37,6 +38,7 @@ func (d *delay) Run(input <-chan *record.Record, output chan<- *record.Record) e } d.SendRecord(r, output) + ack.Release(r.Context) } return nil diff --git a/internal/pkg/pipeline/task/echo/echo.go b/internal/pkg/pipeline/task/echo/echo.go index 2370ffe..ed698e9 100644 --- a/internal/pkg/pipeline/task/echo/echo.go +++ b/internal/pkg/pipeline/task/echo/echo.go @@ -4,6 +4,7 @@ import ( "fmt" "time" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" ) @@ -40,9 +41,8 @@ func (e *echo) Run(input <-chan *record.Record, output chan<- *record.Record) (e // fmt.Println(r.GetContextValue(`names`)) fmt.Println(time.Now().Format(timeNowFormat), `-`, e.Name, `-`, string(item)) - if output != nil { - e.SendRecord(r, output) - } + e.SendRecord(r, output) + ack.Release(r.Context) } return diff --git a/internal/pkg/pipeline/task/file/file.go b/internal/pkg/pipeline/task/file/file.go index 5b92c7c..19bc5d0 100644 --- a/internal/pkg/pipeline/task/file/file.go +++ b/internal/pkg/pipeline/task/file/file.go @@ -11,6 +11,7 @@ import ( "strings" "github.com/patterninc/caterpillar/internal/pkg/config" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" "github.com/patterninc/caterpillar/internal/pkg/textutil" @@ -158,6 +159,21 @@ func (f *file) readFile(output chan<- *record.Record) error { } +// abort settles rc as failed, then returns err, so a source deferring +// acknowledgement redelivers the record rather than waiting forever for a +// write that will never happen. +// +// Only rc is rejected, never the records queued behind it: sibling workers are +// still running and will write those, and the pipeline rejects whatever is +// genuinely left over once every worker has returned. +func (f *file) abort(rc *record.Record, err error) error { + + ack.Reject(rc.Context) + + return err + +} + func (f *file) writeFile(input <-chan *record.Record) error { for { @@ -169,13 +185,13 @@ func (f *file) writeFile(input <-chan *record.Record) error { // Evaluate the path with the record context path, err := f.Path.Get(rc) if err != nil { - return err + return f.abort(rc, err) } // Determine the scheme from the evaluated path parsedURL, err := url.Parse(path) if err != nil { - return err + return f.abort(rc, err) } pathScheme := parsedURL.Scheme if pathScheme == `` { @@ -198,11 +214,13 @@ func (f *file) writeFile(input <-chan *record.Record) error { writerFunction, found := writers[pathScheme] if !found { - return unknownSchemeError(pathScheme) + return f.abort(rc, unknownSchemeError(pathScheme)) } if err := writerFunction(&fs, rc, bytes.NewReader(rc.Data)); err != nil { - return err + return f.abort(rc, err) } + + ack.Release(rc.Context) } return nil diff --git a/internal/pkg/pipeline/task/flatten/flatten.go b/internal/pkg/pipeline/task/flatten/flatten.go index c36ab46..74704fb 100644 --- a/internal/pkg/pipeline/task/flatten/flatten.go +++ b/internal/pkg/pipeline/task/flatten/flatten.go @@ -3,6 +3,7 @@ package flatten import ( "encoding/json" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" ) @@ -26,7 +27,7 @@ func (f *flatten) Run(input <-chan *record.Record, output chan<- *record.Record) var data map[string]any if err := json.Unmarshal(r.Data, &data); err != nil { - return err + return ack.Rejected(r.Context, err) } flat := make(map[string]any) @@ -38,10 +39,11 @@ func (f *flatten) Run(input <-chan *record.Record, output chan<- *record.Record) flatJson, err := json.Marshal(flat) if err != nil { - return err + return ack.Rejected(r.Context, err) } f.SendData(r.Context, flatJson, output) + ack.Release(r.Context) } return nil diff --git a/internal/pkg/pipeline/task/heimdall/heimdall.go b/internal/pkg/pipeline/task/heimdall/heimdall.go index c96f944..98932fe 100644 --- a/internal/pkg/pipeline/task/heimdall/heimdall.go +++ b/internal/pkg/pipeline/task/heimdall/heimdall.go @@ -1,12 +1,14 @@ package heimdall import ( + "context" "encoding/json" "fmt" "net/http" "time" "github.com/patterninc/caterpillar/internal/pkg/duration" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" ) @@ -69,25 +71,26 @@ func (h *heimdall) Run(input <-chan *record.Record, output chan<- *record.Record // Parse the input record to get dynamic context var jobContext map[string]any if err := json.Unmarshal([]byte(rc.Data), &jobContext); err != nil { - return err + return ack.Rejected(rc.Context, err) } // Create a job request with the dynamic context jobReq := h.buildJobRequest(jobContext) - if err := h.submitJob(jobReq, output); err != nil { + if err := h.submitJob(rc.Context, jobReq, output); err != nil { if !h.SkipOnError { - return err + return ack.Rejected(rc.Context, err) } fmt.Printf("WARN: skipping failed record in %s: %s\n", h.GetName(), err) - continue } + + ack.Release(rc.Context) } return nil } // Create a job request with the configured context jobReq := h.buildJobRequest(h.JobRequest.Context) - return h.submitJob(jobReq, output) + return h.submitJob(ctx, jobReq, output) } @@ -107,7 +110,7 @@ func (h *heimdall) buildJobRequest(context map[string]any) *jobRequest { } } -func (h *heimdall) submitJob(jobReq *jobRequest, output chan<- *record.Record) error { +func (h *heimdall) submitJob(srcCtx context.Context, jobReq *jobRequest, output chan<- *record.Record) error { response := &response{} @@ -125,14 +128,14 @@ func (h *heimdall) submitJob(jobReq *jobRequest, output chan<- *record.Record) e if !h.GetResult { return nil } - return h.sendToOutput(response.Result, output) + return h.sendToOutput(srcCtx, response.Result, output) } // For asynchronous jobs, poll until completion - return h.processAsyncJob(response.ID, output) + return h.processAsyncJob(srcCtx, response.ID, output) } -func (h *heimdall) processAsyncJob(jobID string, output chan<- *record.Record) error { +func (h *heimdall) processAsyncJob(srcCtx context.Context, jobID string, output chan<- *record.Record) error { // Set timeout for job polling endTime := time.Now().Add(time.Duration(h.Timeout)) @@ -155,7 +158,7 @@ func (h *heimdall) processAsyncJob(jobID string, output chan<- *record.Record) e if err := h.api(http.MethodGet, fmt.Sprintf(h.Endpoint+endpointJobResult, jobID), nil, result); err != nil { return err } - return h.sendToOutput(result, output) + return h.sendToOutput(srcCtx, result, output) case jobStatusFailed: return fmt.Errorf("job id %s failed", jobID) } diff --git a/internal/pkg/pipeline/task/heimdall/result.go b/internal/pkg/pipeline/task/heimdall/result.go index 52484bd..ec9e891 100644 --- a/internal/pkg/pipeline/task/heimdall/result.go +++ b/internal/pkg/pipeline/task/heimdall/result.go @@ -46,7 +46,10 @@ func (r *result) toSlice() ([][]byte, error) { } -func (h *heimdall) sendToOutput(result *result, output chan<- *record.Record) error { +// srcCtx carries the ack of the record the job was built from, so the results +// descend from it. A fresh context here would emit untracked records and let the +// source acknowledge before any of them was written. +func (h *heimdall) sendToOutput(srcCtx context.Context, result *result, output chan<- *record.Record) error { items, err := result.toSlice() if err != nil { @@ -54,7 +57,7 @@ func (h *heimdall) sendToOutput(result *result, output chan<- *record.Record) er } for _, item := range items { - h.SendData(ctx, item, output) + h.SendData(srcCtx, item, output) } return nil diff --git a/internal/pkg/pipeline/task/http/http.go b/internal/pkg/pipeline/task/http/http.go index 6f1b1d7..d7244af 100644 --- a/internal/pkg/pipeline/task/http/http.go +++ b/internal/pkg/pipeline/task/http/http.go @@ -14,6 +14,7 @@ import ( "github.com/patterninc/caterpillar/internal/pkg/config" "github.com/patterninc/caterpillar/internal/pkg/duration" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task/http/status" @@ -141,11 +142,15 @@ func (h *httpCore) Run(input <-chan *record.Record, output chan<- *record.Record // let's get our http object newHttp, err := h.newFromInput(rc.Data) if err != nil { - return err + return ack.Rejected(rc.Context, err) } if err := newHttp.processItem(rc, output); err != nil { - return err + return ack.Rejected(rc.Context, err) } + + // one release per input record, however many pages it produced: each + // page's send registered its own branch. + ack.Release(rc.Context) } } diff --git a/internal/pkg/pipeline/task/join/join.go b/internal/pkg/pipeline/task/join/join.go index 00d06f5..140676c 100644 --- a/internal/pkg/pipeline/task/join/join.go +++ b/internal/pkg/pipeline/task/join/join.go @@ -7,6 +7,7 @@ import ( "time" "github.com/patterninc/caterpillar/internal/pkg/duration" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" ) @@ -55,11 +56,14 @@ func (j *join) Run(input <-chan *record.Record, output chan<- *record.Record) er tickerCh = ticker.C } + // the input receive has to live inside the select: with a default case the + // select never blocks, so tickerCh is only serviced between records and + // never fires while input is stalled - which is exactly when a partially + // filled buffer needs flushing. With duration unset tickerCh is nil and + // this is a plain blocking receive on input. for { select { - default: - // Try to get a record from input - r, ok := j.GetRecord(input) + case r, ok := <-input: if !ok { // Input channel closed, send any remaining records j.flushBuffer(&buffer, output) @@ -95,13 +99,18 @@ func (j *join) sendJoinedRecords(buffer []*record.Record, output chan<- *record. // Join all data with the specified delimiter var joinedData strings.Builder + ctxs := make([]context.Context, len(buffer)) for i, r := range buffer { if i > 0 { joinedData.WriteString(j.Delimiter) } joinedData.Write(r.Data) + ctxs[i] = r.Context } - j.SendData(ctx, []byte(joinedData.String()), output) + // fan-in: the output record is produced from every buffered input, so its + // ack must transitively complete every one of theirs. + joinedAck := ack.Joined(ctxs...) + j.SendData(ack.WithContext(ctx, joinedAck), []byte(joinedData.String()), output) } diff --git a/internal/pkg/pipeline/task/jq/jq.go b/internal/pkg/pipeline/task/jq/jq.go index 3d06881..dbc0c05 100644 --- a/internal/pkg/pipeline/task/jq/jq.go +++ b/internal/pkg/pipeline/task/jq/jq.go @@ -5,6 +5,7 @@ import ( "fmt" "github.com/patterninc/caterpillar/internal/pkg/config" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" ) @@ -23,59 +24,78 @@ func New() (task.Task, error) { func (j *jq) Run(input <-chan *record.Record, output chan<- *record.Record) (err error) { - if input != nil && output != nil { + if input == nil { + return nil + } + + if output == nil { + // terminal: the transform has no effect, but input still has to be + // drained and each record's ack settled, or a source deferring + // acknowledgement never finishes. for { r, ok := j.GetRecord(input) if !ok { - break + return nil } + ack.Release(r.Context) + } + } - // First evaluate config templates in the path - query, err := j.Path.GetJQ(r) - if err != nil { - return err - } + for { + r, ok := j.GetRecord(input) + if !ok { + break + } - // Execute the JQ query - items, err := query.Execute(r.Data) - if err != nil { - // An ignored query error costs only its own record, since it describes - // that record rather than the pipeline, and warns because the run is not - // failing over it. Otherwise it is critical: return, and let fail_on_error - // judge the verdict as it does for every other task. - if j.IgnoreError { - fmt.Printf("WARN: %s: skipping record %d: %s\n", j.GetName(), r.ID, err) - continue - } - return err - } - if items == nil { + // First evaluate config templates in the path + query, err := j.Path.GetJQ(r) + if err != nil { + return ack.Rejected(r.Context, err) + } + + // Execute the JQ query + items, err := query.Execute(r.Data) + if err != nil { + // An ignored query error costs only its own record, since it describes + // that record rather than the pipeline, and warns because the run is not + // failing over it. Otherwise it is critical: return, and let fail_on_error + // judge the verdict as it does for every other task. + if j.IgnoreError { + fmt.Printf("WARN: %s: skipping record %d: %s\n", j.GetName(), r.ID, err) + ack.Release(r.Context) continue } - if splitItems, ok := items.([]any); j.Explode && ok { - for _, splitItem := range splitItems { - if j.AsRaw { - j.SendData(r.Context, fmt.Appendf(nil, "%v", splitItem), output) - } else { - jsonItem, err := json.Marshal(splitItem) - if err != nil { - return err - } - j.SendData(r.Context, jsonItem, output) - } + return ack.Rejected(r.Context, err) + } + if items == nil { + ack.Release(r.Context) + continue + } + if splitItems, ok := items.([]any); j.Explode && ok { + for _, splitItem := range splitItems { + if j.AsRaw { + j.SendData(r.Context, fmt.Appendf(nil, "%v", splitItem), output) + continue } + jsonItem, err := json.Marshal(splitItem) + if err != nil { + return ack.Rejected(r.Context, err) + } + j.SendData(r.Context, jsonItem, output) + } + } else { + if j.AsRaw { + j.SendData(r.Context, fmt.Appendf(nil, "%v", items), output) } else { - if j.AsRaw { - j.SendData(r.Context, fmt.Appendf(nil, "%v", items), output) - } else { - jsonItem, err := json.Marshal(items) - if err != nil { - return err - } - j.SendData(r.Context, jsonItem, output) + jsonItem, err := json.Marshal(items) + if err != nil { + return ack.Rejected(r.Context, err) } + j.SendData(r.Context, jsonItem, output) } } + + ack.Release(r.Context) } return diff --git a/internal/pkg/pipeline/task/kafka/kafka.go b/internal/pkg/pipeline/task/kafka/kafka.go index a5418d6..40c2739 100644 --- a/internal/pkg/pipeline/task/kafka/kafka.go +++ b/internal/pkg/pipeline/task/kafka/kafka.go @@ -10,6 +10,7 @@ import ( "github.com/google/uuid" "github.com/patterninc/caterpillar/internal/pkg/duration" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" ) @@ -144,15 +145,57 @@ func (k *kafka) write(input <-chan *record.Record) error { wg sync.WaitGroup firstDeliveryErr error ) + + // acks for records the producer has taken but not yet confirmed. A message that + // never produces a delivery report — a flush timeout, typically — would + // otherwise leave its source record unsettled and the pipeline waiting on it + // forever, so whatever is left here at the end is failed explicitly. + var ( + pendingMu sync.Mutex + pending = make(map[*ack.Ack]struct{}) + ) + + settle := func(a *ack.Ack, deliveryErr error) { + + if a == nil { + return + } + + pendingMu.Lock() + delete(pending, a) + pendingMu.Unlock() + + if deliveryErr != nil { + a.Fail() + return + } + + a.Done() + + } + wg.Add(1) go func() { defer wg.Done() for e := range deliveryCh { - if m, ok := e.(*ckafka.Message); ok && m.TopicPartition.Error != nil && firstDeliveryErr == nil { - firstDeliveryErr = m.TopicPartition.Error - fmt.Printf("delivery failed for topic %s partition %d: %v\n", - k.Topic, m.TopicPartition.Partition, m.TopicPartition.Error) + m, ok := e.(*ckafka.Message) + if !ok { + continue + } + a, _ := m.Opaque.(*ack.Ack) + if m.TopicPartition.Error != nil { + if firstDeliveryErr == nil { + firstDeliveryErr = m.TopicPartition.Error + fmt.Printf("delivery failed for topic %s partition %d: %v\n", + k.Topic, m.TopicPartition.Partition, m.TopicPartition.Error) + } + // the source record never made it to the topic, so fail its + // ack: the source must leave it unacknowledged for retry. + settle(a, m.TopicPartition.Error) + continue } + // settle only on broker-confirmed delivery, not on local enqueue. + settle(a, nil) } }() @@ -166,14 +209,27 @@ func (k *kafka) write(input <-chan *record.Record) error { msgBytes, err := codec.serialize(k.Topic, r.Data) if err != nil { produceErr = fmt.Errorf("failed to serialize record for topic %s: %w", k.Topic, err) + ack.Reject(r.Context) break } + var opaque any + a, tracked := ack.FromContext(r.Context) + if tracked { + opaque = a + pendingMu.Lock() + pending[a] = struct{}{} + pendingMu.Unlock() + } + if err = p.Produce(&ckafka.Message{ TopicPartition: ckafka.TopicPartition{Topic: &k.Topic, Partition: ckafka.PartitionAny}, Value: msgBytes, + Opaque: opaque, }, deliveryCh); err != nil { produceErr = fmt.Errorf("failed to enqueue message to topic %s: %w", k.Topic, err) + // the producer never took it, so no delivery report is coming + settle(a, produceErr) break } } @@ -187,6 +243,14 @@ func (k *kafka) write(input <-chan *record.Record) error { close(deliveryCh) wg.Wait() + // no delivery report ever arrived for these, so nothing else will settle them + pendingMu.Lock() + for a := range pending { + a.Fail() + } + pending = nil + pendingMu.Unlock() + if produceErr != nil { return produceErr } diff --git a/internal/pkg/pipeline/task/replace/replace.go b/internal/pkg/pipeline/task/replace/replace.go index e7d724a..f57ebb8 100644 --- a/internal/pkg/pipeline/task/replace/replace.go +++ b/internal/pkg/pipeline/task/replace/replace.go @@ -3,6 +3,7 @@ package replace import ( "regexp" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" ) @@ -24,14 +25,30 @@ func (r *replace) Run(input <-chan *record.Record, output chan<- *record.Record) return err } - if output != nil { + if input == nil { + return nil + } + + if output == nil { + // terminal: the replacement has no effect, but input still has to be + // drained and each record's ack settled, or a source deferring + // acknowledgement never finishes. for { record, ok := r.GetRecord(input) if !ok { - break + return nil } - r.SendData(record.Context, []byte(rx.ReplaceAllString(string(record.Data), r.Replacement)), output) + ack.Release(record.Context) + } + } + + for { + record, ok := r.GetRecord(input) + if !ok { + break } + r.SendData(record.Context, []byte(rx.ReplaceAllString(string(record.Data), r.Replacement)), output) + ack.Release(record.Context) } return nil diff --git a/internal/pkg/pipeline/task/sample/head.go b/internal/pkg/pipeline/task/sample/head.go index 79147ab..cfd55ca 100644 --- a/internal/pkg/pipeline/task/sample/head.go +++ b/internal/pkg/pipeline/task/sample/head.go @@ -1,6 +1,7 @@ package sample import ( + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" ) @@ -26,6 +27,8 @@ func (h *head) filter(r *record.Record, output chan<- *record.Record) error { h.index++ } + ack.Release(r.Context) + return nil } diff --git a/internal/pkg/pipeline/task/sample/nth.go b/internal/pkg/pipeline/task/sample/nth.go index d05618d..eaeaf88 100644 --- a/internal/pkg/pipeline/task/sample/nth.go +++ b/internal/pkg/pipeline/task/sample/nth.go @@ -1,6 +1,7 @@ package sample import ( + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" ) @@ -25,6 +26,8 @@ func (n *nth) filter(r *record.Record, output chan<- *record.Record) error { n.sendRecord(r, output) } + ack.Release(r.Context) + n.index++ return nil diff --git a/internal/pkg/pipeline/task/sample/percent.go b/internal/pkg/pipeline/task/sample/percent.go index 3847a56..d12575c 100644 --- a/internal/pkg/pipeline/task/sample/percent.go +++ b/internal/pkg/pipeline/task/sample/percent.go @@ -5,6 +5,7 @@ import ( "fmt" "math/big" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" ) @@ -44,6 +45,8 @@ func (p *percent) filter(r *record.Record, output chan<- *record.Record) error { p.sendRecord(r, output) } + ack.Release(r.Context) + return nil } diff --git a/internal/pkg/pipeline/task/sample/random.go b/internal/pkg/pipeline/task/sample/random.go index dfb07cd..1de5730 100644 --- a/internal/pkg/pipeline/task/sample/random.go +++ b/internal/pkg/pipeline/task/sample/random.go @@ -4,6 +4,7 @@ import ( "crypto/rand" "math/big" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" ) @@ -29,6 +30,8 @@ func (r *random) filter(row *record.Record, _ chan<- *record.Record) error { if len(r.buffer) < r.size { r.buffer = append(r.buffer, row) + } else { + ack.Release(row.Context) } return nil @@ -38,16 +41,28 @@ func (r *random) filter(row *record.Record, _ chan<- *record.Record) error { func (r *random) drain(output chan<- *record.Record) error { if l := int64(len(r.buffer)); l > 0 { + + // draws are with replacement, so a buffered record can be sent zero, one, or + // many times. Each send registers itself, so no tally is needed. for i := 0; i < r.limit; i++ { index, err := rand.Int(rand.Reader, big.NewInt(l)) if err != nil { + // bailing out mid-draw leaves the buffer unreleased otherwise + for _, row := range r.buffer { + ack.Reject(row.Context) + } return err } r.sendRecord(r.buffer[index.Int64()], output) } + + for _, row := range r.buffer { + ack.Release(row.Context) + } + } return nil diff --git a/internal/pkg/pipeline/task/sample/sample.go b/internal/pkg/pipeline/task/sample/sample.go index 80d6fd2..0f67c4c 100644 --- a/internal/pkg/pipeline/task/sample/sample.go +++ b/internal/pkg/pipeline/task/sample/sample.go @@ -3,6 +3,7 @@ package sample import ( "fmt" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" ) @@ -72,7 +73,7 @@ func (s *sample) Run(input <-chan *record.Record, output chan<- *record.Record) break } if err := sampler.filter(r, output); err != nil { - return err + return ack.Rejected(r.Context, err) } } diff --git a/internal/pkg/pipeline/task/sample/tail.go b/internal/pkg/pipeline/task/sample/tail.go index e394d6b..2892708 100644 --- a/internal/pkg/pipeline/task/sample/tail.go +++ b/internal/pkg/pipeline/task/sample/tail.go @@ -1,6 +1,7 @@ package sample import ( + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" ) @@ -24,6 +25,12 @@ func newTail(s *sample) (sampler, error) { func (t *tail) filter(r *record.Record, _ chan<- *record.Record) error { + // the ring buffer is about to overwrite this slot; the evicted record will + // never be forwarded, so settle its ack here. + if evicted := t.buffer[t.index]; evicted != nil { + ack.Release(evicted.Context) + } + t.buffer[t.index] = r t.index = (t.index + 1) % t.limit t.count++ @@ -45,6 +52,13 @@ func (t *tail) drain(output chan<- *record.Record) error { t.sendRecord(t.buffer[(start+i)%len(t.buffer)], output) } + // the ring held these back across iterations, so it owes each one a release + for _, row := range t.buffer { + if row != nil { + ack.Release(row.Context) + } + } + return nil } diff --git a/internal/pkg/pipeline/task/sftp/operations.go b/internal/pkg/pipeline/task/sftp/operations.go index 5eecfd9..c17f8ce 100644 --- a/internal/pkg/pipeline/task/sftp/operations.go +++ b/internal/pkg/pipeline/task/sftp/operations.go @@ -10,6 +10,7 @@ import ( "github.com/bmatcuk/doublestar" pkgsftp "github.com/pkg/sftp" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" "github.com/patterninc/caterpillar/internal/pkg/textutil" @@ -29,12 +30,14 @@ func (s *sftp) upload(client *pkgsftp.Client, input <-chan *record.Record) error file, err := s.Path.Get(rc) if err != nil { - return err + return ack.Rejected(rc.Context, err) } if err := s.uploadOne(client, file, rc.Data); err != nil { - return err + return ack.Rejected(rc.Context, err) } + + ack.Release(rc.Context) } return nil diff --git a/internal/pkg/pipeline/task/sns/sns.go b/internal/pkg/pipeline/task/sns/sns.go index cec5c03..24f04f5 100644 --- a/internal/pkg/pipeline/task/sns/sns.go +++ b/internal/pkg/pipeline/task/sns/sns.go @@ -11,6 +11,7 @@ import ( "github.com/aws/aws-sdk-go-v2/service/sns/types" "github.com/google/uuid" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" ) @@ -103,8 +104,10 @@ func (s *snsTask) Run(input <-chan *record.Record, output chan<- *record.Record) _, err := s.client.Publish(r.Context, publishInput) if err != nil { - return fmt.Errorf("failed to publish to SNS topic %s: %w", s.TopicArn, err) + return ack.Rejected(r.Context, fmt.Errorf("failed to publish to SNS topic %s: %w", s.TopicArn, err)) } + + ack.Release(r.Context) } return nil diff --git a/internal/pkg/pipeline/task/split/split.go b/internal/pkg/pipeline/task/split/split.go index 51029a7..6b51972 100644 --- a/internal/pkg/pipeline/task/split/split.go +++ b/internal/pkg/pipeline/task/split/split.go @@ -3,6 +3,7 @@ package split import ( "strings" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" ) @@ -32,6 +33,7 @@ func (s *split) Run(input <-chan *record.Record, output chan<- *record.Record) e for _, line := range lines { s.SendData(r.Context, []byte(line), output) } + ack.Release(r.Context) } return nil diff --git a/internal/pkg/pipeline/task/sqs/README.md b/internal/pkg/pipeline/task/sqs/README.md index 696d958..faddf38 100644 --- a/internal/pkg/pipeline/task/sqs/README.md +++ b/internal/pkg/pipeline/task/sqs/README.md @@ -18,7 +18,7 @@ The task automatically determines its mode based on the presence of input/output | `name` | string | - | Task name for identification | | `type` | string | `sqs` | Must be "sqs" | | `queue_url` | string | - | SQS queue URL (required) | -| `concurrency` | int | `10` | Number of concurrent message processors | +| `concurrency` | int | `10` | Number of concurrent workers that acknowledge (delete) fully-processed messages | | `max_messages` | int | `10` | Maximum number of messages to receive per batch | | `wait_time_seconds` | int | `10` | Long polling wait time in seconds | | `exit_on_empty` | bool | `false` | Exit when queue is empty | @@ -70,9 +70,32 @@ tasks: queue_url: {{ env "SQS_QUEUE_URL" }} ``` +## Message Acknowledgment + +When reading from a queue, a message's receipt is deleted only once every downstream task +has finished with the record produced from it. A downstream failure leaves the receipt +alone, so SQS redelivers the message after the visibility timeout rather than losing it. +Delivery is therefore at-least-once: a pipeline may see a message more than once, but never +zero times. + +Two consequences worth tuning for: + +- **`channel_size` bounds how many messages can sit inside the pipeline at once** (see the + root README). A message waiting in a deep channel can exceed the queue's visibility + timeout, at which point SQS redelivers it while the first copy is still in flight and + the eventual delete fails on a stale receipt handle. Keep `channel_size` in proportion + to how long a record takes to traverse the pipeline, relative to the queue's visibility + timeout. Messages that have finished the pipeline but whose delete has not yet returned + are bounded only by `concurrency` (how many `DeleteMessage` calls run at once), not by + `channel_size`. +- **SQS caps in-flight messages** at 120,000 per standard queue and 20,000 per FIFO queue. + A large `channel_size` on a long pipeline can approach that; on a breach `ReceiveMessage` + returns `OverLimit` and the task stops. FIFO queues are stricter still, since + unacknowledged messages block their message group. + ## Sample Pipelines -- `test/pipelines/sqs_with_context_concurrency.yaml` - SQS read with context variables and concurrency +- `test/pipelines/sqs_with_context_concurrency.yaml` - SQS read with context variables and concurrency; run `test/pipelines/setup_localstack_sqs.sh` first to create the queue ## Use Cases diff --git a/internal/pkg/pipeline/task/sqs/sqs.go b/internal/pkg/pipeline/task/sqs/sqs.go index 1dcf05b..1fe06ca 100644 --- a/internal/pkg/pipeline/task/sqs/sqs.go +++ b/internal/pkg/pipeline/task/sqs/sqs.go @@ -6,7 +6,6 @@ import ( "fmt" "regexp" "strings" - "sync" "time" "github.com/aws/aws-sdk-go-v2/aws" @@ -14,16 +13,16 @@ import ( qs "github.com/aws/aws-sdk-go-v2/service/sqs" "github.com/google/uuid" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" ) const ( - defaultConcurrency = 10 - defaultMaxMessages = 10 - defaultWaitTimeSeconds = 10 - receiptsQueueMultiplier = 1000 - defaultRegion = "us-west-2" + defaultConcurrency = 10 + defaultMaxMessages = 10 + defaultWaitTimeSeconds = 10 + defaultRegion = "us-west-2" ) var ( @@ -40,7 +39,8 @@ type sqs struct { ExitOnEmpty bool `yaml:"exit_on_empty,omitempty" json:"exit_on_empty,omitempty"` MessageGroupId string `yaml:"message_group_id,omitempty" json:"message_group_id,omitempty"` // used for FIFO queues - client *qs.Client + client *qs.Client + tracker *ack.Tracker } func New() (task.Task, error) { @@ -67,6 +67,8 @@ func (s *sqs) Init() error { } s.client = qs.NewFromConfig(awsConfig) + s.tracker = ack.NewTracker(s.Concurrency) + return nil } @@ -88,33 +90,34 @@ func (s *sqs) extractRegionFromQueueURL() string { func (s *sqs) Run(input <-chan *record.Record, output chan<- *record.Record) error { - // Client is already initialized in RunPreHook - just use it + // Client is already initialized in Init - just use it if input != nil { return s.sendMessages(input) } - // If input is nil, act as a source: start getMessages and receipt workers - // let's create channel to which getMessages function will communicate messages receipts - receipts := make(chan *string, s.Concurrency*receiptsQueueMultiplier) + // If input is nil, act as a source: read messages and hand each receipt to + // the tracker, which deletes it once every downstream task has finished + // with the record it produced. Finish, not Run, waits for those deletions. + return s.getMessages(ctx, output) - // we set a pool of workers that will delete messages from the queue - var wg sync.WaitGroup - wg.Add(s.Concurrency) - for i := 0; i < s.Concurrency; i++ { - go s.processReceipts(receipts, &wg) - } +} - err := s.getMessages(ctx, output, receipts) +// Finish waits for every deferred deletion before the pipeline treats this +// task as complete, so a shutdown never abandons in-flight acknowledgements. +// It can't happen in Run: downstream tasks that emit only once their input +// closes can't finish with a record until this task's output channel is +// closed, which the pipeline does only after Run returns. +func (s *sqs) Finish() error { - wg.Wait() + if s.tracker != nil { + s.tracker.Wait() + } - return err + return nil } -func (s *sqs) getMessages(ctx context.Context, output chan<- *record.Record, receipts chan *string) error { - - defer close(receipts) +func (s *sqs) getMessages(ctx context.Context, output chan<- *record.Record) error { // do we need to stop pipeline after a while? if s.EndAfter > 0 { @@ -156,34 +159,59 @@ func (s *sqs) getMessages(ctx context.Context, output chan<- *record.Record, rec } for _, m := range receiveMessageOutput.Messages { - // create new record and send it downstream - if output != nil { - s.SendData(ctx, []byte(*m.Body), output) + // nothing to forward to, so there's no downstream ack to wait + // for: delete the receipt right away. + if output == nil { + s.deleteMessage(m.MessageId, m.ReceiptHandle) + continue } - // send receipt to receipts channel for deletion - receipts <- m.ReceiptHandle + msgAck := ack.New() + s.SendData(ack.WithContext(ctx, msgAck), []byte(*m.Body), output) + + s.tracker.Track(msgAck, &messageAck{ + sqs: s, + messageId: m.MessageId, + receiptHandle: m.ReceiptHandle, + }) } } } } -func (s *sqs) processReceipts(receipts <-chan *string, wg *sync.WaitGroup) error { +// messageAck acknowledges one received message on behalf of ack.Tracker. +type messageAck struct { + sqs *sqs + messageId *string + receiptHandle *string +} - defer wg.Done() +// Ack deletes the message's receipt so SQS doesn't redeliver it. On a +// downstream failure it does nothing: leaving the receipt alone lets SQS +// redeliver the message once the visibility timeout expires. +func (m *messageAck) Ack(failed bool) { - for receipt := range receipts { - if _, err := s.client.DeleteMessage(ctx, &qs.DeleteMessageInput{ - QueueUrl: &s.QueueURL, - ReceiptHandle: receipt, - }); err != nil { - return err - } + if failed { + return } - return nil + m.sqs.deleteMessage(m.messageId, m.receiptHandle) + +} + +// deleteMessage acknowledges a message by deleting its receipt. A failure is +// logged rather than returned: the message has already been processed, so the +// worst case is a redelivery after the visibility timeout. +func (s *sqs) deleteMessage(messageId, receiptHandle *string) { + + if _, err := s.client.DeleteMessage(ctx, &qs.DeleteMessageInput{ + QueueUrl: &s.QueueURL, + ReceiptHandle: receiptHandle, + }); err != nil { + fmt.Printf("failed to delete message %s from queue %s: %v\n", aws.ToString(messageId), s.QueueURL, err) + } } @@ -203,8 +231,10 @@ func (s *sqs) sendMessages(input <-chan *record.Record) error { MessageGroupId: s.getMessageGroupID(), }) if err != nil { - return err + return ack.Rejected(r.Context, err) } + + ack.Release(r.Context) } return nil } diff --git a/internal/pkg/pipeline/task/task.go b/internal/pkg/pipeline/task/task.go index 776a067..30e9739 100644 --- a/internal/pkg/pipeline/task/task.go +++ b/internal/pkg/pipeline/task/task.go @@ -7,6 +7,7 @@ import ( "sync" "github.com/patterninc/caterpillar/internal/pkg/jq" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" ) @@ -37,6 +38,18 @@ type Task interface { Init() error // Called once after unmarshaling, before pipeline execution } +// Finisher is implemented by tasks with work to do only once their output +// channel has been closed. Deferred acknowledgement is the case this exists +// for: a source's acks settle only after every downstream task has drained, +// and a downstream task that emits on input close can't drain until the +// source's output channel is closed, so waiting inside Run would deadlock. +// +// The pipeline calls Finish exactly once per task, after every worker of that +// task has returned from Run and after the task's output channel is closed. +type Finisher interface { + Finish() error +} + type Base struct { Name string `yaml:"name,omitempty" json:"name,omitempty"` Type string `yaml:"type,omitempty" json:"type,omitempty"` @@ -84,6 +97,7 @@ func (b *Base) Run(input <-chan *record.Record, output chan<- *record.Record) er for r := range input { b.SendRecord(r, output) + ack.Release(r.Context) } return nil @@ -114,6 +128,13 @@ func (b *Base) SendRecord(r *record.Record, output chan<- *record.Record) /* we return } + // every send registers its own branch, before the record is visible to another + // goroutine. This is what removes the need for a task to declare a fan-out + // count: N sends are N branches, and the task then releases its input once. + if a, ok := ack.FromContext(r.Context); ok { + a.AddBranch(1) + } + defer func() { output <- r }() diff --git a/internal/pkg/pipeline/task/xpath/xpath.go b/internal/pkg/pipeline/task/xpath/xpath.go index 29b6f2f..5c249c7 100644 --- a/internal/pkg/pipeline/task/xpath/xpath.go +++ b/internal/pkg/pipeline/task/xpath/xpath.go @@ -9,6 +9,7 @@ import ( "github.com/antchfx/htmlquery" "golang.org/x/net/html" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task/converter" @@ -37,7 +38,7 @@ func (x *xpath) Run(input <-chan *record.Record, output chan<- *record.Record) e document, err := htmlquery.Parse(bytes.NewReader(r.Data)) if err != nil { - return err + return ack.Rejected(r.Context, err) } containerNodes := []*html.Node{document} @@ -45,9 +46,10 @@ func (x *xpath) Run(input <-chan *record.Record, output chan<- *record.Record) e containerNodes = htmlquery.Find(document, x.Container) if len(containerNodes) == 0 { if !x.IgnoreMissing { - return fmt.Errorf("no nodes found for XPath: %s", x.Container) + return ack.Rejected(r.Context, fmt.Errorf("no nodes found for XPath: %s", x.Container)) } fmt.Println("container is missing - ", x.Container) + ack.Release(r.Context) continue } } @@ -55,7 +57,7 @@ func (x *xpath) Run(input <-chan *record.Record, output chan<- *record.Record) e for i, container := range containerNodes { data, err := x.queryFields(container) if err != nil { - return err + return ack.Rejected(r.Context, err) } if len(data) != 0 { @@ -64,6 +66,8 @@ func (x *xpath) Run(input <-chan *record.Record, output chan<- *record.Record) e x.SendData(r.Context, data, output) } } + + ack.Release(r.Context) } return nil diff --git a/test/pipelines/setup_localstack_fanout.sh b/test/pipelines/setup_localstack_fanout.sh new file mode 100755 index 0000000..550cd95 --- /dev/null +++ b/test/pipelines/setup_localstack_fanout.sh @@ -0,0 +1,54 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Creates/refills the queue used by sqs_fanout_fanin_dag.yaml against a +# LocalStack instance on localhost:4566. Each message carries an array, so the +# pipeline's jq explode step fans one message out into ITEMS_PER_MESSAGE +# records. +# +# Requires: LocalStack running (community image, e.g. +# `docker run -d -p 4566:4566 localstack/localstack:4.0`), plus aws-cli and jq. + +ENDPOINT="http://localhost:4566" +REGION="us-west-2" +QUEUE_NAME="local-sqs-fanout-fanin-queue" # matches queue_url in sqs_fanout_fanin_dag.yaml +# Override either to change the shape of the run. 20 messages keeps it quick; +# above ~50 it also exercises the case where a source with a bounded +# unacknowledged-message window would stall against the join downstream. +TOTAL_MESSAGES="${TOTAL_MESSAGES:-20}" +ITEMS_PER_MESSAGE="${ITEMS_PER_MESSAGE:-5}" +BATCH_SIZE=10 # SQS SendMessageBatch max + +export AWS_ACCESS_KEY_ID="${AWS_ACCESS_KEY_ID:-test}" +export AWS_SECRET_ACCESS_KEY="${AWS_SECRET_ACCESS_KEY:-test}" +export AWS_DEFAULT_REGION="$REGION" +export AWS_PAGER= + +awscli() { + aws --endpoint-url "$ENDPOINT" --region "$REGION" "$@" +} + +echo "Creating queue '$QUEUE_NAME'..." +QUEUE_URL=$(awscli sqs create-queue --queue-name "$QUEUE_NAME" --query 'QueueUrl' --output text) +echo "Queue URL: $QUEUE_URL" + +echo "Pushing $TOTAL_MESSAGES messages of $ITEMS_PER_MESSAGE items each..." +for ((batch_start=0; batch_start/dev/null +done + +expected_records=$((TOTAL_MESSAGES * ITEMS_PER_MESSAGE * 2)) +echo +echo "Sent $TOTAL_MESSAGES messages." +echo "Expected: $expected_records records into join, $((expected_records / 4)) output files." +echo "A correct run ends with both queue depths at 0." diff --git a/test/pipelines/setup_localstack_sqs.sh b/test/pipelines/setup_localstack_sqs.sh new file mode 100755 index 0000000..0497ca8 --- /dev/null +++ b/test/pipelines/setup_localstack_sqs.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Creates/refills the local SQS queue used by sqs_with_context_concurrency.yaml +# against a LocalStack instance running on localhost:4566. +# +# Requires: LocalStack running (`localstack start` or the localstack/localstack +# docker image), plus aws-cli and jq installed locally. + +ENDPOINT="http://localhost:4566" +REGION="us-west-2" +QUEUE_NAME="local-sqs-context-concurrency-queue" # matches queue_url in sqs_with_context_concurrency.yaml +TOTAL_MESSAGES=100 +BATCH_SIZE=10 # SQS SendMessageBatch max + +export AWS_ACCESS_KEY_ID="${AWS_ACCESS_KEY_ID:-test}" +export AWS_SECRET_ACCESS_KEY="${AWS_SECRET_ACCESS_KEY:-test}" +export AWS_DEFAULT_REGION="$REGION" + +awscli() { + aws --endpoint-url "$ENDPOINT" --region "$REGION" "$@" +} + +echo "Creating queue '$QUEUE_NAME'..." +QUEUE_URL=$(awscli sqs create-queue --queue-name "$QUEUE_NAME" --query 'QueueUrl' --output text) +echo "Queue URL: $QUEUE_URL" + +echo "Pushing $TOTAL_MESSAGES random messages..." +for ((batch_start=0; batch_start/dev/null +done + +echo "Done. Sent $TOTAL_MESSAGES messages to queue '$QUEUE_NAME'." +echo +echo "Queue URL (LocalStack): $QUEUE_URL" +echo +echo "To point the pipeline at LocalStack, run it with:" +echo " AWS_ENDPOINT_URL=$ENDPOINT AWS_ACCESS_KEY_ID=test AWS_SECRET_ACCESS_KEY=test AWS_DEFAULT_REGION=$REGION go run cmd/caterpillar/caterpillar.go -conf test/pipelines/sqs_with_context_concurrency.yaml" diff --git a/test/pipelines/sqs_fanout_fanin_dag.yaml b/test/pipelines/sqs_fanout_fanin_dag.yaml new file mode 100644 index 0000000..b715cdb --- /dev/null +++ b/test/pipelines/sqs_fanout_fanin_dag.yaml @@ -0,0 +1,66 @@ +# Exercises acknowledgment across every shape that changes a record's branch +# count, in one DAG. Acks only exist when the source creates them, so this has +# to be SQS-driven to test anything: with any other source, Fanout/Joined are +# no-ops. +# +# Each message body is {"id": N, "items": [1,2,3,4,5]}. +# +# read_queue 1 record per message +# >> explode_items 5 records (task fan-out: jq explode) +# >> [tag_a, tag_b] 10 records (structural fan-out: DAG branches) +# >> batch (fan-in: join, 7 at a time) +# >> save +# +# With the default 20 messages: 20 x 5 x 2 = 200 records into join, so 28 full +# batches of 7 plus a final partial batch of 4 = 29 files, 200 records, and +# each item value appearing exactly 40 times (20 messages x 2 branches). +# +# The message is only deleted once all 10 of its branches have landed, so a +# correct run ends with the queue at 0 visible / 0 in flight. Any leak shows up +# as leftover depth; any premature ack shows up as a short record count. +# +# join's batch size is deliberately NOT a divisor of the record count, so +# records are always left buffered when the source stops reading. That is the +# case that deadlocks if a source waits for its acknowledgements before its +# output channel closes: join can only flush those last records once the +# channel closes, and the channel can only close once the source stops +# waiting. Keep it indivisible or this fixture stops testing that. +# +# Seed the queue first (see test/pipelines/setup_localstack_fanout.sh), then: +# AWS_ENDPOINT_URL=http://localhost:4566 AWS_ACCESS_KEY_ID=test \ +# AWS_SECRET_ACCESS_KEY=test AWS_DEFAULT_REGION=us-west-2 \ +# go run ./cmd/caterpillar -conf test/pipelines/sqs_fanout_fanin_dag.yaml + +tasks: + - name: read_queue + type: sqs + queue_url: http://sqs.us-west-2.localhost.localstack.cloud:4566/000000000000/local-sqs-fanout-fanin-queue + exit_on_empty: true + + # fan-out: one message becomes one record per item in the array + - name: explode_items + type: jq + path: '.items' + explode: true + + # structural fan-out: the DAG duplicates every record into both branches + - name: tag_a + type: jq + path: '{value: ., branch: "a"}' + + - name: tag_b + type: jq + path: '{value: ., branch: "b"}' + + # fan-in: seven records become one, and its ack must transitively complete + # all seven + - name: batch + type: join + number: 7 + delimiter: "\n" + + - name: save + type: file + path: ./output/fanout/{{ macro "uuid" }}.json + +dag: read_queue >> explode_items >> [tag_a,tag_b] >> batch >> save diff --git a/test/pipelines/sqs_with_context_concurrency.yaml b/test/pipelines/sqs_with_context_concurrency.yaml index 27390f0..611ea65 100644 --- a/test/pipelines/sqs_with_context_concurrency.yaml +++ b/test/pipelines/sqs_with_context_concurrency.yaml @@ -1,7 +1,7 @@ tasks: - name: read_queue type: sqs - queue_url: https://sqs.us-west-2.amazonaws.com/123456789012/my-queue + queue_url: http://sqs.us-west-2.localhost.localstack.cloud:4566/000000000000/local-sqs-context-concurrency-queue exit_on_empty: true - name: extract_urls type: jq