Present service logs with colored service prefix

This reproduce docker-compose behaviour to report logs with prefix
also moves log formating out from sdk.go

Signed-off-by: Nicolas De Loof <[email protected]>

ecs/pkg/amazon/api.go

@@ -2,7 +2,7 @@ package amazon
 
 import "context"
 
-//go:generate mockgen -destination=./mock/api.go -package=mock . API
+//go:generate mockgen -destination=./api_mock.go -self_package "github.com/docker/ecs-plugin/pkg/amazon" -package=amazon . API
 
 type API interface {
 	downAPI

ecs/pkg/amazon/mock/api.go → ecs/pkg/amazon/api_mock.go

@@ -1,8 +1,8 @@
 // Code generated by MockGen. DO NOT EDIT.
 // Source: github.com/docker/ecs-plugin/pkg/amazon (interfaces: API)
 
-// Package mock is a generated GoMock package.
-package mock
+// Package amazon is a generated GoMock package.
+package amazon
 
 import (
 	context "context"
@@ -153,17 +153,17 @@ func (mr *MockAPIMockRecorder) GetDefaultVPC(arg0 interface{}) *gomock.Call {
 }
 
 // GetLogs mocks base method
-func (m *MockAPI) GetLogs(arg0 context.Context, arg1 string) error {
+func (m *MockAPI) GetLogs(arg0 context.Context, arg1 string, arg2 LogConsumer) error {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetLogs", arg0, arg1)
+	ret := m.ctrl.Call(m, "GetLogs", arg0, arg1, arg2)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
 // GetLogs indicates an expected call of GetLogs
-func (mr *MockAPIMockRecorder) GetLogs(arg0, arg1 interface{}) *gomock.Call {
+func (mr *MockAPIMockRecorder) GetLogs(arg0, arg1, arg2 interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetLogs", reflect.TypeOf((*MockAPI)(nil).GetLogs), arg0, arg1)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetLogs", reflect.TypeOf((*MockAPI)(nil).GetLogs), arg0, arg1, arg2)
 }
 
 // GetNetworkInterfaces mocks base method

ecs/pkg/amazon/down_test.go

@@ -4,14 +4,13 @@ import (
 	"context"
 	"testing"
 
-	"github.com/docker/ecs-plugin/pkg/amazon/mock"
 	"github.com/golang/mock/gomock"
 )
 
 func TestDownDontDeleteCluster(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
-	m := mock.NewMockAPI(ctrl)
+	m := NewMockAPI(ctrl)
 	c := &client{
 		Cluster: "test_cluster",
 		Region:  "region",
@@ -30,7 +29,7 @@ func TestDownDontDeleteCluster(t *testing.T) {
 func TestDownDeleteCluster(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
-	m := mock.NewMockAPI(ctrl)
+	m := NewMockAPI(ctrl)
 	c := &client{
 		Cluster: "test_cluster",
 		Region:  "region",

ecs/pkg/amazon/logs.go

@@ -2,12 +2,62 @@ package amazon
 
 import (
 	"context"
+	"fmt"
+	"os"
+	"os/signal"
+	"strconv"
+	"strings"
+
+	"github.com/docker/ecs-plugin/pkg/console"
 )
 
 func (c *client) ComposeLogs(ctx context.Context, projectName string) error {
-	return c.api.GetLogs(ctx, projectName)
+	err := c.api.GetLogs(ctx, projectName, &logConsumer{
+		colors: map[string]console.ColorFunc{},
+		width:  0,
+	})
+	if err != nil {
+		return err
+	}
+
+	signalChan := make(chan os.Signal, 1)
+	signal.Notify(signalChan, os.Interrupt)
+	<-signalChan
+	return nil
+}
+
+type logConsumer struct {
+	colors map[string]console.ColorFunc
+	width  int
+}
+
+func (l *logConsumer) Log(service, container, message string) {
+	cf, ok := l.colors[service]
+	if !ok {
+		cf = <-console.Rainbow
+		l.colors[service] = cf
+		l.computeWidth()
+	}
+	prefix := fmt.Sprintf("%-"+strconv.Itoa(l.width)+"s |", service)
+	for _, line := range strings.Split(message, "\n") {
+		fmt.Printf("%s %s\n", cf(prefix), line)
+	}
+}
+
+func (l *logConsumer) computeWidth() {
+	width := 0
+	for n := range l.colors {
+		if len(n) > width {
+			width = len(n)
+		}
+	}
+	l.width = width + 3
+}
+
+type LogConsumer interface {
+	Log(service, container, message string)
 }
 
 type logsAPI interface {
-	GetLogs(ctx context.Context, name string) error
+	GetLogs(ctx context.Context, name string, consumer LogConsumer) error
 }

ecs/pkg/amazon/sdk.go

@@ -3,6 +3,7 @@ package amazon
 import (
 	"context"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -309,7 +310,7 @@ func (s sdk) DeleteSecret(ctx context.Context, id string, recover bool) error {
 	return err
 }
 
-func (s sdk) GetLogs(ctx context.Context, name string) error {
+func (s sdk) GetLogs(ctx context.Context, name string, consumer LogConsumer) error {
 	logGroup := fmt.Sprintf("/docker-compose/%s", name)
 	var startTime int64
 	for {
@@ -331,7 +332,8 @@ func (s sdk) GetLogs(ctx context.Context, name string) error {
 			}
 
 			for _, event := range events.Events {
-				fmt.Println(*event.Message)
+				p := strings.Split(*event.LogStreamName, "/")
+				consumer.Log(p[1], p[2], *event.Message)
 				startTime = *event.IngestionTime
 			}
 		}

ecs/pkg/console/colors.go

@@ -0,0 +1,62 @@
+package console
+
+import (
+	"strconv"
+)
+
+var NAMES = []string{
+	"grey",
+	"red",
+	"green",
+	"yellow",
+	"blue",
+	"magenta",
+	"cyan",
+	"white",
+}
+
+var COLORS map[string]ColorFunc
+
+// ColorFunc use ANSI codes to render colored text on console
+type ColorFunc func(s string) string
+
+var Monochrome = func(s string) string {
+	return s
+}
+
+func makeColorFunc(code string) ColorFunc {
+	return func(s string) string {
+		return ansiColor(code, s)
+	}
+}
+
+var Rainbow = make(chan ColorFunc)
+
+func init() {
+	COLORS = map[string]ColorFunc{}
+	for i, name := range NAMES {
+		COLORS[name] = makeColorFunc(strconv.Itoa(30 + i))
+		COLORS["intense_"+name] = makeColorFunc(strconv.Itoa(30+i) + ";1")
+	}
+
+	go func() {
+		i := 0
+		rainbow := []ColorFunc{
+			COLORS["cyan"],
+			COLORS["yellow"],
+			COLORS["green"],
+			COLORS["magenta"],
+			COLORS["blue"],
+			COLORS["intense_cyan"],
+			COLORS["intense_yellow"],
+			COLORS["intense_green"],
+			COLORS["intense_magenta"],
+			COLORS["intense_blue"],
+		}
+
+		for {
+			Rainbow <- rainbow[i]
+			i = (i + 1) % len(rainbow)
+		}
+	}()
+}