diff --git a/internal/ioutil/ioutil.go b/internal/ioutil/ioutil.go new file mode 100644 index 0000000..75167cb --- /dev/null +++ b/internal/ioutil/ioutil.go @@ -0,0 +1,86 @@ +package ioutil + +import ( + "fmt" + "io" +) + +// HeadWriter keeps the first 'Limit' bytes in memory. +type HeadWriter struct { + Buf []byte + Limit int +} + +var _ io.Writer = &HeadWriter{} + +func (w *HeadWriter) Write(p []byte) (n int, err error) { + if len(w.Buf) >= w.Limit { + return len(p), nil + } + w.Buf = append(w.Buf, p...) + if len(w.Buf) > w.Limit { + w.Buf = w.Buf[:w.Limit] + } + return len(p), nil +} + +func (w *HeadWriter) Bytes() []byte { + return w.Buf +} + +// tailWriter keeps the last 'Limit' bytes in memory. +type TailWriter struct { + Buf []byte + Limit int +} + +var _ io.Writer = &TailWriter{} + +func (w *TailWriter) Write(p []byte) (n int, err error) { + w.Buf = append(w.Buf, p...) + if len(w.Buf) > w.Limit { + w.Buf = w.Buf[len(w.Buf)-w.Limit:] + } + return len(p), nil +} + +func (w *TailWriter) Bytes() []byte { + return w.Buf +} + +type OutputCapturer struct { + HeadWriter + TailWriter + Limit int + totalBytes int +} + +var _ io.Writer = &OutputCapturer{} + +func NewOutputCapturer(limit int) *OutputCapturer { + return &OutputCapturer{ + HeadWriter: HeadWriter{Limit: limit}, + TailWriter: TailWriter{Limit: limit}, + Limit: limit, + } +} + +func (w *OutputCapturer) Write(p []byte) (n int, err error) { + w.HeadWriter.Write(p) + w.TailWriter.Write(p) + w.totalBytes += len(p) + return len(p), nil +} + +func (w *OutputCapturer) String() string { + head := w.HeadWriter.Bytes() + tail := w.TailWriter.Bytes() + if w.totalBytes <= w.Limit { + return string(head) + } + + head = head[:w.Limit/2] + tail = tail[len(tail)-w.Limit/2:] + + return fmt.Sprintf("%s...[%v bytes dropped]...%s", string(head), w.totalBytes-len(head)-len(tail), string(tail)) +} diff --git a/internal/orchestrator/task.go b/internal/orchestrator/task.go index 96441cf..f4de386 100644 --- a/internal/orchestrator/task.go +++ b/internal/orchestrator/task.go @@ -67,8 +67,8 @@ func (t *TaskWithOperation) runWithOpAndContext(ctx context.Context, do func(ctx err := do(ctx, t.op) - if bytes := buf.Bytes(); len(bytes) > 0 { - ref, e := t.orch.logStore.Write(bytes) + if str := capture.String(); len(str) > 0 { + ref, e := t.orch.logStore.Write([]byte(str)) if e != nil { errors.Join(err, fmt.Errorf("failed to write log to logstore: %w", e)) }