| 1 | package runs |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "errors" |
| 6 | "io" |
| 7 | "sync" |
| 8 | "time" |
| 9 | |
| 10 | "github.com/carlosframework/platform/internal/release" |
| 11 | ) |
| 12 | |
| 13 | // errBadStream is what Close reports when a writer was made for a stream |
| 14 | // that is neither StreamOut nor StreamErr. ChunkKey trusts its stream, so the |
| 15 | // chunker refuses anything else before it reaches a key. |
| 16 | var errBadStream = errors.New("runs: chunker stream is not out or err") |
| 17 | |
| 18 | // Chunker turns a run's two streams into chunks (run spec §5.5): |
| 19 | // out/<seq>-<o|e>, one seq for both streams in the order the bytes were |
| 20 | // read. A chunk holds one stream's bytes: a write on the other stream cuts |
| 21 | // it, so the follower interleaves exactly as the agent read. A chunk is |
| 22 | // written every second, or as soon as it holds 256 KB. Past 64 MB in total |
| 23 | // nothing more is written and Close says truncated; writes never fail, so |
| 24 | // the command is never failed by a refused pipe. Each chunk is a |
| 25 | // create-only Put retried until it lands or the context ends, and the |
| 26 | // retries run under the lock. A stalled bucket therefore does block the |
| 27 | // command's pipes, and the command with them once its pipe buffers fill, |
| 28 | // until the Put lands or the context ends. That is deliberate: it keeps |
| 29 | // seqs in order and stops output piling up in memory. |
| 30 | type Chunker struct { |
| 31 | ctx context.Context |
| 32 | s release.Store |
| 33 | acct, app, id string |
| 34 | every time.Duration |
| 35 | max int64 |
| 36 | |
| 37 | mu sync.Mutex |
| 38 | stream string |
| 39 | pending []byte |
| 40 | seq int64 |
| 41 | total int64 |
| 42 | truncated bool |
| 43 | err error |
| 44 | closed bool |
| 45 | |
| 46 | stop chan struct{} |
| 47 | done chan struct{} |
| 48 | closeOnce sync.Once |
| 49 | } |
| 50 | |
| 51 | // NewChunker starts a chunker for run id with the spec's cadence and cap. |
| 52 | // An account, app or run id the store helpers would refuse (ErrBadApp, |
| 53 | // ErrBadID) gives a chunker that writes nothing and whose Close returns that |
| 54 | // error, so it can never write outside runs/<id>/out/. |
| 55 | func NewChunker(ctx context.Context, s release.Store, acct, app, id string) *Chunker { |
| 56 | return newChunker(ctx, s, acct, app, id, ChunkEvery, MaxOutputBytes) |
| 57 | } |
| 58 | |
| 59 | func newChunker(ctx context.Context, s release.Store, acct, app, id string, every time.Duration, max int64) *Chunker { |
| 60 | c := &Chunker{ctx: ctx, s: s, acct: acct, app: app, id: id, every: every, max: max, |
| 61 | stop: make(chan struct{}), done: make(chan struct{})} |
| 62 | if err := checkRun(acct, app, id); err != nil { |
| 63 | c.err, c.truncated = err, true |
| 64 | close(c.done) |
| 65 | return c |
| 66 | } |
| 67 | go c.tick() |
| 68 | return c |
| 69 | } |
| 70 | |
| 71 | func (c *Chunker) tick() { |
| 72 | defer close(c.done) |
| 73 | t := time.NewTicker(c.every) |
| 74 | defer t.Stop() |
| 75 | for { |
| 76 | select { |
| 77 | case <-c.stop: |
| 78 | return |
| 79 | case <-t.C: |
| 80 | c.mu.Lock() |
| 81 | c.flushLocked() |
| 82 | c.mu.Unlock() |
| 83 | } |
| 84 | } |
| 85 | } |
| 86 | |
| 87 | type chunkStream struct { |
| 88 | c *Chunker |
| 89 | stream string |
| 90 | } |
| 91 | |
| 92 | func (w chunkStream) Write(p []byte) (int, error) { |
| 93 | w.c.write(w.stream, p) |
| 94 | return len(p), nil |
| 95 | } |
| 96 | |
| 97 | // Stream is the writer for one stream, StreamOut or StreamErr. Its Write |
| 98 | // never fails and never returns short. |
| 99 | func (c *Chunker) Stream(stream string) io.Writer { return chunkStream{c: c, stream: stream} } |
| 100 | |
| 101 | func (c *Chunker) write(stream string, p []byte) { |
| 102 | c.mu.Lock() |
| 103 | defer c.mu.Unlock() |
| 104 | if c.closed { |
| 105 | return |
| 106 | } |
| 107 | if stream != StreamOut && stream != StreamErr { |
| 108 | if c.err == nil { |
| 109 | c.err, c.truncated, c.pending = errBadStream, true, nil |
| 110 | } |
| 111 | return |
| 112 | } |
| 113 | if c.err != nil { |
| 114 | c.truncated = c.truncated || len(p) > 0 |
| 115 | return |
| 116 | } |
| 117 | if room := c.max - c.total; int64(len(p)) > room { |
| 118 | p = p[:room] |
| 119 | c.truncated = true |
| 120 | } |
| 121 | for len(p) > 0 { |
| 122 | if c.stream != stream && len(c.pending) > 0 { |
| 123 | c.flushLocked() |
| 124 | if c.err != nil { |
| 125 | return |
| 126 | } |
| 127 | } |
| 128 | c.stream = stream |
| 129 | n := min(len(p), ChunkBytes-len(c.pending)) |
| 130 | c.pending = append(c.pending, p[:n]...) |
| 131 | c.total += int64(n) |
| 132 | p = p[n:] |
| 133 | if len(c.pending) >= ChunkBytes { |
| 134 | c.flushLocked() |
| 135 | if c.err != nil { |
| 136 | c.truncated = true |
| 137 | return |
| 138 | } |
| 139 | } |
| 140 | } |
| 141 | } |
| 142 | |
| 143 | // flushLocked writes the pending chunk under the next seq. A Put that |
| 144 | // answers ErrCASConflict was written by an earlier attempt whose answer was |
| 145 | // lost; the seq is spent either way. |
| 146 | func (c *Chunker) flushLocked() { |
| 147 | if len(c.pending) == 0 || c.err != nil { |
| 148 | return |
| 149 | } |
| 150 | seq := c.seq + 1 |
| 151 | key := ChunkKey(c.acct, c.app, c.id, seq, c.stream) |
| 152 | backoff := 100 * time.Millisecond |
| 153 | for { |
| 154 | _, err := c.s.Put(c.ctx, key, c.pending, "") |
| 155 | if err == nil || errors.Is(err, release.ErrCASConflict) { |
| 156 | break |
| 157 | } |
| 158 | select { |
| 159 | case <-c.ctx.Done(): |
| 160 | c.err, c.truncated, c.pending = err, true, nil |
| 161 | return |
| 162 | case <-time.After(backoff): |
| 163 | } |
| 164 | backoff = min(backoff*2, 5*time.Second) |
| 165 | } |
| 166 | c.seq, c.pending = seq, nil |
| 167 | } |
| 168 | |
| 169 | // Close stops the ticker and writes the last chunk. Call it after both |
| 170 | // streams reached their end (CmdHandle.Stop drains them); writes after Close |
| 171 | // are dropped. It does not return until every chunk is written or chunking |
| 172 | // gave up, so a caller may write result.json once it returns. It returns the |
| 173 | // last seq written (0: no output), whether output was cut off, and the |
| 174 | // failure that ended chunking, if one did. It is safe to call twice. |
| 175 | func (c *Chunker) Close() (int64, bool, error) { |
| 176 | c.closeOnce.Do(func() { |
| 177 | close(c.stop) |
| 178 | <-c.done |
| 179 | }) |
| 180 | c.mu.Lock() |
| 181 | defer c.mu.Unlock() |
| 182 | c.flushLocked() |
| 183 | c.closed = true |
| 184 | return c.seq, c.truncated, c.err |
| 185 | } |
| 186 | |
amadan