From 4fb713935f3d85389e9e022b1e1f4f817cc17d18 Mon Sep 17 00:00:00 2001 From: Mahesh Date: Fri, 17 Jul 2026 13:26:33 +0530 Subject: [PATCH 1/3] feat: implement acknowledgment tracking for pipeline records --- internal/pkg/pipeline/ack/ack.go | 177 ++++++++++++++++++ internal/pkg/pipeline/pipeline.go | 12 ++ internal/pkg/pipeline/task/archive/tar.go | 34 +++- internal/pkg/pipeline/task/archive/zip.go | 25 ++- .../aws/parameter_store/parameter_store.go | 9 + .../pkg/pipeline/task/compress/compress.go | 7 + internal/pkg/pipeline/task/echo/echo.go | 3 + internal/pkg/pipeline/task/file/file.go | 5 + .../pkg/pipeline/task/heimdall/heimdall.go | 9 + internal/pkg/pipeline/task/http/http.go | 9 + internal/pkg/pipeline/task/join/join.go | 9 +- internal/pkg/pipeline/task/jq/jq.go | 3 + internal/pkg/pipeline/task/kafka/kafka.go | 33 +++- internal/pkg/pipeline/task/sample/head.go | 3 + internal/pkg/pipeline/task/sample/nth.go | 3 + internal/pkg/pipeline/task/sample/percent.go | 3 + internal/pkg/pipeline/task/sample/random.go | 23 ++- internal/pkg/pipeline/task/sample/tail.go | 8 + internal/pkg/pipeline/task/sftp/operations.go | 5 + internal/pkg/pipeline/task/sns/sns.go | 5 + internal/pkg/pipeline/task/split/split.go | 2 + internal/pkg/pipeline/task/sqs/sqs.go | 96 ++++++---- test/pipelines/setup_localstack_sqs.sh | 49 +++++ .../sqs_with_context_concurrency.yaml | 2 +- 24 files changed, 493 insertions(+), 41 deletions(-) create mode 100644 internal/pkg/pipeline/ack/ack.go create mode 100755 test/pipelines/setup_localstack_sqs.sh diff --git a/internal/pkg/pipeline/ack/ack.go b/internal/pkg/pipeline/ack/ack.go new file mode 100644 index 0000000..f225d10 --- /dev/null +++ b/internal/pkg/pipeline/ack/ack.go @@ -0,0 +1,177 @@ +// Package ack tracks completion of a single source record (e.g. an SQS +// message) as it flows through a pipeline, so the source task can defer +// acknowledging it (e.g. deleting the SQS receipt) 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 attached to the record's +// context.Context (see WithContext/FromContext), so tasks that just +// transform a record in place need no explicit wiring at all. +// +// Tasks that fan a single input record out into multiple output records +// (e.g. split, or jq with explode) must call AddBranch(n-1) before sending +// the n outputs, so Wait's channel only closes once all n have completed +// (see the Fanout helper below). Tasks that decide not to forward a record +// at all (a filter, an empty query result) must call Done or Fail exactly +// once for it (see Drop). Tasks that fan multiple input records IN into a +// single output record (e.g. join, or archiving many records into one +// file) must attach the Ack returned by Joined to that output record +// instead of any one input's Ack, so completing the joined output +// transitively completes every record that went into it (see Joined). +// Terminal tasks (those with a nil output channel) must call Done, or +// Fail, once they finish processing each record they consume. +type Ack struct { + remaining atomic.Int32 + failed atomic.Bool + done chan struct{} + children []*Ack // completed (Done or Fail, per failed) once this Ack itself 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 this instead of Done when a +// branch didn't actually make it to where it needed to go (e.g. a Kafka +// delivery failure), 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 + } + + // this Ack itself is now fully complete: propagate that to whatever + // Acks it was joined from, based on whether ANY of its own branches + // failed (not just this particular call), since a single downstream + // success/failure on the joined record applies to all of them equally. + 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 as the record moves downstream (record.Record.Context is +// forwarded by tasks even when they construct a new *record.Record). A nil +// ctx (e.g. an aggregating task, like archive pack, that never actually +// received a record to inherit a context from) is treated as +// context.Background() instead of panicking. +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 incoming branch ctx currently carries. +// Call it once, before sending any of the n outputs: this guarantees the +// adjustment is visible before any of those outputs can reach a downstream +// Done/Fail call, which would otherwise be able to race ahead of it. n may +// be 0 (the incoming record produces no output and is immediately +// completed) or repeat a record multiple times (n counts sends, not +// distinct records). +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() + } +} + +// Joined returns a new Ack with a single pending branch representing one +// output record produced by combining n inputs (a "fan-in"), such as join +// or archiving several records into one file. Attach it to that output +// record via WithContext before sending it. Completing the returned Ack +// (Done or Fail, however the output record's own journey downstream ends) +// transitively completes every Ack found among ctxs, so a downstream +// success or failure on the combined record is correctly 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 + +} diff --git a/internal/pkg/pipeline/pipeline.go b/internal/pkg/pipeline/pipeline.go index 8540082..33ee2cf 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" @@ -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 { + // this is a structural fan-out: the same 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 diff --git a/internal/pkg/pipeline/task/archive/tar.go b/internal/pkg/pipeline/task/archive/tar.go index e2d1198..4cc666c 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,11 +29,33 @@ func (t *tarArchive) Read() { } if len(rc.Data) == 0 { + ack.Drop(rc.Context) continue } b := rc.Data + // this is a fan-out: one archive record can expand into multiple + // file records, so the ack must represent all of them - counted up + // front, before any of them is sent - or a downstream Done/Fail for + // the first file could race ahead of a later count adjustment. tar + // readers are forward-only, so counting takes its own pass over a + // fresh reader. + regularFiles := 0 + for counter := tar.NewReader(bytes.NewReader(b)); ; { + header, err := counter.Next() + if err == io.EOF { + break + } + if err != nil { + log.Fatal(err) + } + if header.Typeflag == tar.TypeReg { + regularFiles++ + } + } + ack.Fanout(rc.Context, regularFiles) + r := tar.NewReader(bytes.NewReader(b)) for { @@ -62,12 +86,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 +132,10 @@ func (t *tarArchive) Write() { log.Fatal(err) } - t.SendData(rc.Context, buf.Bytes(), t.OutputChan) + // this is a fan-in: one archive record is produced from every input + // record consumed above, so its ack must transitively complete all of + // theirs instead of discarding all but the last. + 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..eecf249 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.Drop(rc.Context) continue } @@ -35,6 +38,19 @@ func (z *zipArchive) Read() { if err != nil { log.Fatal(err) } + + // this is a fan-out: one archive record can expand into multiple + // file records, so the ack must represent all of them - computed + // up front, before any of them is sent - or a downstream Done/Fail + // for the first file could race ahead of a later count adjustment. + regularFiles := 0 + for _, f := range r.File { + if f.FileInfo().Mode().IsRegular() { + regularFiles++ + } + } + ack.Fanout(rc.Context, regularFiles) + for _, f := range r.File { // check the file type is regular file @@ -66,6 +82,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 +111,19 @@ func (z *zipArchive) Write() { } rc.Context = rec.Context + ctxs = append(ctxs, rec.Context) } if err := zipWriter.Close(); err != nil { log.Fatal(err) } + // this is a fan-in: one archive record is produced from every input + // record consumed above, so its ack must transitively complete all of + // theirs instead of discarding all but the last. + 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 365998b..c98c9f9 100644 --- a/internal/pkg/pipeline/task/aws/parameter_store/parameter_store.go +++ b/internal/pkg/pipeline/task/aws/parameter_store/parameter_store.go @@ -10,6 +10,7 @@ import ( "github.com/aws/aws-sdk-go-v2/service/ssm/types" "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" ) @@ -91,6 +92,14 @@ func (p *parameterStore) Run(input <-chan *record.Record, output chan<- *record. p.SendRecord(r, output) } } + + // terminal (sink) mode: nothing forwards r downstream, so this task + // is the last one to touch it, once all its parameters are set. + if output == nil { + if a, ok := ack.FromContext(r.Context); ok { + a.Done() + } + } } return nil diff --git a/internal/pkg/pipeline/task/compress/compress.go b/internal/pkg/pipeline/task/compress/compress.go index be7e20a..352052c 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.Drop(r.Context) continue } @@ -78,11 +80,16 @@ func (c *core) Run(input <-chan *record.Record, output chan<- *record.Record) (e // skip empty transformed data if len(transformedData) == 0 { + ack.Drop(r.Context) continue } if output != nil { c.SendData(r.Context, transformedData, output) + } else if a, ok := ack.FromContext(r.Context); ok { + // terminal (sink) mode: nothing forwards this record + // downstream, so this task is the last one to touch it. + a.Done() } } diff --git a/internal/pkg/pipeline/task/echo/echo.go b/internal/pkg/pipeline/task/echo/echo.go index 2370ffe..524e224 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" ) @@ -42,6 +43,8 @@ func (e *echo) Run(input <-chan *record.Record, output chan<- *record.Record) (e if output != nil { e.SendRecord(r, output) + } else if a, ok := ack.FromContext(r.Context); ok { + a.Done() } } diff --git a/internal/pkg/pipeline/task/file/file.go b/internal/pkg/pipeline/task/file/file.go index 5b92c7c..007b1e4 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" @@ -203,6 +204,10 @@ func (f *file) writeFile(input <-chan *record.Record) error { if err := writerFunction(&fs, rc, bytes.NewReader(rc.Data)); err != nil { return err } + + if a, ok := ack.FromContext(rc.Context); ok { + a.Done() + } } return nil diff --git a/internal/pkg/pipeline/task/heimdall/heimdall.go b/internal/pkg/pipeline/task/heimdall/heimdall.go index 02d13c5..27e85f2 100644 --- a/internal/pkg/pipeline/task/heimdall/heimdall.go +++ b/internal/pkg/pipeline/task/heimdall/heimdall.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" ) @@ -76,6 +77,14 @@ func (h *heimdall) Run(input <-chan *record.Record, output chan<- *record.Record if err := h.submitJob(jobReq, output); err != nil { return err } + + // terminal (sink) mode: nothing forwards rc downstream, so this + // task is the last one to touch it. + if output == nil { + if a, ok := ack.FromContext(rc.Context); ok { + a.Done() + } + } } return nil } diff --git a/internal/pkg/pipeline/task/http/http.go b/internal/pkg/pipeline/task/http/http.go index 41449c8..f385626 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" @@ -163,6 +164,14 @@ func (h *httpCore) Run(input <-chan *record.Record, output chan<- *record.Record if err := newHttp.processItem(rc, output); err != nil { return err } + + // terminal (sink) mode: nothing forwards rc downstream, so this + // task is the last one to touch it. + if output == nil { + if a, ok := ack.FromContext(rc.Context); ok { + a.Done() + } + } } } diff --git a/internal/pkg/pipeline/task/join/join.go b/internal/pkg/pipeline/task/join/join.go index 00d06f5..6de0bf2 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" ) @@ -95,13 +96,19 @@ 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) + // this is a fan-in: one output record is produced from len(buffer) + // inputs, so its ack must transitively complete every one of theirs + // instead of discarding them. + 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 bb3c9d1..675b9ab 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" ) @@ -41,9 +42,11 @@ func (j *jq) Run(input <-chan *record.Record, output chan<- *record.Record) (err return err } if items == nil { + ack.Drop(r.Context) continue } if splitItems, ok := items.([]any); j.Explode && ok { + ack.Fanout(r.Context, len(splitItems)) for _, splitItem := range splitItems { if j.AsRaw { j.SendData(r.Context, fmt.Appendf(nil, "%v", splitItem), output) diff --git a/internal/pkg/pipeline/task/kafka/kafka.go b/internal/pkg/pipeline/task/kafka/kafka.go index a5418d6..01f8f0b 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" ) @@ -148,10 +149,28 @@ func (k *kafka) write(input <-chan *record.Record) error { 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 + } + 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 failed to make it to Kafka: fail its ack + // instead of dropping it silently, so the source knows not + // to acknowledge/delete it (letting it be retried upstream). + if a, ok := m.Opaque.(*ack.Ack); ok { + a.Fail() + } + continue + } + // only mark the source record complete once the broker has confirmed + // delivery, not merely once it's been enqueued locally. + if a, ok := m.Opaque.(*ack.Ack); ok { + a.Done() } } }() @@ -169,9 +188,15 @@ func (k *kafka) write(input <-chan *record.Record) error { break } + var opaque any + if a, ok := ack.FromContext(r.Context); ok { + opaque = a + } + 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) break diff --git a/internal/pkg/pipeline/task/sample/head.go b/internal/pkg/pipeline/task/sample/head.go index 79147ab..2881097 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" ) @@ -24,6 +25,8 @@ func (h *head) filter(r *record.Record, output chan<- *record.Record) error { if h.index < h.limit { h.sendRecord(r, output) h.index++ + } else { + ack.Drop(r.Context) } return nil diff --git a/internal/pkg/pipeline/task/sample/nth.go b/internal/pkg/pipeline/task/sample/nth.go index d05618d..3b06ce3 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" ) @@ -23,6 +24,8 @@ func (n *nth) filter(r *record.Record, output chan<- *record.Record) error { if n.index%n.divider == 0 { n.sendRecord(r, output) + } else { + ack.Drop(r.Context) } n.index++ diff --git a/internal/pkg/pipeline/task/sample/percent.go b/internal/pkg/pipeline/task/sample/percent.go index 3847a56..9c3b857 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" ) @@ -42,6 +43,8 @@ func (p *percent) filter(r *record.Record, output chan<- *record.Record) error { if n.Int64() < int64(p.cutoff) { p.sendRecord(r, output) + } else { + ack.Drop(r.Context) } return nil diff --git a/internal/pkg/pipeline/task/sample/random.go b/internal/pkg/pipeline/task/sample/random.go index dfb07cd..cb40c6e 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.Drop(row.Context) } return nil @@ -38,6 +41,13 @@ 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 the same buffered record can be + // sent zero, one, or multiple times. Tally every draw first, then + // adjust each record's ack for its final send count before sending + // any of them - otherwise a downstream Done/Fail for an earlier + // send could race ahead of a later AddBranch call for the same ack. + counts := make([]int, l) for i := 0; i < r.limit; i++ { index, err := rand.Int(rand.Reader, big.NewInt(l)) @@ -45,9 +55,20 @@ func (r *random) drain(output chan<- *record.Record) error { return err } - r.sendRecord(r.buffer[index.Int64()], output) + counts[index.Int64()]++ + + } + for i, count := range counts { + ack.Fanout(r.buffer[i].Context, count) } + + for i, count := range counts { + for range count { + r.sendRecord(r.buffer[i], output) + } + } + } return nil diff --git a/internal/pkg/pipeline/task/sample/tail.go b/internal/pkg/pipeline/task/sample/tail.go index e394d6b..c46dc64 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,13 @@ func newTail(s *sample) (sampler, error) { func (t *tail) filter(r *record.Record, _ chan<- *record.Record) error { + // the ring buffer is about to overwrite whatever record currently + // occupies this slot (if any); it will never be forwarded, so its ack + // must be completed here instead of leaking forever. + if evicted := t.buffer[t.index]; evicted != nil { + ack.Drop(evicted.Context) + } + t.buffer[t.index] = r t.index = (t.index + 1) % t.limit t.count++ diff --git a/internal/pkg/pipeline/task/sftp/operations.go b/internal/pkg/pipeline/task/sftp/operations.go index 5eecfd9..110cf90 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" @@ -35,6 +36,10 @@ func (s *sftp) upload(client *pkgsftp.Client, input <-chan *record.Record) error if err := s.uploadOne(client, file, rc.Data); err != nil { return err } + + if a, ok := ack.FromContext(rc.Context); ok { + a.Done() + } } return nil diff --git a/internal/pkg/pipeline/task/sns/sns.go b/internal/pkg/pipeline/task/sns/sns.go index cec5c03..a771074 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" ) @@ -105,6 +106,10 @@ func (s *snsTask) Run(input <-chan *record.Record, output chan<- *record.Record) if err != nil { return fmt.Errorf("failed to publish to SNS topic %s: %w", s.TopicArn, err) } + + if a, ok := ack.FromContext(r.Context); ok { + a.Done() + } } return nil diff --git a/internal/pkg/pipeline/task/split/split.go b/internal/pkg/pipeline/task/split/split.go index 51029a7..4925a10 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" ) @@ -29,6 +30,7 @@ func (s *split) Run(input <-chan *record.Record, output chan<- *record.Record) e break } lines := strings.Split(strings.TrimSuffix(string(r.Data), s.Delimiter), s.Delimiter) + ack.Fanout(r.Context, len(lines)) for _, line := range lines { s.SendData(r.Context, []byte(line), output) } diff --git a/internal/pkg/pipeline/task/sqs/sqs.go b/internal/pkg/pipeline/task/sqs/sqs.go index 1dcf05b..16a6c9d 100644 --- a/internal/pkg/pipeline/task/sqs/sqs.go +++ b/internal/pkg/pipeline/task/sqs/sqs.go @@ -14,16 +14,23 @@ 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" + + // inFlight bounds how many messages can be unacknowledged (received but + // not yet deleted) at once. It's a multiple of Concurrency rather than + // Concurrency itself so fetching can run some distance ahead of full + // downstream completion instead of stalling every time Concurrency + // messages are simultaneously in flight. + inFlightMultiplier = 5 ) var ( @@ -93,28 +100,24 @@ func (s *sqs) Run(input <-chan *record.Record, output chan<- *record.Record) err 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) - - // we set a pool of workers that will delete messages from the queue + // If input is nil, act as a source: read messages and, once every + // downstream task has finished with a given message, delete its + // receipt so it isn't redelivered. + inFlight := make(chan struct{}, s.Concurrency*inFlightMultiplier) 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) + err := s.getMessages(ctx, output, inFlight, &wg) + // wait for every deleteOnComplete goroutine spawned below to finish + // (and its receipt to be deleted or left alone) before this task's Run + // returns, so a shutdown never abandons in-flight acknowledgements. wg.Wait() return err } -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, inFlight chan struct{}, wg *sync.WaitGroup) 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, same as when + // there's no consumer at all. + if output == nil { + if _, err := s.client.DeleteMessage(ctx, &qs.DeleteMessageInput{ + QueueUrl: &s.QueueURL, + ReceiptHandle: m.ReceiptHandle, + }); err != nil { + fmt.Printf("failed to delete message %s from queue %s: %v\n", aws.ToString(m.MessageId), s.QueueURL, err) + } + continue } - // send receipt to receipts channel for deletion - receipts <- m.ReceiptHandle + inFlight <- struct{}{} + + msgAck := ack.New() + s.SendData(ack.WithContext(ctx, msgAck), []byte(*m.Body), output) + + wg.Add(1) + go s.deleteOnComplete(msgAck, m.MessageId, m.ReceiptHandle, inFlight, wg) } } } } -func (s *sqs) processReceipts(receipts <-chan *string, wg *sync.WaitGroup) error { +// deleteOnComplete waits until every downstream task has finished +// processing the record derived from this message, deletes its receipt so +// it isn't redelivered, and frees its inFlight slot. It runs detached from +// getMessages, since getMessages must return (closing the task's output +// channel) before downstream tasks can drain and signal completion; wg lets +// Run wait for it to finish before returning. +func (s *sqs) deleteOnComplete(msgAck *ack.Ack, messageId, receiptHandle *string, inFlight chan struct{}, wg *sync.WaitGroup) { defer wg.Done() + defer func() { <-inFlight }() - for receipt := range receipts { - if _, err := s.client.DeleteMessage(ctx, &qs.DeleteMessageInput{ - QueueUrl: &s.QueueURL, - ReceiptHandle: receipt, - }); err != nil { - return err - } + <-msgAck.Wait() + + // a downstream failure means this message wasn't fully processed: + // leave its receipt alone so SQS redelivers it after the visibility + // timeout instead of losing it. + if msgAck.Failed() { + return } - return nil + 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) + } } @@ -205,6 +233,10 @@ func (s *sqs) sendMessages(input <-chan *record.Record) error { if err != nil { return err } + + if a, ok := ack.FromContext(r.Context); ok { + a.Done() + } } return nil } 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_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 From be08956a8662d2aed1077f16b9aa95f7172c22e4 Mon Sep 17 00:00:00 2001 From: Mahesh Date: Tue, 4 Aug 2026 15:25:09 +0530 Subject: [PATCH 2/3] fix(ack): remove deadlocks and ack leaks from deferred acknowledgment Addresses review feedback on the acknowledgment tracking added in 4fb7139. Auditing the goroutine-per-message design surfaced two deadlocks and a class of ack leaks, all reachable from pipeline shapes that run in production today. Deadlocks - The source gated its receive loop on a semaphore of Concurrency*5 that was released only on downstream completion. Any fan-in that accumulates records before emitting closed the cycle: the source stalls at the limit, the fan-in never reaches its flush threshold, nothing acks, no slot frees. The bound is gone; messages in flight are bounded by channel_size, which is what applies backpressure anyway. - tracker.Wait() ran inside Run, but a task's output channel closes only after Run returns. join, archive pack and sample tail/random emit only when their input closes, so waiting inside Run waited for a flush that could not happen until the wait ended. The wait moved to a new optional task.Finisher hook the pipeline calls after closing the output channel. Ack leaks - Base.SendRecord dropped the ack when output was nil, so any task terminating a pipeline without its own guard never settled it. Settled centrally now. jq and replace additionally skipped their whole loop when terminal, never draining their input; they now drain and drop. - 14 sites across 13 tasks abandoned the record they were holding on a mid-loop error, stranding its ack. New ack.Reject/ack.Rejected settle it as failed so the broker redelivers instead of the pipeline hanging. One malformed CSV row was enough to hang a pipeline indefinitely. - The pipeline rejects anything left in a task's input once every worker has returned. Doing this inside the task instead discarded records that healthy sibling workers would have written - measured at 997 of 1998 lost on a file-source pipeline with task_concurrency 4. Data loss - converter and xpath emitted one record per csv row / container node without adjusting the ack, so the first output to complete settled the whole record and the message was deleted while the rest were still in flight. Both now count before sending, matching split and archive. Also - join's duration: never fired. The select had a default case, so it never blocked and the ticker was only serviced between records - never while input was stalled, which is exactly when a partial batch needs flushing. - archive tar unpack reads the archive once instead of twice. - concurrency now sizes the pool of concurrent deletions, restoring the meaning it has on main. No pipeline YAML changes. Verified against origin/main: 20 fixtures with matching exit codes (15 byte-identical), 6 converter goldens byte-identical, race detector clean, 10k records at 38-118 MB, exactly-once with 0 duplicates. Co-Authored-By: Claude Opus 5 (1M context) --- internal/pkg/pipeline/ack/ack.go | 31 ++++++ internal/pkg/pipeline/ack/tracker.go | 89 +++++++++++++++ internal/pkg/pipeline/pipeline.go | 35 +++++- internal/pkg/pipeline/task/archive/tar.go | 51 +++++---- .../aws/parameter_store/parameter_store.go | 6 +- .../pkg/pipeline/task/compress/compress.go | 4 +- .../pkg/pipeline/task/converter/converter.go | 20 +++- internal/pkg/pipeline/task/file/file.go | 24 ++++- internal/pkg/pipeline/task/flatten/flatten.go | 5 +- .../pkg/pipeline/task/heimdall/heimdall.go | 4 +- internal/pkg/pipeline/task/http/http.go | 4 +- internal/pkg/pipeline/task/join/join.go | 10 +- internal/pkg/pipeline/task/jq/jq.go | 94 ++++++++++------ internal/pkg/pipeline/task/replace/replace.go | 22 +++- internal/pkg/pipeline/task/sample/sample.go | 3 +- internal/pkg/pipeline/task/sftp/operations.go | 4 +- internal/pkg/pipeline/task/sns/sns.go | 2 +- internal/pkg/pipeline/task/sqs/README.md | 27 ++++- internal/pkg/pipeline/task/sqs/sqs.go | 101 +++++++++--------- internal/pkg/pipeline/task/task.go | 19 ++++ internal/pkg/pipeline/task/xpath/xpath.go | 33 ++++-- test/pipelines/setup_localstack_fanout.sh | 54 ++++++++++ test/pipelines/sqs_fanout_fanin_dag.yaml | 66 ++++++++++++ 23 files changed, 565 insertions(+), 143 deletions(-) create mode 100644 internal/pkg/pipeline/ack/tracker.go create mode 100755 test/pipelines/setup_localstack_fanout.sh create mode 100644 test/pipelines/sqs_fanout_fanin_dag.yaml diff --git a/internal/pkg/pipeline/ack/ack.go b/internal/pkg/pipeline/ack/ack.go index f225d10..051bcba 100644 --- a/internal/pkg/pipeline/ack/ack.go +++ b/internal/pkg/pipeline/ack/ack.go @@ -154,6 +154,37 @@ func Drop(ctx context.Context) { } } +// Reject completes the Ack embedded in ctx, if any, as FAILED, for a record a +// task could not process. It is the counterpart of Drop: Drop means "this +// record is legitimately finished with", Reject means "this record never made +// it", so the source leaves it unacknowledged and the broker redelivers it +// instead of the pipeline waiting forever for a completion that can't come. +// +// A task bailing out mid-stream should Reject both the record it failed on and +// every record still queued behind it, since it won't be processing those +// either. +func Reject(ctx context.Context) { + if a, ok := FromContext(ctx); ok { + a.Fail() + } +} + +// Rejected is Reject followed by err, for the common case of a task bailing +// out on the record it is holding: +// +// if err != nil { +// return ack.Rejected(r.Context, err) +// } +// +// Keeping the settle and the return on one line stops the two drifting apart - +// a bare `return err` here strands the record, and the symptom is the whole +// pipeline hanging at shutdown rather than anything that points back to this +// line. +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 produced by combining n inputs (a "fan-in"), such as join // or archiving several records into one file. Attach it to that output diff --git a/internal/pkg/pipeline/ack/tracker.go b/internal/pkg/pipeline/ack/tracker.go new file mode 100644 index 0000000..294a861 --- /dev/null +++ b/internal/pkg/pipeline/ack/tracker.go @@ -0,0 +1,89 @@ +package ack + +import "sync" + +// Acknowledger is the broker-specific half of deferred acknowledgement: how +// one particular source settles a single message once the pipeline is done +// with it. SQS deletes the message's receipt, Kafka stores its offset, +// another broker does something else again - Tracker owns the bookkeeping +// that's common to all of them, an Acknowledger owns what isn't. +// +// Implementations are typically a small per-message struct holding whatever +// handle the client needs (a receipt, a partition and offset, ...). +type Acknowledger interface { + // Ack settles one message. failed reports whether any downstream branch + // signalled Fail rather than Done, in which case the message wasn't + // fully processed: implementations should normally leave it + // unacknowledged so the broker redelivers it instead of losing it. + // + // Ack is called at most once per message, from its own goroutine, and no + // more than the Tracker's concurrency at a time. It should log rather + // than panic on client errors, since by the time it runs the record is + // already through the pipeline and the worst case is a redelivery. + Ack(failed bool) +} + +// Tracker lets a source task defer acknowledging a message (deleting an SQS +// receipt, committing a Kafka offset) 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 limit on unacknowledged +// messages sounds prudent, but releasing it depends on downstream completion, +// which deadlocks against any fan-in task that must accumulate records before +// it can emit anything: the source stops receiving at the limit, the fan-in +// never reaches its flush threshold, so nothing ever completes and no slot +// ever frees. The number of messages in flight is instead bounded by the +// pipeline's channel capacity, which is what 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() + + // bound how many broker calls run at once. Taking the slot here + // rather than before the wait is what keeps this safe: by now the + // record is through the pipeline, so releasing the slot depends only + // on Ack returning, 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 33ee2cf..34f9714 100644 --- a/internal/pkg/pipeline/pipeline.go +++ b/internal/pkg/pipeline/pipeline.go @@ -263,11 +263,42 @@ 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() + + // every worker has returned. If any of them bailed out early there can + // be records left in this task's input that nobody is going to process + // and, now, nobody left to consume - which would block whatever is + // still writing upstream. Drain them and reject their acks so a source + // deferring acknowledgement redelivers them instead of waiting forever. + // This only ever finds anything when a worker returned an error; on the + // normal path the workers have already drained the channel. + 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: this is 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/archive/tar.go b/internal/pkg/pipeline/task/archive/tar.go index 4cc666c..531e3b8 100644 --- a/internal/pkg/pipeline/task/archive/tar.go +++ b/internal/pkg/pipeline/task/archive/tar.go @@ -20,6 +20,13 @@ type tarArchive struct { *channelStruct } +// tarFile is a regular file extracted from an archive, held until the total +// file count is known and the fan-out ack can be sized. +type tarFile struct { + name string + data []byte +} + func (t *tarArchive) Read() { for { @@ -39,22 +46,9 @@ func (t *tarArchive) Read() { // file records, so the ack must represent all of them - counted up // front, before any of them is sent - or a downstream Done/Fail for // the first file could race ahead of a later count adjustment. tar - // readers are forward-only, so counting takes its own pass over a - // fresh reader. - regularFiles := 0 - for counter := tar.NewReader(bytes.NewReader(b)); ; { - header, err := counter.Next() - if err == io.EOF { - break - } - if err != nil { - log.Fatal(err) - } - if header.Typeflag == tar.TypeReg { - regularFiles++ - } - } - ack.Fanout(rc.Context, regularFiles) + // readers are forward-only, so the single pass below extracts every + // regular file first and only sends them once the count is known. + files := make([]tarFile, 0) r := tar.NewReader(bytes.NewReader(b)) @@ -68,15 +62,26 @@ 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) } + files = append(files, tarFile{ + name: textutil.SlugifyFileName(filepath.Base(header.Name)), + data: buf, + }) + } + + ack.Fanout(rc.Context, len(files)) + + for _, f := range files { + rc.SetContextValue(string(task.CtxKeyArchiveFileNameWrite), f.name) + t.SendData(rc.Context, f.data, t.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 c98c9f9..0bff05e 100644 --- a/internal/pkg/pipeline/task/aws/parameter_store/parameter_store.go +++ b/internal/pkg/pipeline/task/aws/parameter_store/parameter_store.go @@ -66,12 +66,12 @@ func (p *parameterStore) Run(input <-chan *record.Record, output chan<- *record. for parameterName, parameterQuery := range p.SetParameters { parameterValue, err := parameterQuery.Execute(r.Data) if err != nil { - return err + return ack.Rejected(r.Context, err) } parameterValueString, isString := parameterValue.(string) if !isString { - return fmt.Errorf("%s parameter value is not string", parameterName) + return ack.Rejected(r.Context, fmt.Errorf("%s parameter value is not string", parameterName)) } putParameterInput := &ssm.PutParameterInput{ @@ -85,7 +85,7 @@ func (p *parameterStore) Run(input <-chan *record.Record, output chan<- *record. } if _, err := p.client.PutParameter(ctx, putParameterInput); err != nil { - return err + return ack.Rejected(r.Context, err) } if output != nil { diff --git a/internal/pkg/pipeline/task/compress/compress.go b/internal/pkg/pipeline/task/compress/compress.go index 352052c..83dcf22 100644 --- a/internal/pkg/pipeline/task/compress/compress.go +++ b/internal/pkg/pipeline/task/compress/compress.go @@ -70,11 +70,11 @@ 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) } } diff --git a/internal/pkg/pipeline/task/converter/converter.go b/internal/pkg/pipeline/task/converter/converter.go index 40827b1..8f59886 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,9 +80,26 @@ 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 + // this record is being abandoned, so settle its ack as failed: + // leaving it pending would strand a source that defers + // acknowledgement, which waits for a completion that can no longer + // come. Rejecting sends the message back for redelivery instead. + return ack.Rejected(r.Context, err) } + // this is a fan-out: one input record can convert into many outputs - + // a csv row each, an xlsx sheet each - so the ack must represent all + // of them, counted before any is sent, or a downstream Done for the + // first could settle the whole record while the rest are still in + // flight. A count of 0 (nothing converted) completes it right here. + sends := 0 + for _, out := range outputs { + if out.Data != nil { + sends++ + } + } + ack.Fanout(r.Context, sends) + for _, out := range outputs { if out.Data != nil { // Add metadata to context diff --git a/internal/pkg/pipeline/task/file/file.go b/internal/pkg/pipeline/task/file/file.go index 007b1e4..aeb0f00 100644 --- a/internal/pkg/pipeline/task/file/file.go +++ b/internal/pkg/pipeline/task/file/file.go @@ -159,6 +159,22 @@ func (f *file) readFile(output chan<- *record.Record) error { } +// abort settles rc as failed, then returns err. Bailing out without doing this +// leaves a source that defers acknowledgement waiting forever for a record +// this task is never going to write; rejecting it lets the broker redeliver +// it instead. +// +// Only rc is rejected, never the records still queued behind it: other workers +// of this task are still running and will write those. 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 { @@ -170,13 +186,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 == `` { @@ -199,10 +215,10 @@ 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) } if a, ok := ack.FromContext(rc.Context); ok { diff --git a/internal/pkg/pipeline/task/flatten/flatten.go b/internal/pkg/pipeline/task/flatten/flatten.go index c36ab46..e431330 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,7 +39,7 @@ 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) diff --git a/internal/pkg/pipeline/task/heimdall/heimdall.go b/internal/pkg/pipeline/task/heimdall/heimdall.go index 27e85f2..737f5c0 100644 --- a/internal/pkg/pipeline/task/heimdall/heimdall.go +++ b/internal/pkg/pipeline/task/heimdall/heimdall.go @@ -69,13 +69,13 @@ 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 { - return err + return ack.Rejected(rc.Context, err) } // terminal (sink) mode: nothing forwards rc downstream, so this diff --git a/internal/pkg/pipeline/task/http/http.go b/internal/pkg/pipeline/task/http/http.go index f385626..ea273fa 100644 --- a/internal/pkg/pipeline/task/http/http.go +++ b/internal/pkg/pipeline/task/http/http.go @@ -159,10 +159,10 @@ 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) } // terminal (sink) mode: nothing forwards rc downstream, so this diff --git a/internal/pkg/pipeline/task/join/join.go b/internal/pkg/pipeline/task/join/join.go index 6de0bf2..526ee75 100644 --- a/internal/pkg/pipeline/task/join/join.go +++ b/internal/pkg/pipeline/task/join/join.go @@ -56,11 +56,15 @@ 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 control would fall straight through to a bare + // channel receive and tickerCh would only be serviced between records - + // never firing while input is stalled, which is exactly when a partially + // filled buffer needs flushing. With duration unset tickerCh is nil, so + // this degenerates to today's 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) diff --git a/internal/pkg/pipeline/task/jq/jq.go b/internal/pkg/pipeline/task/jq/jq.go index 675b9ab..7afafc1 100644 --- a/internal/pkg/pipeline/task/jq/jq.go +++ b/internal/pkg/pipeline/task/jq/jq.go @@ -23,51 +23,75 @@ 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: a jq transform with nowhere to send has no effect, but the + // 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.Drop(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 + } + + // 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 { - return err + // Execute the JQ query + items, err := query.Execute(r.Data) + if err != nil { + return ack.Rejected(r.Context, err) + } + if items == nil { + ack.Drop(r.Context) + continue + } + if splitItems, ok := items.([]any); j.Explode && ok { + // marshal every item before adjusting the ack: failing partway + // through afterwards would leave branches counted but never sent, + // and there's no clean way to unwind a partial fan-out. + payloads := make([][]byte, 0, len(splitItems)) + for _, splitItem := range splitItems { + if j.AsRaw { + payloads = append(payloads, fmt.Appendf(nil, "%v", splitItem)) + continue + } + jsonItem, err := json.Marshal(splitItem) + if err != nil { + return ack.Rejected(r.Context, err) + } + payloads = append(payloads, jsonItem) } - if items == nil { - ack.Drop(r.Context) - continue + + ack.Fanout(r.Context, len(payloads)) + + for _, payload := range payloads { + j.SendData(r.Context, payload, output) } - if splitItems, ok := items.([]any); j.Explode && ok { - ack.Fanout(r.Context, len(splitItems)) - 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) - } - } + } 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) } } } diff --git a/internal/pkg/pipeline/task/replace/replace.go b/internal/pkg/pipeline/task/replace/replace.go index e7d724a..e65fc31 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,29 @@ 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: a replace with nowhere to send has no effect, but the + // 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.Drop(record.Context) + } + } + + for { + record, ok := r.GetRecord(input) + if !ok { + break } + r.SendData(record.Context, []byte(rx.ReplaceAllString(string(record.Data), r.Replacement)), output) } 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/sftp/operations.go b/internal/pkg/pipeline/task/sftp/operations.go index 110cf90..21c07b1 100644 --- a/internal/pkg/pipeline/task/sftp/operations.go +++ b/internal/pkg/pipeline/task/sftp/operations.go @@ -30,11 +30,11 @@ 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) } if a, ok := ack.FromContext(rc.Context); ok { diff --git a/internal/pkg/pipeline/task/sns/sns.go b/internal/pkg/pipeline/task/sns/sns.go index a771074..e127757 100644 --- a/internal/pkg/pipeline/task/sns/sns.go +++ b/internal/pkg/pipeline/task/sns/sns.go @@ -104,7 +104,7 @@ 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)) } if a, ok := ack.FromContext(r.Context); ok { diff --git a/internal/pkg/pipeline/task/sqs/README.md b/internal/pkg/pipeline/task/sqs/README.md index 56d2e9a..a530103 100644 --- a/internal/pkg/pipeline/task/sqs/README.md +++ b/internal/pkg/pipeline/task/sqs/README.md @@ -18,9 +18,9 @@ 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` | int | `10` | Long polling wait time in seconds | +| `wait_time_seconds` | int | `10` | Long polling wait time in seconds | | `exit_on_empty` | bool | `false` | Exit when queue is empty | | `message_group_id` | string | - | Message group ID for FIFO queues | | `fail_on_error` | bool | `false` | Whether to stop the pipeline if this task encounters an error | @@ -64,9 +64,30 @@ 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 be unacknowledged 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. +- **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_reader.yaml` - SQS message reading example +- `test/pipelines/sqs_with_context_concurrency.yaml` - SQS reading with per-task + 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 16a6c9d..c16b58c 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" @@ -24,13 +23,6 @@ const ( defaultMaxMessages = 10 defaultWaitTimeSeconds = 10 defaultRegion = "us-west-2" - - // inFlight bounds how many messages can be unacknowledged (received but - // not yet deleted) at once. It's a multiple of Concurrency rather than - // Concurrency itself so fetching can run some distance ahead of full - // downstream completion instead of stalling every time Concurrency - // messages are simultaneously in flight. - inFlightMultiplier = 5 ) var ( @@ -47,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) { @@ -74,6 +67,8 @@ func (s *sqs) Init() error { } s.client = qs.NewFromConfig(awsConfig) + s.tracker = ack.NewTracker(s.Concurrency) + return nil } @@ -95,29 +90,36 @@ 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: read messages and, once every - // downstream task has finished with a given message, delete its - // receipt so it isn't redelivered. - inFlight := make(chan struct{}, s.Concurrency*inFlightMultiplier) - var wg sync.WaitGroup + // If input is nil, act as a source: read messages and hand each one's + // 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; see Finish. + return s.getMessages(ctx, output) - err := s.getMessages(ctx, output, inFlight, &wg) +} - // wait for every deleteOnComplete goroutine spawned below to finish - // (and its receipt to be deleted or left alone) before this task's Run - // returns, so a shutdown never abandons in-flight acknowledgements. - wg.Wait() +// Finish waits for every deferred deletion to run (and each receipt to be +// deleted or left alone) 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 only emit 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 { - return err + if s.tracker != nil { + s.tracker.Wait() + } + + return nil } -func (s *sqs) getMessages(ctx context.Context, output chan<- *record.Record, inFlight chan struct{}, wg *sync.WaitGroup) error { +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 { @@ -164,48 +166,51 @@ func (s *sqs) getMessages(ctx context.Context, output chan<- *record.Record, inF // wait for: delete the receipt right away, same as when // there's no consumer at all. if output == nil { - if _, err := s.client.DeleteMessage(ctx, &qs.DeleteMessageInput{ - QueueUrl: &s.QueueURL, - ReceiptHandle: m.ReceiptHandle, - }); err != nil { - fmt.Printf("failed to delete message %s from queue %s: %v\n", aws.ToString(m.MessageId), s.QueueURL, err) - } + s.deleteMessage(m.MessageId, m.ReceiptHandle) continue } - inFlight <- struct{}{} - msgAck := ack.New() s.SendData(ack.WithContext(ctx, msgAck), []byte(*m.Body), output) - wg.Add(1) - go s.deleteOnComplete(msgAck, m.MessageId, m.ReceiptHandle, inFlight, wg) + s.tracker.Track(msgAck, &messageAck{ + sqs: s, + messageId: m.MessageId, + receiptHandle: m.ReceiptHandle, + }) } } } } -// deleteOnComplete waits until every downstream task has finished -// processing the record derived from this message, deletes its receipt so -// it isn't redelivered, and frees its inFlight slot. It runs detached from -// getMessages, since getMessages must return (closing the task's output -// channel) before downstream tasks can drain and signal completion; wg lets -// Run wait for it to finish before returning. -func (s *sqs) deleteOnComplete(msgAck *ack.Ack, messageId, receiptHandle *string, inFlight chan struct{}, wg *sync.WaitGroup) { - - defer wg.Done() - defer func() { <-inFlight }() +// messageAck acknowledges one received message on behalf of ack.Tracker. +type messageAck struct { + sqs *sqs + messageId *string + receiptHandle *string +} - <-msgAck.Wait() +// Ack deletes the message's receipt so SQS doesn't redeliver it. On a +// downstream failure it does nothing: the message wasn't fully processed, so +// leaving the receipt alone lets SQS redeliver it once the visibility +// timeout expires instead of losing it. +func (m *messageAck) Ack(failed bool) { - // a downstream failure means this message wasn't fully processed: - // leave its receipt alone so SQS redelivers it after the visibility - // timeout instead of losing it. - if msgAck.Failed() { + if failed { return } + m.sqs.deleteMessage(m.messageId, m.receiptHandle) + +} + +// deleteMessage acknowledges a message by deleting its receipt so it isn't +// redelivered. A failure to delete is logged rather than returned: the +// message has already been processed, and 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, @@ -231,7 +236,7 @@ func (s *sqs) sendMessages(input <-chan *record.Record) error { MessageGroupId: s.getMessageGroupID(), }) if err != nil { - return err + return ack.Rejected(r.Context, err) } if a, ok := ack.FromContext(r.Context); ok { diff --git a/internal/pkg/pipeline/task/task.go b/internal/pkg/pipeline/task/task.go index 776a067..0c6d530 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,19 @@ 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. A source deferring acknowledgement is the case +// this exists for: its acks can only settle after every downstream task has +// drained, and downstream tasks that emit on input close - join, archive +// pack, sample tail/random - can only drain once the source's output channel +// is closed. Waiting inside Run would therefore 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"` @@ -111,6 +125,11 @@ func (b *Base) SendData(ctx context.Context, data []byte, output chan<- *record. func (b *Base) SendRecord(r *record.Record, output chan<- *record.Record) /* we should return error here */ { if output == nil { + // terminal task: nothing forwards r downstream, so this is the last + // place that will ever touch it. Settle its ack here rather than + // dropping it, or a source deferring acknowledgement waits forever + // for a completion that can no longer come. + ack.Drop(r.Context) return } diff --git a/internal/pkg/pipeline/task/xpath/xpath.go b/internal/pkg/pipeline/task/xpath/xpath.go index 29b6f2f..50dedd2 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,25 +46,45 @@ 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.Drop(r.Context) continue } } + // this is a fan-out: one input record yields one output per container + // node, so the ack must represent all of them, counted before any is + // sent - otherwise a downstream Done for the first node could settle + // the whole record while later nodes are still in flight. queryFields + // is pure, so extracting every node up front costs only memory, and + // the original node position is carried along since node_index is + // part of the output contract. + type nodePayload struct { + index int + data []byte + } + + payloads := make([]nodePayload, 0, len(containerNodes)) + 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 { - index := fmt.Sprintf("%d", i+1) - r.SetContextValue(nodeIndexKey, index) - x.SendData(r.Context, data, output) + payloads = append(payloads, nodePayload{index: i + 1, data: data}) } } + + ack.Fanout(r.Context, len(payloads)) + + for _, p := range payloads { + r.SetContextValue(nodeIndexKey, fmt.Sprintf("%d", p.index)) + x.SendData(r.Context, p.data, output) + } } 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/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 From 738b6444fe54522423f67e8fee6c0db7c89006c5 Mon Sep 17 00:00:00 2001 From: Divyanshu Tiwari Date: Thu, 6 Aug 2026 14:48:44 +0530 Subject: [PATCH 3/3] docs(ack): tighten comments to durable why Trim the acknowledgment comments to the repo's comment conventions: cut essay-length blocks to the constraint a reader can't recover from the code, drop "instead of " framing that goes stale the moment this merges, and stop naming tasks from other packages in ack/ and task.Finisher. Also removes guidance in ack.Reject that contradicted file.abort about who rejects the records queued behind a failed one. --- internal/pkg/pipeline/ack/ack.go | 112 +++++++----------- internal/pkg/pipeline/ack/tracker.go | 49 +++----- internal/pkg/pipeline/pipeline.go | 21 ++-- internal/pkg/pipeline/task/archive/tar.go | 19 ++- internal/pkg/pipeline/task/archive/zip.go | 12 +- .../aws/parameter_store/parameter_store.go | 4 +- .../pkg/pipeline/task/compress/compress.go | 4 +- .../pkg/pipeline/task/converter/converter.go | 12 +- internal/pkg/pipeline/task/file/file.go | 13 +- .../pkg/pipeline/task/heimdall/heimdall.go | 3 +- internal/pkg/pipeline/task/http/http.go | 3 +- internal/pkg/pipeline/task/join/join.go | 14 +-- internal/pkg/pipeline/task/jq/jq.go | 8 +- internal/pkg/pipeline/task/kafka/kafka.go | 8 +- internal/pkg/pipeline/task/replace/replace.go | 6 +- internal/pkg/pipeline/task/sample/random.go | 9 +- internal/pkg/pipeline/task/sample/tail.go | 5 +- internal/pkg/pipeline/task/sqs/sqs.go | 35 +++--- internal/pkg/pipeline/task/task.go | 15 +-- internal/pkg/pipeline/task/xpath/xpath.go | 11 +- 20 files changed, 142 insertions(+), 221 deletions(-) diff --git a/internal/pkg/pipeline/ack/ack.go b/internal/pkg/pipeline/ack/ack.go index 051bcba..8450db0 100644 --- a/internal/pkg/pipeline/ack/ack.go +++ b/internal/pkg/pipeline/ack/ack.go @@ -1,7 +1,6 @@ -// Package ack tracks completion of a single source record (e.g. an SQS -// message) as it flows through a pipeline, so the source task can defer -// acknowledging it (e.g. deleting the SQS receipt) until every downstream -// branch produced from it has finished processing. +// 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 ( @@ -13,28 +12,19 @@ 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 attached to the record's -// context.Context (see WithContext/FromContext), so tasks that just -// transform a record in place need no explicit wiring at all. +// 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 fan a single input record out into multiple output records -// (e.g. split, or jq with explode) must call AddBranch(n-1) before sending -// the n outputs, so Wait's channel only closes once all n have completed -// (see the Fanout helper below). Tasks that decide not to forward a record -// at all (a filter, an empty query result) must call Done or Fail exactly -// once for it (see Drop). Tasks that fan multiple input records IN into a -// single output record (e.g. join, or archiving many records into one -// file) must attach the Ack returned by Joined to that output record -// instead of any one input's Ack, so completing the joined output -// transitively completes every record that went into it (see Joined). -// Terminal tasks (those with a nil output channel) must call Done, or -// Fail, once they finish processing each record they consume. +// 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 // completed (Done or Fail, per failed) once this Ack itself completes; see Joined + children []*Ack // settled with this Ack's own outcome once it completes; see Joined } // New returns an Ack with a single pending branch. @@ -55,10 +45,9 @@ 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 this instead of Done when a -// branch didn't actually make it to where it needed to go (e.g. a Kafka -// delivery failure), so the source knows not to acknowledge the record. +// 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) } @@ -73,10 +62,9 @@ func (a *Ack) complete(failed bool) { return } - // this Ack itself is now fully complete: propagate that to whatever - // Acks it was joined from, based on whether ANY of its own branches - // failed (not just this particular call), since a single downstream - // success/failure on the joined record applies to all of them equally. + // 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 { @@ -103,11 +91,8 @@ func (a *Ack) Wait() <-chan struct{} { } // WithContext returns a copy of ctx carrying a, recoverable later via -// FromContext as the record moves downstream (record.Record.Context is -// forwarded by tasks even when they construct a new *record.Record). A nil -// ctx (e.g. an aggregating task, like archive pack, that never actually -// received a record to inherit a context from) is treated as -// context.Background() instead of panicking. +// 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() @@ -121,14 +106,11 @@ func FromContext(ctx context.Context) (*Ack, bool) { return a, ok } -// Fanout adjusts the Ack embedded in ctx, if any, so it represents n -// branches derived from the single incoming branch ctx currently carries. -// Call it once, before sending any of the n outputs: this guarantees the -// adjustment is visible before any of those outputs can reach a downstream -// Done/Fail call, which would otherwise be able to race ahead of it. n may -// be 0 (the incoming record produces no output and is immediately -// completed) or repeat a record multiple times (n counts sends, not -// distinct records). +// 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) @@ -145,8 +127,8 @@ func Fanout(ctx context.Context, n int) { } -// 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 +// 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 { @@ -154,45 +136,31 @@ func Drop(ctx context.Context) { } } -// Reject completes the Ack embedded in ctx, if any, as FAILED, for a record a -// task could not process. It is the counterpart of Drop: Drop means "this -// record is legitimately finished with", Reject means "this record never made -// it", so the source leaves it unacknowledged and the broker redelivers it -// instead of the pipeline waiting forever for a completion that can't come. -// -// A task bailing out mid-stream should Reject both the record it failed on and -// every record still queued behind it, since it won't be processing those -// either. +// 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 the common case of a task bailing -// out on the record it is holding: -// -// if err != nil { -// return ack.Rejected(r.Context, err) -// } -// -// Keeping the settle and the return on one line stops the two drifting apart - -// a bare `return err` here strands the record, and the symptom is the whole -// pipeline hanging at shutdown rather than anything that points back to this -// line. +// 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 produced by combining n inputs (a "fan-in"), such as join -// or archiving several records into one file. Attach it to that output -// record via WithContext before sending it. Completing the returned Ack -// (Done or Fail, however the output record's own journey downstream ends) -// transitively completes every Ack found among ctxs, so a downstream -// success or failure on the combined record is correctly attributed back -// to each record that went into it. ctxs with no Ack attached are ignored. +// 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() diff --git a/internal/pkg/pipeline/ack/tracker.go b/internal/pkg/pipeline/ack/tracker.go index 294a861..76714b4 100644 --- a/internal/pkg/pipeline/ack/tracker.go +++ b/internal/pkg/pipeline/ack/tracker.go @@ -2,39 +2,27 @@ package ack import "sync" -// Acknowledger is the broker-specific half of deferred acknowledgement: how -// one particular source settles a single message once the pipeline is done -// with it. SQS deletes the message's receipt, Kafka stores its offset, -// another broker does something else again - Tracker owns the bookkeeping -// that's common to all of them, an Acknowledger owns what isn't. -// -// Implementations are typically a small per-message struct holding whatever -// handle the client needs (a receipt, a partition and offset, ...). +// 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 rather than Done, in which case the message wasn't - // fully processed: implementations should normally leave it - // unacknowledged so the broker redelivers it instead of losing it. - // - // Ack is called at most once per message, from its own goroutine, and no - // more than the Tracker's concurrency at a time. It should log rather - // than panic on client errors, since by the time it runs the record is - // already through the pipeline and the worst case is a redelivery. + // 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 (deleting an SQS -// receipt, committing a Kafka offset) until every downstream task has -// finished with the record produced from it, without each source having to -// re-implement the bookkeeping. +// 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 limit on unacknowledged -// messages sounds prudent, but releasing it depends on downstream completion, -// which deadlocks against any fan-in task that must accumulate records before -// it can emit anything: the source stops receiving at the limit, the fan-in -// never reaches its flush threshold, so nothing ever completes and no slot -// ever frees. The number of messages in flight is instead bounded by the -// pipeline's channel capacity, which is what applies backpressure already. +// 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 { @@ -69,10 +57,9 @@ func (t *Tracker) Track(a *Ack, target Acknowledger) { <-a.Wait() - // bound how many broker calls run at once. Taking the slot here - // rather than before the wait is what keeps this safe: by now the - // record is through the pipeline, so releasing the slot depends only - // on Ack returning, never on the pipeline making further progress. + // 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 }() diff --git a/internal/pkg/pipeline/pipeline.go b/internal/pkg/pipeline/pipeline.go index 34f9714..267e841 100644 --- a/internal/pkg/pipeline/pipeline.go +++ b/internal/pkg/pipeline/pipeline.go @@ -198,9 +198,9 @@ func (p *Pipeline) distributeToChannels(input <-chan *record.Record, outputs []c } for rec := range input { - // this is a structural fan-out: the same record is duplicated to - // every parallel DAG branch, so its ack must represent all of them - // before any branch can complete it. + // 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 { @@ -267,13 +267,10 @@ func (p *Pipeline) runTaskConcurrently(t task.Task, input <-chan *record.Record, wg.Wait() - // every worker has returned. If any of them bailed out early there can - // be records left in this task's input that nobody is going to process - // and, now, nobody left to consume - which would block whatever is - // still writing upstream. Drain them and reject their acks so a source - // deferring acknowledgement redelivers them instead of waiting forever. - // This only ever finds anything when a worker returned an error; on the - // normal path the workers have already drained the channel. + // 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) @@ -285,8 +282,8 @@ func (p *Pipeline) runTaskConcurrently(t task.Task, input <-chan *record.Record, } // the output channel is closed, so downstream tasks can now drain to - // completion: this is the only safe point at which a source can wait - // for its deferred acknowledgements. + // 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) diff --git a/internal/pkg/pipeline/task/archive/tar.go b/internal/pkg/pipeline/task/archive/tar.go index 531e3b8..ecf84c8 100644 --- a/internal/pkg/pipeline/task/archive/tar.go +++ b/internal/pkg/pipeline/task/archive/tar.go @@ -20,8 +20,8 @@ type tarArchive struct { *channelStruct } -// tarFile is a regular file extracted from an archive, held until the total -// file count is known and the fan-out ack can be sized. +// tarFile is a regular file extracted from an archive, buffered until the +// total file count is known and the fan-out ack can be sized. type tarFile struct { name string data []byte @@ -42,12 +42,10 @@ func (t *tarArchive) Read() { b := rc.Data - // this is a fan-out: one archive record can expand into multiple - // file records, so the ack must represent all of them - counted up - // front, before any of them is sent - or a downstream Done/Fail for - // the first file could race ahead of a later count adjustment. tar - // readers are forward-only, so the single pass below extracts every - // regular file first and only sends them once the count is known. + // fan-out: the ack must cover every file before any of them is sent, + // or a downstream Done/Fail for the first could race ahead of a later + // count adjustment. tar readers are forward-only, so extract in a + // single pass and send once the count is known. files := make([]tarFile, 0) r := tar.NewReader(bytes.NewReader(b)) @@ -137,9 +135,8 @@ func (t *tarArchive) Write() { log.Fatal(err) } - // this is a fan-in: one archive record is produced from every input - // record consumed above, so its ack must transitively complete all of - // theirs instead of discarding all but the last. + // 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 eecf249..187d906 100644 --- a/internal/pkg/pipeline/task/archive/zip.go +++ b/internal/pkg/pipeline/task/archive/zip.go @@ -39,10 +39,9 @@ func (z *zipArchive) Read() { log.Fatal(err) } - // this is a fan-out: one archive record can expand into multiple - // file records, so the ack must represent all of them - computed - // up front, before any of them is sent - or a downstream Done/Fail - // for the first file could race ahead of a later count adjustment. + // fan-out: the ack must cover every file before any of them is sent, + // or a downstream Done/Fail for the first could race ahead of a later + // count adjustment. regularFiles := 0 for _, f := range r.File { if f.FileInfo().Mode().IsRegular() { @@ -118,9 +117,8 @@ func (z *zipArchive) Write() { log.Fatal(err) } - // this is a fan-in: one archive record is produced from every input - // record consumed above, so its ack must transitively complete all of - // theirs instead of discarding all but the last. + // 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 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 0bff05e..7c531a0 100644 --- a/internal/pkg/pipeline/task/aws/parameter_store/parameter_store.go +++ b/internal/pkg/pipeline/task/aws/parameter_store/parameter_store.go @@ -93,8 +93,8 @@ func (p *parameterStore) Run(input <-chan *record.Record, output chan<- *record. } } - // terminal (sink) mode: nothing forwards r downstream, so this task - // is the last one to touch it, once all its parameters are set. + // terminal (sink) mode: nothing forwards r on, so settle it here once + // all its parameters are set. if output == nil { if a, ok := ack.FromContext(r.Context); ok { a.Done() diff --git a/internal/pkg/pipeline/task/compress/compress.go b/internal/pkg/pipeline/task/compress/compress.go index 83dcf22..2c00f32 100644 --- a/internal/pkg/pipeline/task/compress/compress.go +++ b/internal/pkg/pipeline/task/compress/compress.go @@ -87,8 +87,8 @@ func (c *core) Run(input <-chan *record.Record, output chan<- *record.Record) (e if output != nil { c.SendData(r.Context, transformedData, output) } else if a, ok := ack.FromContext(r.Context); ok { - // terminal (sink) mode: nothing forwards this record - // downstream, so this task is the last one to touch it. + // terminal (sink) mode: nothing forwards this record on, so + // settle it here. a.Done() } } diff --git a/internal/pkg/pipeline/task/converter/converter.go b/internal/pkg/pipeline/task/converter/converter.go index 8f59886..cf32e83 100644 --- a/internal/pkg/pipeline/task/converter/converter.go +++ b/internal/pkg/pipeline/task/converter/converter.go @@ -80,18 +80,12 @@ func (c *core) Run(input <-chan *record.Record, output chan<- *record.Record) er outputs, err := c.convert(r.Data, c.Delimiter) if err != nil { - // this record is being abandoned, so settle its ack as failed: - // leaving it pending would strand a source that defers - // acknowledgement, which waits for a completion that can no longer - // come. Rejecting sends the message back for redelivery instead. return ack.Rejected(r.Context, err) } - // this is a fan-out: one input record can convert into many outputs - - // a csv row each, an xlsx sheet each - so the ack must represent all - // of them, counted before any is sent, or a downstream Done for the - // first could settle the whole record while the rest are still in - // flight. A count of 0 (nothing converted) completes it right here. + // fan-out: one input can convert into many outputs, so count them all + // before sending any, or a downstream Done for the first could settle + // the whole record while the rest are still in flight. sends := 0 for _, out := range outputs { if out.Data != nil { diff --git a/internal/pkg/pipeline/task/file/file.go b/internal/pkg/pipeline/task/file/file.go index aeb0f00..407f1ee 100644 --- a/internal/pkg/pipeline/task/file/file.go +++ b/internal/pkg/pipeline/task/file/file.go @@ -159,14 +159,13 @@ func (f *file) readFile(output chan<- *record.Record) error { } -// abort settles rc as failed, then returns err. Bailing out without doing this -// leaves a source that defers acknowledgement waiting forever for a record -// this task is never going to write; rejecting it lets the broker redeliver -// it instead. +// 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 still queued behind it: other workers -// of this task are still running and will write those. The pipeline rejects -// whatever is genuinely left over once every worker has returned. +// 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) diff --git a/internal/pkg/pipeline/task/heimdall/heimdall.go b/internal/pkg/pipeline/task/heimdall/heimdall.go index 737f5c0..f3695f0 100644 --- a/internal/pkg/pipeline/task/heimdall/heimdall.go +++ b/internal/pkg/pipeline/task/heimdall/heimdall.go @@ -78,8 +78,7 @@ func (h *heimdall) Run(input <-chan *record.Record, output chan<- *record.Record return ack.Rejected(rc.Context, err) } - // terminal (sink) mode: nothing forwards rc downstream, so this - // task is the last one to touch it. + // terminal (sink) mode: nothing forwards rc on, so settle it here. if output == nil { if a, ok := ack.FromContext(rc.Context); ok { a.Done() diff --git a/internal/pkg/pipeline/task/http/http.go b/internal/pkg/pipeline/task/http/http.go index ea273fa..3ead243 100644 --- a/internal/pkg/pipeline/task/http/http.go +++ b/internal/pkg/pipeline/task/http/http.go @@ -165,8 +165,7 @@ func (h *httpCore) Run(input <-chan *record.Record, output chan<- *record.Record return ack.Rejected(rc.Context, err) } - // terminal (sink) mode: nothing forwards rc downstream, so this - // task is the last one to touch it. + // terminal (sink) mode: nothing forwards rc on, so settle it here. if output == nil { if a, ok := ack.FromContext(rc.Context); ok { a.Done() diff --git a/internal/pkg/pipeline/task/join/join.go b/internal/pkg/pipeline/task/join/join.go index 526ee75..140676c 100644 --- a/internal/pkg/pipeline/task/join/join.go +++ b/internal/pkg/pipeline/task/join/join.go @@ -57,11 +57,10 @@ func (j *join) Run(input <-chan *record.Record, output chan<- *record.Record) er } // the input receive has to live inside the select: with a default case the - // select never blocks, so control would fall straight through to a bare - // channel receive and tickerCh would only be serviced between records - - // never firing while input is stalled, which is exactly when a partially - // filled buffer needs flushing. With duration unset tickerCh is nil, so - // this degenerates to today's plain blocking receive on input. + // 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 { case r, ok := <-input: @@ -109,9 +108,8 @@ func (j *join) sendJoinedRecords(buffer []*record.Record, output chan<- *record. ctxs[i] = r.Context } - // this is a fan-in: one output record is produced from len(buffer) - // inputs, so its ack must transitively complete every one of theirs - // instead of discarding them. + // 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 7afafc1..52fbe8c 100644 --- a/internal/pkg/pipeline/task/jq/jq.go +++ b/internal/pkg/pipeline/task/jq/jq.go @@ -28,9 +28,9 @@ func (j *jq) Run(input <-chan *record.Record, output chan<- *record.Record) (err } if output == nil { - // terminal: a jq transform with nowhere to send has no effect, but the - // input still has to be drained and each record's ack settled, or a - // source deferring acknowledgement never finishes. + // 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 { @@ -64,7 +64,7 @@ func (j *jq) Run(input <-chan *record.Record, output chan<- *record.Record) (err if splitItems, ok := items.([]any); j.Explode && ok { // marshal every item before adjusting the ack: failing partway // through afterwards would leave branches counted but never sent, - // and there's no clean way to unwind a partial fan-out. + // and a partial fan-out can't be unwound. payloads := make([][]byte, 0, len(splitItems)) for _, splitItem := range splitItems { if j.AsRaw { diff --git a/internal/pkg/pipeline/task/kafka/kafka.go b/internal/pkg/pipeline/task/kafka/kafka.go index 01f8f0b..6ad6ddf 100644 --- a/internal/pkg/pipeline/task/kafka/kafka.go +++ b/internal/pkg/pipeline/task/kafka/kafka.go @@ -159,16 +159,14 @@ func (k *kafka) write(input <-chan *record.Record) error { fmt.Printf("delivery failed for topic %s partition %d: %v\n", k.Topic, m.TopicPartition.Partition, m.TopicPartition.Error) } - // the source record failed to make it to Kafka: fail its ack - // instead of dropping it silently, so the source knows not - // to acknowledge/delete it (letting it be retried upstream). + // the source record never made it to the topic, so fail its + // ack: the source must leave it unacknowledged for retry. if a, ok := m.Opaque.(*ack.Ack); ok { a.Fail() } continue } - // only mark the source record complete once the broker has confirmed - // delivery, not merely once it's been enqueued locally. + // settle only on broker-confirmed delivery, not on local enqueue. if a, ok := m.Opaque.(*ack.Ack); ok { a.Done() } diff --git a/internal/pkg/pipeline/task/replace/replace.go b/internal/pkg/pipeline/task/replace/replace.go index e65fc31..cda3109 100644 --- a/internal/pkg/pipeline/task/replace/replace.go +++ b/internal/pkg/pipeline/task/replace/replace.go @@ -30,9 +30,9 @@ func (r *replace) Run(input <-chan *record.Record, output chan<- *record.Record) } if output == nil { - // terminal: a replace with nowhere to send has no effect, but the - // input still has to be drained and each record's ack settled, or a - // source deferring acknowledgement never finishes. + // 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 { diff --git a/internal/pkg/pipeline/task/sample/random.go b/internal/pkg/pipeline/task/sample/random.go index cb40c6e..21b5d7e 100644 --- a/internal/pkg/pipeline/task/sample/random.go +++ b/internal/pkg/pipeline/task/sample/random.go @@ -42,11 +42,10 @@ func (r *random) drain(output chan<- *record.Record) error { if l := int64(len(r.buffer)); l > 0 { - // draws are with replacement, so the same buffered record can be - // sent zero, one, or multiple times. Tally every draw first, then - // adjust each record's ack for its final send count before sending - // any of them - otherwise a downstream Done/Fail for an earlier - // send could race ahead of a later AddBranch call for the same ack. + // draws are with replacement, so a buffered record can be sent zero, + // one, or many times. Tally every draw and size each ack for its final + // send count before sending any, or a downstream Done/Fail for an + // earlier send could race ahead of a later AddBranch on the same ack. counts := make([]int, l) for i := 0; i < r.limit; i++ { diff --git a/internal/pkg/pipeline/task/sample/tail.go b/internal/pkg/pipeline/task/sample/tail.go index c46dc64..156c5dc 100644 --- a/internal/pkg/pipeline/task/sample/tail.go +++ b/internal/pkg/pipeline/task/sample/tail.go @@ -25,9 +25,8 @@ func newTail(s *sample) (sampler, error) { func (t *tail) filter(r *record.Record, _ chan<- *record.Record) error { - // the ring buffer is about to overwrite whatever record currently - // occupies this slot (if any); it will never be forwarded, so its ack - // must be completed here instead of leaking forever. + // 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.Drop(evicted.Context) } diff --git a/internal/pkg/pipeline/task/sqs/sqs.go b/internal/pkg/pipeline/task/sqs/sqs.go index c16b58c..ab83d32 100644 --- a/internal/pkg/pipeline/task/sqs/sqs.go +++ b/internal/pkg/pipeline/task/sqs/sqs.go @@ -95,20 +95,18 @@ func (s *sqs) Run(input <-chan *record.Record, output chan<- *record.Record) err return s.sendMessages(input) } - // If input is nil, act as a source: read messages and hand each one's - // 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; see Finish. + // 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) } -// Finish waits for every deferred deletion to run (and each receipt to be -// deleted or left alone) 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 only emit 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. +// 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 { if s.tracker != nil { @@ -162,9 +160,8 @@ func (s *sqs) getMessages(ctx context.Context, output chan<- *record.Record) err for _, m := range receiveMessageOutput.Messages { - // nothing to forward to, so there's no downstream ack to - // wait for: delete the receipt right away, same as when - // there's no consumer at all. + // 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 @@ -192,9 +189,8 @@ type messageAck struct { } // Ack deletes the message's receipt so SQS doesn't redeliver it. On a -// downstream failure it does nothing: the message wasn't fully processed, so -// leaving the receipt alone lets SQS redeliver it once the visibility -// timeout expires instead of losing it. +// 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) { if failed { @@ -205,10 +201,9 @@ func (m *messageAck) Ack(failed bool) { } -// deleteMessage acknowledges a message by deleting its receipt so it isn't -// redelivered. A failure to delete is logged rather than returned: the -// message has already been processed, and the worst case is a redelivery -// after the visibility timeout. +// 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{ diff --git a/internal/pkg/pipeline/task/task.go b/internal/pkg/pipeline/task/task.go index 0c6d530..9fc339d 100644 --- a/internal/pkg/pipeline/task/task.go +++ b/internal/pkg/pipeline/task/task.go @@ -39,11 +39,10 @@ type Task interface { } // Finisher is implemented by tasks with work to do only once their output -// channel has been closed. A source deferring acknowledgement is the case -// this exists for: its acks can only settle after every downstream task has -// drained, and downstream tasks that emit on input close - join, archive -// pack, sample tail/random - can only drain once the source's output channel -// is closed. Waiting inside Run would therefore deadlock. +// 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. @@ -125,10 +124,8 @@ func (b *Base) SendData(ctx context.Context, data []byte, output chan<- *record. func (b *Base) SendRecord(r *record.Record, output chan<- *record.Record) /* we should return error here */ { if output == nil { - // terminal task: nothing forwards r downstream, so this is the last - // place that will ever touch it. Settle its ack here rather than - // dropping it, or a source deferring acknowledgement waits forever - // for a completion that can no longer come. + // terminal task: nothing forwards r downstream, so settle its ack here + // or a source deferring acknowledgement waits forever for it. ack.Drop(r.Context) return } diff --git a/internal/pkg/pipeline/task/xpath/xpath.go b/internal/pkg/pipeline/task/xpath/xpath.go index 50dedd2..825df6b 100644 --- a/internal/pkg/pipeline/task/xpath/xpath.go +++ b/internal/pkg/pipeline/task/xpath/xpath.go @@ -54,13 +54,10 @@ func (x *xpath) Run(input <-chan *record.Record, output chan<- *record.Record) e } } - // this is a fan-out: one input record yields one output per container - // node, so the ack must represent all of them, counted before any is - // sent - otherwise a downstream Done for the first node could settle - // the whole record while later nodes are still in flight. queryFields - // is pure, so extracting every node up front costs only memory, and - // the original node position is carried along since node_index is - // part of the output contract. + // fan-out: one output per container node, counted before any is sent, + // or a downstream Done for the first could settle the whole record + // while later nodes are still in flight. node_index is part of the + // output contract, so the original position travels with the data. type nodePayload struct { index int data []byte