diff --git a/internal/pkg/pipeline/ack/ack.go b/internal/pkg/pipeline/ack/ack.go new file mode 100644 index 0000000..8450db0 --- /dev/null +++ b/internal/pkg/pipeline/ack/ack.go @@ -0,0 +1,176 @@ +// Package ack tracks completion of a single source record as it flows through +// a pipeline, so the source task can defer acknowledging it until every +// downstream branch produced from it has finished processing. +package ack + +import ( + "context" + "sync/atomic" +) + +type catterpillarAckKey string + +const CATERPILLAR_ACK catterpillarAckKey = "CATERPILLAR_ACK" + +// Ack is created once per source record with exactly one pending branch: the +// record itself. It rides downstream on the record's context.Context, so tasks +// that only transform a record need no explicit wiring. +// +// Tasks that change a record's branch count must say so before sending: Fanout +// for one-to-many, Joined for many-to-one, Drop or Reject for a record they +// don't forward. A task with a nil output channel settles every record it +// consumes. +type Ack struct { + remaining atomic.Int32 + failed atomic.Bool + done chan struct{} + children []*Ack // settled with this Ack's own outcome once it completes; see Joined +} + +// New returns an Ack with a single pending branch. +func New() *Ack { + a := &Ack{done: make(chan struct{})} + a.remaining.Store(1) + return a +} + +// AddBranch registers cnt additional branches that must call Done or Fail +// before Wait's channel closes. +func (a *Ack) AddBranch(cnt int32) { + a.remaining.Add(cnt) +} + +// Done marks one branch as complete. +func (a *Ack) Done() { + a.complete(false) +} + +// Fail marks one branch as complete but unsuccessful, so Failed reports true +// once every branch has finished. Use it when a branch didn't make it to where +// it needed to go, so the source knows not to acknowledge the record. +func (a *Ack) Fail() { + a.complete(true) +} + +func (a *Ack) complete(failed bool) { + + if failed { + a.failed.Store(true) + } + + if a.remaining.Add(-1) != 0 { + return + } + + // a joined record's single downstream outcome applies equally to every + // record that went into it, so children get this Ack's overall result + // rather than the outcome of this particular call. + anyFailed := a.failed.Load() + for _, c := range a.children { + if anyFailed { + c.Fail() + } else { + c.Done() + } + } + + close(a.done) + +} + +// Failed reports whether any branch called Fail instead of Done. Only +// meaningful after Wait's channel has closed. +func (a *Ack) Failed() bool { + return a.failed.Load() +} + +// Wait returns a channel that closes once every branch has called Done or +// Fail. +func (a *Ack) Wait() <-chan struct{} { + return a.done +} + +// WithContext returns a copy of ctx carrying a, recoverable later via +// FromContext. A nil ctx is treated as context.Background(), since an +// aggregating task may never have received a record to inherit one from. +func WithContext(ctx context.Context, a *Ack) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, CATERPILLAR_ACK, a) +} + +// FromContext recovers the Ack embedded in ctx, if any. +func FromContext(ctx context.Context) (*Ack, bool) { + a, ok := ctx.Value(CATERPILLAR_ACK).(*Ack) + return a, ok +} + +// Fanout adjusts the Ack embedded in ctx, if any, so it represents n branches +// derived from the single branch ctx currently carries. Call it once, before +// sending any of the n outputs, so no downstream Done/Fail can race ahead of +// the adjustment. n counts sends rather than distinct records; n == 0 +// completes the Ack immediately. +func Fanout(ctx context.Context, n int) { + + a, ok := FromContext(ctx) + if !ok { + return + } + + switch { + case n == 0: + a.Done() + case n > 1: + a.AddBranch(int32(n - 1)) + } + +} + +// Drop completes the Ack embedded in ctx, if any, for a record a task decided +// not to forward downstream (e.g. it was filtered out). It is the +// single-record equivalent of Fanout(ctx, 0). +func Drop(ctx context.Context) { + if a, ok := FromContext(ctx); ok { + a.Done() + } +} + +// Reject completes the Ack embedded in ctx, if any, as failed, for a record a +// task could not process. Where Drop means the record is legitimately finished +// with, Reject means it never made it: the source leaves it unacknowledged and +// the broker redelivers it, rather than the pipeline waiting for a completion +// that can't come. +func Reject(ctx context.Context) { + if a, ok := FromContext(ctx); ok { + a.Fail() + } +} + +// Rejected is Reject followed by err, for a task bailing out on the record it +// is holding. Keeping the settle and the return on one line stops the two +// drifting apart: a bare return strands the record, and the symptom is the +// whole pipeline hanging at shutdown rather than anything pointing here. +func Rejected(ctx context.Context, err error) error { + Reject(ctx) + return err +} + +// Joined returns a new Ack with a single pending branch representing one output +// record combined from several inputs. Attach it to that output via +// WithContext: completing it transitively completes every Ack found among +// ctxs, so the combined record's outcome is attributed back to each record +// that went into it. ctxs with no Ack attached are ignored. +func Joined(ctxs ...context.Context) *Ack { + + a := New() + + for _, c := range ctxs { + if child, ok := FromContext(c); ok { + a.children = append(a.children, child) + } + } + + return a + +} diff --git a/internal/pkg/pipeline/ack/tracker.go b/internal/pkg/pipeline/ack/tracker.go new file mode 100644 index 0000000..76714b4 --- /dev/null +++ b/internal/pkg/pipeline/ack/tracker.go @@ -0,0 +1,76 @@ +package ack + +import "sync" + +// Acknowledger is the broker-specific half of deferred acknowledgement: how a +// source settles a single message once the pipeline is done with it. Tracker +// owns the bookkeeping common to every broker, an Acknowledger owns what isn't. +type Acknowledger interface { + // Ack settles one message. failed reports whether any downstream branch + // signalled Fail, in which case implementations should normally leave the + // message unacknowledged so the broker redelivers it. Called at most once + // per message, from its own goroutine, and at most the Tracker's + // concurrency at a time. + Ack(failed bool) +} + +// Tracker lets a source task defer acknowledging a message until every +// downstream task has finished with the record produced from it, without each +// source having to re-implement the bookkeeping. +// +// Nothing here gates the source's receive loop: a cap on unacknowledged +// messages would deadlock against any fan-in task that must accumulate records +// before it can emit, since freeing a slot depends on the very completion the +// fan-in is waiting to produce. Messages in flight are bounded by the +// pipeline's channel capacity, which applies backpressure already. +// +// A Tracker must be created with NewTracker; the zero value is not usable. +type Tracker struct { + slots chan struct{} // bounds concurrent Ack calls; acquired only after settling + wg sync.WaitGroup +} + +// NewTracker returns a Tracker that runs at most concurrency Ack calls at a +// time. A value below 1 is treated as 1. +func NewTracker(concurrency int) *Tracker { + + if concurrency < 1 { + concurrency = 1 + } + + return &Tracker{slots: make(chan struct{}, concurrency)} + +} + +// Track watches a in the background and calls target.Ack once every +// downstream task has signalled Done or Fail for it, passing on whether any +// of them failed. It does not block. +// +// Every Track for a given Tracker must happen before its Wait. +func (t *Tracker) Track(a *Ack, target Acknowledger) { + + t.wg.Add(1) + + go func() { + + defer t.wg.Done() + + <-a.Wait() + + // take the slot after the wait, not before: the record is already + // through the pipeline, so freeing it depends only on Ack returning + // and never on the pipeline making further progress. + t.slots <- struct{}{} + defer func() { <-t.slots }() + + target.Ack(a.Failed()) + + }() + +} + +// Wait blocks until every tracked Ack has settled and its acknowledgement +// has been carried out. It is idempotent. +func (t *Tracker) Wait() { + t.wg.Wait() +} diff --git a/internal/pkg/pipeline/pipeline.go b/internal/pkg/pipeline/pipeline.go index 8540082..267e841 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 { + // structural fan-out: the record is duplicated to every parallel DAG + // branch, so its ack must represent all of them before any branch can + // complete it. + ack.Fanout(rec.Context, branches) for _, ch := range outputs { if ch != nil { ch <- rec @@ -251,11 +263,39 @@ func (p *Pipeline) runTaskConcurrently(t task.Task, input <-chan *record.Record, }(t, input, output) } - go func(wg *sync.WaitGroup, out chan<- *record.Record) { + go func(t task.Task, wg *sync.WaitGroup, in <-chan *record.Record, out chan<- *record.Record) { + wg.Wait() + + // a worker that bailed out early can leave records in this task's + // input with nobody left to consume them, blocking upstream writers. + // Reject them so a source deferring acknowledgement redelivers them + // rather than waiting forever. + if in != nil { + for r := range in { + ack.Reject(r.Context) + } + } + if out != nil { close(out) } + + // the output channel is closed, so downstream tasks can now drain to + // completion: the only safe point at which a source can wait for its + // deferred acknowledgements. + if f, ok := t.(task.Finisher); ok { + if err := f.Finish(); err != nil { + fmt.Printf("error finishing %s: %s\n", t.GetName(), err) + if t.GetFailOnError() { + p.locker.Lock() + p.errors[t.GetName()] = err + p.locker.Unlock() + } + } + } + p.wg.Done() - }(&taskWg, output) + + }(t, &taskWg, input, output) } diff --git a/internal/pkg/pipeline/task/archive/tar.go b/internal/pkg/pipeline/task/archive/tar.go index e2d1198..ecf84c8 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" @@ -18,6 +20,13 @@ type tarArchive struct { *channelStruct } +// 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 +} + func (t *tarArchive) Read() { for { @@ -27,11 +36,18 @@ func (t *tarArchive) Read() { } if len(rc.Data) == 0 { + ack.Drop(rc.Context) continue } b := rc.Data + // 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)) for { @@ -44,15 +60,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) } } } @@ -62,12 +89,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 +135,9 @@ func (t *tarArchive) Write() { log.Fatal(err) } - t.SendData(rc.Context, buf.Bytes(), t.OutputChan) + // fan-in: the archive record is produced from every input consumed above, + // so its ack must transitively complete all of theirs. + joinedAck := ack.Joined(ctxs...) + + t.SendData(ack.WithContext(rc.Context, joinedAck), buf.Bytes(), t.OutputChan) } diff --git a/internal/pkg/pipeline/task/archive/zip.go b/internal/pkg/pipeline/task/archive/zip.go index bf0b694..187d906 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,18 @@ func (z *zipArchive) Read() { if err != nil { log.Fatal(err) } + + // 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() { + regularFiles++ + } + } + ack.Fanout(rc.Context, regularFiles) + for _, f := range r.File { // check the file type is regular file @@ -66,6 +81,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 +110,18 @@ func (z *zipArchive) Write() { } rc.Context = rec.Context + ctxs = append(ctxs, rec.Context) } if err := zipWriter.Close(); err != nil { log.Fatal(err) } + // fan-in: the archive record is produced from every input consumed above, + // so its ack must transitively complete all of theirs. + joinedAck := ack.Joined(ctxs...) + // Send the complete ZIP archive - z.SendData(rc.Context, zipBuf.Bytes(), z.OutputChan) + z.SendData(ack.WithContext(rc.Context, joinedAck), zipBuf.Bytes(), z.OutputChan) } diff --git a/internal/pkg/pipeline/task/aws/parameter_store/parameter_store.go b/internal/pkg/pipeline/task/aws/parameter_store/parameter_store.go index 365998b..7c531a0 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" ) @@ -65,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{ @@ -84,13 +85,21 @@ 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 { p.SendRecord(r, output) } } + + // 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() + } + } } return nil diff --git a/internal/pkg/pipeline/task/compress/compress.go b/internal/pkg/pipeline/task/compress/compress.go index be7e20a..2c00f32 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 } @@ -68,21 +70,26 @@ func (c *core) Run(input <-chan *record.Record, output chan<- *record.Record) (e var err error if c.Action == defaultAction { if transformedData, err = c.compress(r); err != nil { - return err + return ack.Rejected(r.Context, err) } } else { if transformedData, err = c.decompress(r); err != nil { - return err + return ack.Rejected(r.Context, err) } } // skip empty transformed data if len(transformedData) == 0 { + ack.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 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 40827b1..cf32e83 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,20 @@ func (c *core) Run(input <-chan *record.Record, output chan<- *record.Record) er outputs, err := c.convert(r.Data, c.Delimiter) if err != nil { - return err + return ack.Rejected(r.Context, err) } + // 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 { + 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/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..407f1ee 100644 --- a/internal/pkg/pipeline/task/file/file.go +++ b/internal/pkg/pipeline/task/file/file.go @@ -11,6 +11,7 @@ import ( "strings" "github.com/patterninc/caterpillar/internal/pkg/config" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" "github.com/patterninc/caterpillar/internal/pkg/textutil" @@ -158,6 +159,21 @@ func (f *file) readFile(output chan<- *record.Record) error { } +// abort settles rc as failed, then returns err, so a source deferring +// acknowledgement redelivers the record rather than waiting forever for a +// write that will never happen. +// +// Only rc is rejected, never the records queued behind it: sibling workers are +// still running and will write those, and the pipeline rejects whatever is +// genuinely left over once every worker has returned. +func (f *file) abort(rc *record.Record, err error) error { + + ack.Reject(rc.Context) + + return err + +} + func (f *file) writeFile(input <-chan *record.Record) error { for { @@ -169,13 +185,13 @@ func (f *file) writeFile(input <-chan *record.Record) error { // Evaluate the path with the record context path, err := f.Path.Get(rc) if err != nil { - return err + return f.abort(rc, err) } // Determine the scheme from the evaluated path parsedURL, err := url.Parse(path) if err != nil { - return err + return f.abort(rc, err) } pathScheme := parsedURL.Scheme if pathScheme == `` { @@ -198,10 +214,14 @@ 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 { + a.Done() } } 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 02d13c5..f3695f0 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" ) @@ -68,13 +69,20 @@ 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 on, so settle it here. + 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..3ead243 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" @@ -158,10 +159,17 @@ 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 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 00d06f5..140676c 100644 --- a/internal/pkg/pipeline/task/join/join.go +++ b/internal/pkg/pipeline/task/join/join.go @@ -7,6 +7,7 @@ import ( "time" "github.com/patterninc/caterpillar/internal/pkg/duration" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" ) @@ -55,11 +56,14 @@ func (j *join) Run(input <-chan *record.Record, output chan<- *record.Record) er tickerCh = ticker.C } + // the input receive has to live inside the select: with a default case the + // select never blocks, so tickerCh is only serviced between records and + // never fires while input is stalled - which is exactly when a partially + // filled buffer needs flushing. With duration unset tickerCh is nil and + // this is a plain blocking receive on input. for { select { - default: - // Try to get a record from input - r, ok := j.GetRecord(input) + case r, ok := <-input: if !ok { // Input channel closed, send any remaining records j.flushBuffer(&buffer, output) @@ -95,13 +99,18 @@ func (j *join) sendJoinedRecords(buffer []*record.Record, output chan<- *record. // Join all data with the specified delimiter var joinedData strings.Builder + ctxs := make([]context.Context, len(buffer)) for i, r := range buffer { if i > 0 { joinedData.WriteString(j.Delimiter) } joinedData.Write(r.Data) + ctxs[i] = r.Context } - j.SendData(ctx, []byte(joinedData.String()), output) + // fan-in: the output record is produced from every buffered input, so its + // ack must transitively complete every one of theirs. + joinedAck := ack.Joined(ctxs...) + j.SendData(ack.WithContext(ctx, joinedAck), []byte(joinedData.String()), output) } diff --git a/internal/pkg/pipeline/task/jq/jq.go b/internal/pkg/pipeline/task/jq/jq.go index bb3c9d1..52fbe8c 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" ) @@ -22,49 +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: the transform has no effect, but input still has to be + // drained and each record's ack settled, or a source deferring + // acknowledgement never finishes. for { r, ok := j.GetRecord(input) if !ok { - break + return nil } + ack.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 a partial fan-out can't be unwound. + 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 { - 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 { - 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/kafka/kafka.go b/internal/pkg/pipeline/task/kafka/kafka.go index a5418d6..6ad6ddf 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,26 @@ 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 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 + } + // settle only on broker-confirmed delivery, not on local enqueue. + if a, ok := m.Opaque.(*ack.Ack); ok { + a.Done() } } }() @@ -169,9 +186,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/replace/replace.go b/internal/pkg/pipeline/task/replace/replace.go index e7d724a..cda3109 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: the replacement has no effect, but input still has to be + // drained and each record's ack settled, or a source deferring + // acknowledgement never finishes. for { record, ok := r.GetRecord(input) if !ok { - break + return nil } - r.SendData(record.Context, []byte(rx.ReplaceAllString(string(record.Data), r.Replacement)), output) + ack.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/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..21b5d7e 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,12 @@ func (r *random) filter(row *record.Record, _ chan<- *record.Record) error { func (r *random) drain(output chan<- *record.Record) error { if l := int64(len(r.buffer)); l > 0 { + + // draws are with replacement, so a buffered record can be sent zero, + // one, or many times. 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++ { index, err := rand.Int(rand.Reader, big.NewInt(l)) @@ -45,9 +54,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/sample.go b/internal/pkg/pipeline/task/sample/sample.go index 80d6fd2..0f67c4c 100644 --- a/internal/pkg/pipeline/task/sample/sample.go +++ b/internal/pkg/pipeline/task/sample/sample.go @@ -3,6 +3,7 @@ package sample import ( "fmt" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" ) @@ -72,7 +73,7 @@ func (s *sample) Run(input <-chan *record.Record, output chan<- *record.Record) break } if err := sampler.filter(r, output); err != nil { - return err + return ack.Rejected(r.Context, err) } } diff --git a/internal/pkg/pipeline/task/sample/tail.go b/internal/pkg/pipeline/task/sample/tail.go index e394d6b..156c5dc 100644 --- a/internal/pkg/pipeline/task/sample/tail.go +++ b/internal/pkg/pipeline/task/sample/tail.go @@ -1,6 +1,7 @@ package sample import ( + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" ) @@ -24,6 +25,12 @@ func newTail(s *sample) (sampler, error) { func (t *tail) filter(r *record.Record, _ chan<- *record.Record) error { + // the ring buffer is about to overwrite this slot; the evicted record will + // never be forwarded, so settle its ack here. + if evicted := t.buffer[t.index]; evicted != nil { + ack.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..21c07b1 100644 --- a/internal/pkg/pipeline/task/sftp/operations.go +++ b/internal/pkg/pipeline/task/sftp/operations.go @@ -10,6 +10,7 @@ import ( "github.com/bmatcuk/doublestar" pkgsftp "github.com/pkg/sftp" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" "github.com/patterninc/caterpillar/internal/pkg/textutil" @@ -29,11 +30,15 @@ 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 { + a.Done() } } diff --git a/internal/pkg/pipeline/task/sns/sns.go b/internal/pkg/pipeline/task/sns/sns.go index cec5c03..e127757 100644 --- a/internal/pkg/pipeline/task/sns/sns.go +++ b/internal/pkg/pipeline/task/sns/sns.go @@ -11,6 +11,7 @@ import ( "github.com/aws/aws-sdk-go-v2/service/sns/types" "github.com/google/uuid" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" ) @@ -103,7 +104,11 @@ 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 { + a.Done() } } 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/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 1dcf05b..ab83d32 100644 --- a/internal/pkg/pipeline/task/sqs/sqs.go +++ b/internal/pkg/pipeline/task/sqs/sqs.go @@ -6,7 +6,6 @@ import ( "fmt" "regexp" "strings" - "sync" "time" "github.com/aws/aws-sdk-go-v2/aws" @@ -14,16 +13,16 @@ import ( qs "github.com/aws/aws-sdk-go-v2/service/sqs" "github.com/google/uuid" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" "github.com/patterninc/caterpillar/internal/pkg/pipeline/task" ) const ( - defaultConcurrency = 10 - defaultMaxMessages = 10 - defaultWaitTimeSeconds = 10 - receiptsQueueMultiplier = 1000 - defaultRegion = "us-west-2" + defaultConcurrency = 10 + defaultMaxMessages = 10 + defaultWaitTimeSeconds = 10 + defaultRegion = "us-west-2" ) var ( @@ -40,7 +39,8 @@ type sqs struct { ExitOnEmpty bool `yaml:"exit_on_empty,omitempty" json:"exit_on_empty,omitempty"` MessageGroupId string `yaml:"message_group_id,omitempty" json:"message_group_id,omitempty"` // used for FIFO queues - client *qs.Client + client *qs.Client + tracker *ack.Tracker } func New() (task.Task, error) { @@ -67,6 +67,8 @@ func (s *sqs) Init() error { } s.client = qs.NewFromConfig(awsConfig) + s.tracker = ack.NewTracker(s.Concurrency) + return nil } @@ -88,33 +90,34 @@ func (s *sqs) extractRegionFromQueueURL() string { func (s *sqs) Run(input <-chan *record.Record, output chan<- *record.Record) error { - // Client is already initialized in RunPreHook - just use it + // Client is already initialized in Init - just use it if input != nil { return s.sendMessages(input) } - // If input is nil, act as a source: start getMessages and receipt workers - // let's create channel to which getMessages function will communicate messages receipts - receipts := make(chan *string, s.Concurrency*receiptsQueueMultiplier) + // If input is nil, act as a source: read messages and hand each receipt to + // the tracker, which deletes it once every downstream task has finished + // with the record it produced. Finish, not Run, waits for those deletions. + return s.getMessages(ctx, output) - // we set a pool of workers that will delete messages from the queue - var wg sync.WaitGroup - wg.Add(s.Concurrency) - for i := 0; i < s.Concurrency; i++ { - go s.processReceipts(receipts, &wg) - } +} - err := s.getMessages(ctx, output, receipts) +// Finish waits for every deferred deletion before the pipeline treats this +// task as complete, so a shutdown never abandons in-flight acknowledgements. +// It can't happen in Run: downstream tasks that emit only once their input +// closes can't finish with a record until this task's output channel is +// closed, which the pipeline does only after Run returns. +func (s *sqs) Finish() error { - wg.Wait() + if s.tracker != nil { + s.tracker.Wait() + } - return err + return nil } -func (s *sqs) getMessages(ctx context.Context, output chan<- *record.Record, receipts chan *string) error { - - defer close(receipts) +func (s *sqs) getMessages(ctx context.Context, output chan<- *record.Record) error { // do we need to stop pipeline after a while? if s.EndAfter > 0 { @@ -156,34 +159,59 @@ func (s *sqs) getMessages(ctx context.Context, output chan<- *record.Record, rec } for _, m := range receiveMessageOutput.Messages { - // create new record and send it downstream - if output != nil { - s.SendData(ctx, []byte(*m.Body), output) + // nothing to forward to, so there's no downstream ack to wait + // for: delete the receipt right away. + if output == nil { + s.deleteMessage(m.MessageId, m.ReceiptHandle) + continue } - // send receipt to receipts channel for deletion - receipts <- m.ReceiptHandle + msgAck := ack.New() + s.SendData(ack.WithContext(ctx, msgAck), []byte(*m.Body), output) + + s.tracker.Track(msgAck, &messageAck{ + sqs: s, + messageId: m.MessageId, + receiptHandle: m.ReceiptHandle, + }) } } } } -func (s *sqs) processReceipts(receipts <-chan *string, wg *sync.WaitGroup) error { +// messageAck acknowledges one received message on behalf of ack.Tracker. +type messageAck struct { + sqs *sqs + messageId *string + receiptHandle *string +} - defer wg.Done() +// Ack deletes the message's receipt so SQS doesn't redeliver it. On a +// downstream failure it does nothing: leaving the receipt alone lets SQS +// redeliver the message once the visibility timeout expires. +func (m *messageAck) Ack(failed bool) { - for receipt := range receipts { - if _, err := s.client.DeleteMessage(ctx, &qs.DeleteMessageInput{ - QueueUrl: &s.QueueURL, - ReceiptHandle: receipt, - }); err != nil { - return err - } + if failed { + return } - return nil + m.sqs.deleteMessage(m.messageId, m.receiptHandle) + +} + +// deleteMessage acknowledges a message by deleting its receipt. A failure is +// logged rather than returned: the message has already been processed, so the +// worst case is a redelivery after the visibility timeout. +func (s *sqs) deleteMessage(messageId, receiptHandle *string) { + + if _, err := s.client.DeleteMessage(ctx, &qs.DeleteMessageInput{ + QueueUrl: &s.QueueURL, + ReceiptHandle: receiptHandle, + }); err != nil { + fmt.Printf("failed to delete message %s from queue %s: %v\n", aws.ToString(messageId), s.QueueURL, err) + } } @@ -203,7 +231,11 @@ 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 { + a.Done() } } return nil diff --git a/internal/pkg/pipeline/task/task.go b/internal/pkg/pipeline/task/task.go index 776a067..9fc339d 100644 --- a/internal/pkg/pipeline/task/task.go +++ b/internal/pkg/pipeline/task/task.go @@ -7,6 +7,7 @@ import ( "sync" "github.com/patterninc/caterpillar/internal/pkg/jq" + "github.com/patterninc/caterpillar/internal/pkg/pipeline/ack" "github.com/patterninc/caterpillar/internal/pkg/pipeline/record" ) @@ -37,6 +38,18 @@ type Task interface { Init() error // Called once after unmarshaling, before pipeline execution } +// Finisher is implemented by tasks with work to do only once their output +// channel has been closed. Deferred acknowledgement is the case this exists +// for: a source's acks settle only after every downstream task has drained, +// and a downstream task that emits on input close can't drain until the +// source's output channel is closed, so waiting inside Run would deadlock. +// +// The pipeline calls Finish exactly once per task, after every worker of that +// task has returned from Run and after the task's output channel is closed. +type Finisher interface { + Finish() error +} + type Base struct { Name string `yaml:"name,omitempty" json:"name,omitempty"` Type string `yaml:"type,omitempty" json:"type,omitempty"` @@ -111,6 +124,9 @@ 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 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 29b6f2f..825df6b 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,42 @@ 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 } } + // 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 + } + + 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/setup_localstack_sqs.sh b/test/pipelines/setup_localstack_sqs.sh new file mode 100755 index 0000000..0497ca8 --- /dev/null +++ b/test/pipelines/setup_localstack_sqs.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Creates/refills the local SQS queue used by sqs_with_context_concurrency.yaml +# against a LocalStack instance running on localhost:4566. +# +# Requires: LocalStack running (`localstack start` or the localstack/localstack +# docker image), plus aws-cli and jq installed locally. + +ENDPOINT="http://localhost:4566" +REGION="us-west-2" +QUEUE_NAME="local-sqs-context-concurrency-queue" # matches queue_url in sqs_with_context_concurrency.yaml +TOTAL_MESSAGES=100 +BATCH_SIZE=10 # SQS SendMessageBatch max + +export AWS_ACCESS_KEY_ID="${AWS_ACCESS_KEY_ID:-test}" +export AWS_SECRET_ACCESS_KEY="${AWS_SECRET_ACCESS_KEY:-test}" +export AWS_DEFAULT_REGION="$REGION" + +awscli() { + aws --endpoint-url "$ENDPOINT" --region "$REGION" "$@" +} + +echo "Creating queue '$QUEUE_NAME'..." +QUEUE_URL=$(awscli sqs create-queue --queue-name "$QUEUE_NAME" --query 'QueueUrl' --output text) +echo "Queue URL: $QUEUE_URL" + +echo "Pushing $TOTAL_MESSAGES random messages..." +for ((batch_start=0; batch_start/dev/null +done + +echo "Done. Sent $TOTAL_MESSAGES messages to queue '$QUEUE_NAME'." +echo +echo "Queue URL (LocalStack): $QUEUE_URL" +echo +echo "To point the pipeline at LocalStack, run it with:" +echo " AWS_ENDPOINT_URL=$ENDPOINT AWS_ACCESS_KEY_ID=test AWS_SECRET_ACCESS_KEY=test AWS_DEFAULT_REGION=$REGION go run cmd/caterpillar/caterpillar.go -conf test/pipelines/sqs_with_context_concurrency.yaml" diff --git a/test/pipelines/sqs_fanout_fanin_dag.yaml b/test/pipelines/sqs_fanout_fanin_dag.yaml new file mode 100644 index 0000000..b715cdb --- /dev/null +++ b/test/pipelines/sqs_fanout_fanin_dag.yaml @@ -0,0 +1,66 @@ +# Exercises acknowledgment across every shape that changes a record's branch +# count, in one DAG. Acks only exist when the source creates them, so this has +# to be SQS-driven to test anything: with any other source, Fanout/Joined are +# no-ops. +# +# Each message body is {"id": N, "items": [1,2,3,4,5]}. +# +# read_queue 1 record per message +# >> explode_items 5 records (task fan-out: jq explode) +# >> [tag_a, tag_b] 10 records (structural fan-out: DAG branches) +# >> batch (fan-in: join, 7 at a time) +# >> save +# +# With the default 20 messages: 20 x 5 x 2 = 200 records into join, so 28 full +# batches of 7 plus a final partial batch of 4 = 29 files, 200 records, and +# each item value appearing exactly 40 times (20 messages x 2 branches). +# +# The message is only deleted once all 10 of its branches have landed, so a +# correct run ends with the queue at 0 visible / 0 in flight. Any leak shows up +# as leftover depth; any premature ack shows up as a short record count. +# +# join's batch size is deliberately NOT a divisor of the record count, so +# records are always left buffered when the source stops reading. That is the +# case that deadlocks if a source waits for its acknowledgements before its +# output channel closes: join can only flush those last records once the +# channel closes, and the channel can only close once the source stops +# waiting. Keep it indivisible or this fixture stops testing that. +# +# Seed the queue first (see test/pipelines/setup_localstack_fanout.sh), then: +# AWS_ENDPOINT_URL=http://localhost:4566 AWS_ACCESS_KEY_ID=test \ +# AWS_SECRET_ACCESS_KEY=test AWS_DEFAULT_REGION=us-west-2 \ +# go run ./cmd/caterpillar -conf test/pipelines/sqs_fanout_fanin_dag.yaml + +tasks: + - name: read_queue + type: sqs + queue_url: http://sqs.us-west-2.localhost.localstack.cloud:4566/000000000000/local-sqs-fanout-fanin-queue + exit_on_empty: true + + # fan-out: one message becomes one record per item in the array + - name: explode_items + type: jq + path: '.items' + explode: true + + # structural fan-out: the DAG duplicates every record into both branches + - name: tag_a + type: jq + path: '{value: ., branch: "a"}' + + - name: tag_b + type: jq + path: '{value: ., branch: "b"}' + + # fan-in: seven records become one, and its ack must transitively complete + # all seven + - name: batch + type: join + number: 7 + delimiter: "\n" + + - name: save + type: file + path: ./output/fanout/{{ macro "uuid" }}.json + +dag: read_queue >> explode_items >> [tag_a,tag_b] >> batch >> save diff --git a/test/pipelines/sqs_with_context_concurrency.yaml b/test/pipelines/sqs_with_context_concurrency.yaml index 27390f0..611ea65 100644 --- a/test/pipelines/sqs_with_context_concurrency.yaml +++ b/test/pipelines/sqs_with_context_concurrency.yaml @@ -1,7 +1,7 @@ tasks: - name: read_queue type: sqs - queue_url: https://sqs.us-west-2.amazonaws.com/123456789012/my-queue + queue_url: http://sqs.us-west-2.localhost.localstack.cloud:4566/000000000000/local-sqs-context-concurrency-queue exit_on_empty: true - name: extract_urls type: jq