diff --git a/probe/kubernetes/client.go b/probe/kubernetes/client.go index 7e3f2f156..da53ceccf 100644 --- a/probe/kubernetes/client.go +++ b/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 { diff --git a/probe/kubernetes/logreadcloser.go b/probe/kubernetes/logreadcloser.go index 865ce73f4..d2e780bd4 100644 --- a/probe/kubernetes/logreadcloser.go +++ b/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 { diff --git a/probe/kubernetes/logreadcloser_test.go b/probe/kubernetes/logreadcloser_test.go index 55801376b..66806cfe8 100644 --- a/probe/kubernetes/logreadcloser_test.go +++ b/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 } }