fix: include ioutil helpers

This commit is contained in:
garethgeorge
2024-04-06 14:42:45 -07:00
parent e2be189f9e
commit 88a926b0a3
2 changed files with 88 additions and 2 deletions

86
internal/ioutil/ioutil.go Normal file
View File

@@ -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))
}

View File

@@ -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))
}