Add option to timestamp lines

agent/header_times_streamer.go

@@ -73,18 +73,16 @@ func (h *HeaderTimesStreamer) Scan(line string) {
 	h.scanWaitGroup.Add(1)
 	defer h.scanWaitGroup.Done()
 
-	if h.lineIsHeader(line) {
-		logger.Debug("[HeaderTimesStreamer] Found header %q", line)
+	logger.Debug("[HeaderTimesStreamer] Found header %q", line)
 
-		// Aquire a lock on the times and then add the current time to
-		// our times slice.
-		h.timesMutex.Lock()
-		h.times = append(h.times, time.Now().UTC().Format(time.RFC3339Nano))
-		h.timesMutex.Unlock()
+	// Aquire a lock on the times and then add the current time to
+	// our times slice.
+	h.timesMutex.Lock()
+	h.times = append(h.times, time.Now().UTC().Format(time.RFC3339Nano))
+	h.timesMutex.Unlock()
 
-		// Add the time to the wait group
-		h.uploadWaitGroup.Add(1)
-	}
+	// Add the time to the wait group
+	h.uploadWaitGroup.Add(1)
 }
 
 func (h *HeaderTimesStreamer) Upload() {
@@ -136,7 +134,7 @@ func (h *HeaderTimesStreamer) Stop() {
 	h.streamingMutex.Unlock()
 }
 
-func (h *HeaderTimesStreamer) lineIsHeader(line string) bool {
+func (h *HeaderTimesStreamer) LineIsHeader(line string) bool {
 	// Make sure all ANSI colors are removed from the string before we
 	// check to see if it's a header (sometimes a color escape sequence may
 	// be the first thing on the line, which will cause the regex to ignore

agent/job_runner.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"strconv"
 	"sync"
 	"time"
 
@@ -59,13 +60,20 @@ func (r JobRunner) Create() (runner *JobRunner, err error) {
 	// the Buildkite Agent API
 	runner.logStreamer = LogStreamer{MaxChunkSizeBytes: r.Job.ChunksMaxSizeBytes, Callback: r.onUploadChunk}.New()
 
+	timestamp, err := strconv.ParseBool(os.Getenv("BUILDKITE_TIMESTAMP_LINES"))
+	if err != nil {
+		timestamp = false
+	}
+
 	// The process that will run the bootstrap script
 	runner.process = process.Process{
-		Script:        r.AgentConfiguration.BootstrapScript,
-		Env:           r.createEnvironment(),
-		PTY:           r.AgentConfiguration.RunInPty,
-		StartCallback: r.onProcessStartCallback,
-		LineCallback:  runner.headerTimesStreamer.Scan,
+		Script:             r.AgentConfiguration.BootstrapScript,
+		Env:                r.createEnvironment(),
+		PTY:                r.AgentConfiguration.RunInPty,
+		Timestamp:          timestamp,
+		StartCallback:      r.onProcessStartCallback,
+		LineCallback:       runner.headerTimesStreamer.Scan,
+		LineCallbackFilter: runner.headerTimesStreamer.LineIsHeader,
 	}.Create()
 
 	return

process/process.go

@@ -23,6 +23,7 @@ import (
 type Process struct {
 	Pid        int
 	PTY        bool
+	Timestamp  bool
 	Script     string
 	Env        []string
 	ExitStatus string
@@ -35,8 +36,9 @@ type Process struct {
 	StartCallback func()
 
 	// For every line in the process output, this callback will be called
-	// with the contents of the line
-	LineCallback func(string)
+	// with the contents of the line if its filter returns true.
+	LineCallback       func(string)
+	LineCallbackFilter func(string) bool
 
 	// Running is stored as an int32 so we can use atomic operations to
 	// set/get it (it's accessed by multiple goroutines)
@@ -67,7 +69,12 @@ func (p *Process) Start() error {
 
 	lineReaderPipe, lineWriterPipe := io.Pipe()
 
-	multiWriter := io.MultiWriter(&p.buffer, lineWriterPipe)
+	var multiWriter io.Writer
+	if p.Timestamp {
+		multiWriter = io.MultiWriter(lineWriterPipe)
+	} else {
+		multiWriter = io.MultiWriter(&p.buffer, lineWriterPipe)
+	}
 
 	logger.Info("Starting to run script: %s", p.command.Path)
 
@@ -171,11 +178,37 @@ func (p *Process) Start() error {
 				}
 			}
 
-			lineCallbackWaitGroup.Add(1)
-			go func(line string) {
-				defer lineCallbackWaitGroup.Done()
-				p.LineCallback(line)
-			}(string(line))
+			// If we're timestamping this main thread will take
+			// the hit of running the regex so we can build up
+			// the timestamped buffer without breaking headers,
+			// otherwise we let the goroutines take the perf hit.
+
+			checkedForCallback := false
+			lineHasCallback := false
+			lineString := string(line)
+
+			// Create the prefixed buffer
+			if p.Timestamp {
+				lineHasCallback = p.LineCallbackFilter(lineString)
+				checkedForCallback = true
+				if lineHasCallback {
+					// Don't timestamp special lines (e.g. header)
+					p.buffer.WriteString(fmt.Sprintf("%s\n", line))
+				} else {
+					currentTime := time.Now().UTC().Format(time.RFC3339)
+					p.buffer.WriteString(fmt.Sprintf("[%s] %s\n", currentTime, line))
+				}
+			}
+
+			if lineHasCallback || !checkedForCallback {
+				lineCallbackWaitGroup.Add(1)
+				go func(line string) {
+					defer lineCallbackWaitGroup.Done()
+					if (checkedForCallback && lineHasCallback) || p.LineCallbackFilter(lineString) {
+						p.LineCallback(line)
+					}
+				}(lineString)
+			}
 		}
 
 		// We need to make sure all the line callbacks have finish before