From 8d12f928ad8b2325f1e7031f88ce6ed9f847d37c Mon Sep 17 00:00:00 2001
From: Knut Ahlers <knut.ahlers@jimdo.com>
Date: Tue, 29 Nov 2016 13:19:31 +0100
Subject: [PATCH] Add line prefixing to see which check logs lines

refs Jimdo/elb-instance-status#5
---
 main.go                 |  4 ++--
 prefixed_logger.go      | 51 +++++++++++++++++++++++++++++++++++++++++
 prefixed_logger_test.go | 34 +++++++++++++++++++++++++++
 3 files changed, 87 insertions(+), 2 deletions(-)
 create mode 100644 prefixed_logger.go
 create mode 100644 prefixed_logger_test.go

diff --git a/main.go b/main.go
index 4a182ec..0604a1d 100644
--- a/main.go
+++ b/main.go
@@ -138,9 +138,9 @@ func executeAndRegisterCheck(ctx context.Context, checkID string) {
 	start := time.Now()
 
 	cmd := exec.Command("/bin/bash", "-e", "-o", "pipefail", "-c", check.Command)
-	cmd.Stderr = os.Stderr
+	cmd.Stderr = newPrefixedLogger(os.Stderr, checkID+":STDERR")
 	if cfg.Verbose {
-		cmd.Stdout = os.Stdout
+		cmd.Stdout = newPrefixedLogger(os.Stderr, checkID+":STDOUT")
 	}
 	err := cmd.Start()
 
diff --git a/prefixed_logger.go b/prefixed_logger.go
new file mode 100644
index 0000000..3bd1f80
--- /dev/null
+++ b/prefixed_logger.go
@@ -0,0 +1,51 @@
+package main
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"sync"
+)
+
+type prefixedLogger struct {
+	channel       string
+	wrappedWriter io.Writer
+
+	buffer     []byte
+	bufferLock sync.Mutex
+}
+
+func newPrefixedLogger(wrappedWriter io.Writer, channel string) *prefixedLogger {
+	return &prefixedLogger{
+		channel:       channel,
+		wrappedWriter: wrappedWriter,
+		buffer:        []byte{},
+	}
+}
+
+func (p *prefixedLogger) dropCR(data []byte) []byte {
+	if len(data) > 0 && data[len(data)-1] == '\r' {
+		return data[0 : len(data)-1]
+	}
+	return data
+}
+
+func (p *prefixedLogger) Write(in []byte) (n int, err error) {
+	p.bufferLock.Lock()
+	defer p.bufferLock.Unlock()
+
+	n = len(in)
+	p.buffer = append(p.buffer, in...)
+
+	for {
+		if i := bytes.IndexByte(p.buffer, '\n'); i >= 0 {
+			// We have a full newline-terminated line.
+			fmt.Fprintf(p.wrappedWriter, "[%s] %s\n", p.channel, string(p.dropCR(p.buffer[0:i])))
+			p.buffer = p.buffer[i+1 : len(p.buffer)]
+		} else {
+			break
+		}
+	}
+
+	return
+}
diff --git a/prefixed_logger_test.go b/prefixed_logger_test.go
new file mode 100644
index 0000000..dfd7840
--- /dev/null
+++ b/prefixed_logger_test.go
@@ -0,0 +1,34 @@
+package main
+
+import (
+	"bytes"
+	"testing"
+)
+
+func TestPrefixedLogger(t *testing.T) {
+	var (
+		buf = bytes.NewBuffer([]byte{})
+		pl  = newPrefixedLogger(buf, "baum")
+		n   int
+		err error
+	)
+
+	n, err = pl.Write([]byte("non-newline terminated string"))
+	if n != 29 || err != nil {
+		t.Fatalf("Write to prefixedLogger had unexpected results: n=29 != %d, err=nil != %s", n, err)
+	}
+
+	if n = len(buf.Bytes()); n != 0 {
+		t.Fatalf("Buffer contains %d characters, should contain 0", n)
+	}
+
+	pl.Write([]byte("now a newline\nand something"))
+	if n = len(buf.Bytes()); n != 50 {
+		t.Fatalf("Buffer contains %d characters, should contain 50", n)
+	}
+
+	pl.Write([]byte(" more to log\nwith multiple\nnewlines\n"))
+	if n = len(buf.Bytes()); n != 120 {
+		t.Fatalf("Buffer contains %d characters, should contain 120", n)
+	}
+}