fix(shell): keep partial end marker across reads so jobs don't hang

pkg/shell/output_buffer.go

@@ -6,6 +6,7 @@ import (
 	"math"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 	"unicode/utf8"
 
@@ -38,7 +39,7 @@ type OutputBuffer struct {
 	Consumer     func(string)
 	bytes        []byte
 	mu           sync.Mutex
-	done         bool
+	done         atomic.Bool
 	lastAppend   *time.Time
 	flushTimeout time.Duration
 }
@@ -73,14 +74,16 @@ func (b *OutputBuffer) Append(bytes []byte) {
 }
 
 func (b *OutputBuffer) IsEmpty() bool {
+	b.mu.Lock()
+	defer b.mu.Unlock()
 	return len(b.bytes) == 0
 }
 
 func (b *OutputBuffer) Flush() {
 	backoffStrategy := b.exponentialBackoff()
 
 	for {
-		if b.done {
+		if b.done.Load() {
 			log.Debugf("The output buffer was closed - stopping")
 			break
 		}
@@ -157,7 +160,7 @@ func (b *OutputBuffer) flush() {
 	 *       is because we are cutting it here to fit it into the chunk.
 	 */
 
-	if !b.done || (b.done && cutLength == chunkSize) {
+	if !b.done.Load() || (b.done.Load() && cutLength == chunkSize) {
 		for i := 0; i < 4; i++ {
 			if utf8.Valid(b.bytes[0:cutLength]) {
 				break
@@ -227,15 +230,15 @@ func (b *OutputBuffer) timeSinceLastAppend() time.Duration {
 func (b *OutputBuffer) chunkSize() int {
 	// If the output buffer was already closed, we should
 	// use bigger chunks to make sure we can flush everything left in time.
-	if b.done {
+	if b.done.Load() {
 		return OutputBufferDefaultCutLength * 10
 	}
 
 	return OutputBufferDefaultCutLength
 }
 
 func (b *OutputBuffer) Close() error {
-	b.done = true
+	b.done.Store(true)
 
 	ctx, cancelFunc := context.WithTimeout(context.Background(), b.flushTimeout)
 	defer cancelFunc()

pkg/shell/process.go

@@ -52,11 +52,16 @@ type Process struct {
 	Pid               int
 	startMark         string
 	endMark           string
+	commandStartRegex *regexp.Regexp
 	commandEndRegex   *regexp.Regexp
 	inputBuffer       []byte
 	outputBuffer      *OutputBuffer
 	SysProcAttr       *syscall.SysProcAttr
 	UseBase64Encoding bool
+
+	// readSource, when non-nil, replaces reading from the shell TTY. Tests set
+	// it to drive scan() one read at a time with deterministic read boundaries.
+	readSource func() ([]byte, error)
 }
 
 func randomMagicMark() string {
@@ -66,7 +71,13 @@ func randomMagicMark() string {
 func NewProcess(config Config) *Process {
 	startMark := randomMagicMark() + "-start"
 	endMark := randomMagicMark() + "-end"
-	commandEndRegex := regexp.MustCompile(endMark + " " + `(\d+)` + "[\r\n]+")
+	commandStartRegex := regexp.MustCompile(startMark + `\r?\n`)
+
+	// The end marker is emitted as `\001 <endMark> <exit code>\n` (see
+	// constructShellInstruction). Matching it *including* the \001 prefix lets
+	// scan locate the complete marker anywhere in the buffer and flush anything
+	// before it - even a stray \001 - as ordinary output rather than dropping it.
+	commandEndRegex := regexp.MustCompile("\001 " + endMark + " " + `(\d+)` + "[\r\n]+")
 	outputBuffer, _ := NewOutputBuffer(config.OnOutput)
 
 	return &Process{
@@ -76,6 +87,7 @@ func NewProcess(config Config) *Process {
 		ExitCode:          1,
 		startMark:         startMark,
 		endMark:           endMark,
+		commandStartRegex: commandStartRegex,
 		commandEndRegex:   commandEndRegex,
 		outputBuffer:      outputBuffer,
 		UseBase64Encoding: config.UseBase64Encoding,
@@ -94,6 +106,40 @@ func (p *Process) flushInputAll() {
 	p.flushInputBufferTill(len(p.inputBuffer))
 }
 
+// flushInputBufferKeepingMarkerTail is used in buffering mode when the buffer
+// holds a \001 (a possible end-marker start) but no complete end marker has
+// matched yet. It flushes buffered output but retains a trailing slice that
+// could be a not-yet-complete end marker: everything from the LAST \001 onward.
+//
+// The leading \001 that put us in buffering mode may be a stray byte from the
+// command's own output (e.g. binary data printed to the log), with the real
+// marker arriving right after it. Flushing the whole buffer here would leak a
+// partially-received real marker into the job output and the marker would
+// never be recognized (the original hang). Keeping the tail from the last
+// \001 preserves the partial marker across reads.
+//
+// The tail is capped at maxEndMarkerLen: a real marker is never longer than
+// that, so if the tail already exceeds it the leading \001 cannot be a partial
+// marker and the whole buffer is safe to flush. This keeps memory bounded even
+// under a pathological stream of \001 bytes.
+func (p *Process) flushInputBufferKeepingMarkerTail() {
+	lastMark := strings.LastIndex(string(p.inputBuffer), "\001")
+
+	if lastMark < 0 || len(p.inputBuffer)-lastMark > p.maxEndMarkerLen() {
+		p.flushInputAll()
+		return
+	}
+
+	p.flushInputBufferTill(lastMark)
+}
+
+// maxEndMarkerLen is the largest number of bytes a complete end marker can
+// occupy: "\001 " + endMark + " " + exit code digits + line terminator. The
+// +10 slack covers the prefix, separators, exit code, and CR/LF.
+func (p *Process) maxEndMarkerLen() int {
+	return len(p.endMark) + 10
+}
+
 func (p *Process) flushInputBufferTill(index int) {
 	if index == 0 {
 		return
@@ -387,6 +433,12 @@ func (p *Process) readNonPTY(reader *io.PipeReader, done chan bool) {
 
 // Read state from shell into the inputBuffer
 func (p *Process) read() error {
+	if p.readSource != nil {
+		data, err := p.readSource()
+		p.inputBuffer = append(p.inputBuffer, data...)
+		return err
+	}
+
 	buffer := make([]byte, p.readBufferSize())
 
 	log.Debug("Reading started")
@@ -415,20 +467,22 @@ func (p *Process) waitForStartMarker() error {
 		}
 
 		//
-		// If the inputBuffer has a start marker, the wait is done
+		// If the inputBuffer has a start marker, the wait is done.
 		//
-		index := strings.Index(string(p.inputBuffer), p.startMark+"\r\n")
-
-		if index >= 0 {
+		// The marker's line terminator may be \r\n or a bare \n (a previous
+		// command can leave the PTY in -onlcr mode), so we match \r?\n rather
+		// than a literal \r\n.
+		//
+		if loc := p.commandStartRegex.FindStringIndex(string(p.inputBuffer)); loc != nil {
 			//
-			// Cut everything from the buffer before the marker
+			// Cut everything from the buffer up to and including the marker.
 			// Example:
 			//
 			// buffer before: some test <***marker**> rest of the test
 			// buffer after :  rest of the test
 			//
 
-			p.inputBuffer = p.inputBuffer[index+len(p.startMark)+2 : len(p.inputBuffer)]
+			p.inputBuffer = p.inputBuffer[loc[1]:]
 
 			break
 		}
@@ -454,29 +508,40 @@ func (p *Process) scan() error {
 	exitCode := ""
 
 	for {
-		if index := p.endMarkerHeaderIndex(); index >= 0 {
-			if index > 0 {
-				// publish everything until the end mark
-				p.flushInputBufferTill(index)
-			}
-
-			log.Debug("Start of end marker detected, entering buffering mode.")
-
-			if match := p.commandEndRegex.FindStringSubmatch(string(p.inputBuffer)); len(match) == 2 {
-				log.Debug("End marker detected. Exit code: ", match[1])
+		//
+		// A complete end marker (with its \001 prefix) ends the command.
+		// Matching the whole marker means anything before it - including a
+		// stray \001 emitted by the command's own output - is flushed as
+		// ordinary output instead of being discarded.
+		//
+		buffered := string(p.inputBuffer)
+		if loc := p.commandEndRegex.FindStringSubmatchIndex(buffered); loc != nil {
+			exitCode = buffered[loc[2]:loc[3]]
+			log.Debug("End marker detected. Exit code: ", exitCode)
 
-				exitCode = match[1]
-				break
-			}
+			// Everything before the marker is genuine command output.
+			p.flushInputBufferTill(loc[0])
+			break
+		}
 
+		if index := p.endMarkerHeaderIndex(); index >= 0 {
 			//
-			// The buffer is much longer than the end mark, at least by 10
-			// characters.
+			// A \001 is present but no complete marker has arrived yet - it may
+			// still be streaming in across reads. Flush everything before the
+			// first \001 as ordinary output and stay in buffering mode.
+			//
+			log.Debug("Start of end marker detected, entering buffering mode.")
+			p.flushInputBufferTill(index)
+
 			//
-			// If it is not matching the full end mark, it is safe to dump.
+			// The buffer has grown past the longest possible marker without
+			// matching, so the leading \001 was a stray byte. Flush it, keeping
+			// only the tail from the last \001, which may be a real marker that
+			// split across reads - flushing it here would lose the marker and
+			// hang the job.
 			//
-			if len(p.inputBuffer) >= len(p.endMark)+10 {
-				p.flushInputAll()
+			if len(p.inputBuffer) >= p.maxEndMarkerLen() {
+				p.flushInputBufferKeepingMarkerTail()
 			}
 		} else {
 			p.flushInputAll()