LogReadCloser interleaves 'by line' for each container

Also, prepend each line with '[ContainerName]'.

probe/kubernetes/client.go

@@ -327,8 +327,8 @@ func (c *client) WalkNodes(f func(*apiv1.Node) error) error {
 }
 
 func (c *client) GetLogs(namespaceID, podID string, containerNames []string) (io.ReadCloser, error) {
-	readClosers := make([]io.ReadCloser, len(containerNames))
-	for i, container := range containerNames {
+	readClosersWithLabel := map[io.ReadCloser]string{}
+	for _, container := range containerNames {
 		req := c.client.CoreV1().Pods(namespaceID).GetLogs(
 			podID,
 			&apiv1.PodLogOptions{
@@ -339,14 +339,15 @@ func (c *client) GetLogs(namespaceID, podID string, containerNames []string) (io
 		)
 		readCloser, err := req.Stream()
 		if err != nil {
-			for _, rc := range readClosers {
+			for rc := range readClosersWithLabel {
 				rc.Close()
 			}
 			return nil, err
 		}
-		readClosers[i] = readCloser
+		readClosersWithLabel[readCloser] = container
 	}
-	return NewLogReadCloser(readClosers...), nil
+
+	return NewLogReadCloser(readClosersWithLabel), nil
 }
 
 func (c *client) DeletePod(namespaceID, podID string) error {

probe/kubernetes/logreadcloser.go

@@ -1,17 +1,16 @@
 package kubernetes
 
 import (
+	"bufio"
 	"bytes"
+	"fmt"
 	"io"
-
-	log "github.com/Sirupsen/logrus"
-)
-
-const (
-	internalBufferSize = 1024
+	"math"
 )
 
 type logReadCloser struct {
+	labels       []string
+	labelLength  int
 	readClosers  []io.ReadCloser
 	eof          []bool
 	buffer       bytes.Buffer
@@ -20,15 +19,27 @@ type logReadCloser struct {
 	eofChannel   chan int
 }
 
-// NewLogReadCloser takes multiple io.ReadCloser and reads where data is available.
-func NewLogReadCloser(readClosers ...io.ReadCloser) io.ReadCloser {
-	stopChannels := make([]chan struct{}, len(readClosers))
-	for i := range readClosers {
+// NewLogReadCloser reads from multiple io.ReadCloser, where data is available,
+// and annotates each line with the reader's label
+func NewLogReadCloser(readClosersWithLabel map[io.ReadCloser]string) io.ReadCloser {
+	stopChannels := make([]chan struct{}, len(readClosersWithLabel))
+	labels := make([]string, len(readClosersWithLabel))
+	readClosers := make([]io.ReadCloser, len(readClosersWithLabel))
+
+	i := 0
+	labelLength := 0
+	for readCloser, label := range readClosersWithLabel {
 		stopChannels[i] = make(chan struct{})
+		readClosers[i] = readCloser
+		labels[i] = label
+		labelLength = int(math.Max(float64(labelLength), float64(len(label))))
+		i++
 	}
 
 	l := logReadCloser{
 		readClosers:  readClosers,
+		labels:       labels,
+		labelLength:  labelLength,
 		dataChannel:  make(chan []byte),
 		stopChannels: stopChannels,
 		eofChannel:   make(chan int),
@@ -112,32 +123,39 @@ func (l *logReadCloser) readInternalBuffer(p []byte) (int, error) {
 	if err == io.EOF && !l.isEOF() {
 		return n, nil
 	}
-
 	return n, err
 }
 
 func (l *logReadCloser) readInput(idx int) {
-	tmpBuffer := make([]byte, internalBufferSize)
+	reader := bufio.NewReader(l.readClosers[idx])
 	for {
-		n, err := l.readClosers[idx].Read(tmpBuffer)
+		line, err := reader.ReadBytes('\n')
 		if err == io.EOF {
-			if n > 0 {
-				l.dataChannel <- tmpBuffer[:n]
+			if len(line) > 0 {
+				l.dataChannel <- l.annotateLine(idx, line)
 			}
 			l.eofChannel <- idx
 			break
 		}
 		if err != nil {
-			log.Errorf("Failed to read: %v", err)
+			// error, exit
 			break
 		}
-		l.dataChannel <- tmpBuffer[:n]
+		l.dataChannel <- l.annotateLine(idx, line)
 	}
 
 	// signal the routine won't write to dataChannel
 	l.stopChannels[idx] <- struct{}{}
 }
 
+func (l *logReadCloser) annotateLine(idx int, line []byte) []byte {
+	// do not annotate if it's the only reader
+	if len(l.labels) == 1 {
+		return line
+	}
+	return []byte(fmt.Sprintf("[%-*s] %v", l.labelLength, l.labels[idx], string(line)))
+}
+
 func (l *logReadCloser) isEOF() bool {
 	for _, e := range l.eof {
 		if !e {

probe/kubernetes/logreadcloser_test.go

@@ -2,23 +2,34 @@ package kubernetes_test
 
 import (
 	"bytes"
+	"fmt"
 	"io"
 	"io/ioutil"
+	"strings"
 	"testing"
 
 	"github.com/weaveworks/scope/probe/kubernetes"
 )
 
 func TestLogReadCloser(t *testing.T) {
-	s0 := []byte("abcdefghijklmnopqrstuvwxyz")
-	s1 := []byte("ABCDEFGHIJKLMNOPQRSTUVWXYZ")
-	s2 := []byte("0123456789012345")
+	data0 := []byte("abcdefghijklmno\npqrstuvwxyz\n")
+	data1 := []byte("ABCDEFGHI\nJKLMNOPQRSTUVWXYZ\n")
+	data2 := []byte("012345678901\n2345\n\n678\n")
 
-	r0 := ioutil.NopCloser(bytes.NewReader(s0))
-	r1 := ioutil.NopCloser(bytes.NewReader(s1))
-	r2 := ioutil.NopCloser(bytes.NewReader(s2))
+	label0 := "zero"
+	label1 := "one"
+	label2 := "two"
+	longestlabelLength := len(label0)
 
-	l := kubernetes.NewLogReadCloser(r0, r1, r2)
+	readClosersWithLabel := map[io.ReadCloser]string{}
+	r0 := ioutil.NopCloser(bytes.NewReader(data0))
+	readClosersWithLabel[r0] = label0
+	r1 := ioutil.NopCloser(bytes.NewReader(data1))
+	readClosersWithLabel[r1] = label1
+	r2 := ioutil.NopCloser(bytes.NewReader(data2))
+	readClosersWithLabel[r2] = label2
+
+	l := kubernetes.NewLogReadCloser(readClosersWithLabel)
 
 	buf := make([]byte, 3000)
 	count := 0
@@ -33,29 +44,23 @@ func TestLogReadCloser(t *testing.T) {
 		count += n
 	}
 
-	total := len(s0) + len(s1) + len(s2)
-	if count != total {
-		t.Errorf("Must read %v characters, but got %v", total, count)
-	}
+	// convert to string for easier comparison
+	result := map[string]int{}
+	lineCounter(result, longestlabelLength, label0, data0)
+	lineCounter(result, longestlabelLength, label1, data1)
+	lineCounter(result, longestlabelLength, label2, data2)
 
-	// check every byte
-	byteCounter := map[byte]int{}
-	byteCount(byteCounter, s0)
-	byteCount(byteCounter, s1)
-	byteCount(byteCounter, s2)
-
-	for i := 0; i < count; i++ {
-		b := buf[i]
-		v, ok := byteCounter[b]
+	str := string(buf[:count])
+	for _, line := range strings.SplitAfter(str, "\n") {
+		v, ok := result[line]
 		if ok {
-			v--
-			byteCounter[b] = v
+			result[line] = v - 1
 		}
 	}
 
-	for b, c := range byteCounter {
-		if c != 0 {
-			t.Errorf("%v should be 0 instead of %v", b, c)
+	for line, v := range result {
+		if v != 0 {
+			t.Errorf("Line %v has not be read from reader", line)
 		}
 	}
 
@@ -65,13 +70,18 @@ func TestLogReadCloser(t *testing.T) {
 	}
 }
 
-func byteCount(accumulator map[byte]int, s []byte) {
-	for _, b := range s {
-		v, ok := accumulator[b]
+func lineCounter(counter map[string]int, pad int, label string, data []byte) {
+	for _, str := range strings.SplitAfter(string(data), "\n") {
+		if len(str) == 0 {
+			// SplitAfter ends with an empty string if the last character is '\n'
+			continue
+		}
+		line := fmt.Sprintf("[%-*s] %v", pad, label, str)
+		v, ok := counter[line]
 		if !ok {
 			v = 0
 		}
 		v++
-		accumulator[b] = v
+		counter[line] = v
 	}
 }