From 4fcaa296145faa27eaa4d076417f92131c98a816 Mon Sep 17 00:00:00 2001
From: Nicolas De Loof <nicolas.deloof@gmail.com>
Date: Mon, 15 Feb 2021 11:14:09 +0100
Subject: [PATCH] fix color assignment on status messages

Signed-off-by: Nicolas De Loof <nicolas.deloof@gmail.com>
---
 api/compose/api.go      |  7 ++++---
 cli/cmd/compose/up.go   |  6 +++---
 cli/formatter/logs.go   | 44 ++++++++++++++++++++---------------------
 local/compose/attach.go |  5 +++--
 local/compose/logs.go   |  7 +++++--
 local/compose/start.go  |  3 ++-
 6 files changed, 38 insertions(+), 34 deletions(-)

diff --git a/api/compose/api.go b/api/compose/api.go
index 5efeb3700..5847a0cf5 100644
--- a/api/compose/api.go
+++ b/api/compose/api.go
@@ -182,9 +182,9 @@ type Stack struct {
 
 // LogConsumer is a callback to process log messages from services
 type LogConsumer interface {
-	Log(service, container, message string)
-	Status(service, container, msg string)
-	Register(service string, source string)
+	Log(name, container, message string)
+	Status(name, container, msg string)
+	Register(name string, source string)
 }
 
 // ContainerEventListener is a callback to process ContainerEvent from services
@@ -195,6 +195,7 @@ type ContainerEvent struct {
 	Type     int
 	Source   string
 	Service  string
+	Name     string
 	Line     string
 	ExitCode int
 }
diff --git a/cli/cmd/compose/up.go b/cli/cmd/compose/up.go
index 4da5e1bc1..792a81fbe 100644
--- a/cli/cmd/compose/up.go
+++ b/cli/cmd/compose/up.go
@@ -320,11 +320,11 @@ func (p printer) run(ctx context.Context, cascadeStop bool, exitCodeFrom string,
 		event := <-p.queue
 		switch event.Type {
 		case compose.ContainerEventAttach:
-			consumer.Register(event.Service, event.Source)
+			consumer.Register(event.Name, event.Source)
 			count++
 		case compose.ContainerEventExit:
 			if !aborting {
-				consumer.Status(event.Service, event.Source, fmt.Sprintf("exited with code %d", event.ExitCode))
+				consumer.Status(event.Name, event.Source, fmt.Sprintf("exited with code %d", event.ExitCode))
 			}
 			if cascadeStop {
 				if !aborting {
@@ -347,7 +347,7 @@ func (p printer) run(ctx context.Context, cascadeStop bool, exitCodeFrom string,
 			}
 		case compose.ContainerEventLog:
 			if !aborting {
-				consumer.Log(event.Service, event.Source, event.Line)
+				consumer.Log(event.Name, event.Source, event.Line)
 			}
 		}
 	}
diff --git a/cli/formatter/logs.go b/cli/formatter/logs.go
index c5e2e0ab5..3d792123b 100644
--- a/cli/formatter/logs.go
+++ b/cli/formatter/logs.go
@@ -38,21 +38,20 @@ func NewLogConsumer(ctx context.Context, w io.Writer, color bool, prefix bool) c
 	}
 }
 
-func (l *logConsumer) Register(service string, source string) {
-	l.register(service, source)
+func (l *logConsumer) Register(name string, id string) {
+	l.register(name, id)
 }
 
-func (l *logConsumer) register(service string, source string) *presenter {
+func (l *logConsumer) register(name string, id string) *presenter {
 	cf := monochrome
 	if l.color {
 		cf = <-loop
 	}
 	p := &presenter{
-		colors:    cf,
-		service:   service,
-		container: source,
+		colors: cf,
+		name:   name,
 	}
-	l.presenters[source] = p
+	l.presenters[id] = p
 	if l.prefix {
 		l.computeWidth()
 		for _, p := range l.presenters {
@@ -62,34 +61,34 @@ func (l *logConsumer) register(service string, source string) *presenter {
 	return p
 }
 
-// Log formats a log message as received from service/container
-func (l *logConsumer) Log(service, container, message string) {
+// Log formats a log message as received from name/container
+func (l *logConsumer) Log(name, id, message string) {
 	if l.ctx.Err() != nil {
 		return
 	}
-	p, ok := l.presenters[container]
+	p, ok := l.presenters[id]
 	if !ok { // should have been registered, but ¯\_(ツ)_/¯
-		p = l.register(service, container)
+		p = l.register(name, id)
 	}
 	for _, line := range strings.Split(message, "\n") {
 		fmt.Fprintf(l.writer, "%s %s\n", p.prefix, line) // nolint:errcheck
 	}
 }
 
-func (l *logConsumer) Status(service, container, msg string) {
-	p, ok := l.presenters[container]
+func (l *logConsumer) Status(name, id, msg string) {
+	p, ok := l.presenters[id]
 	if !ok {
-		p = l.register(service, container)
+		p = l.register(name, id)
 	}
-	s := p.colors(fmt.Sprintf("%s %s\n", container, msg))
+	s := p.colors(fmt.Sprintf("%s %s\n", name, msg))
 	l.writer.Write([]byte(s)) // nolint:errcheck
 }
 
 func (l *logConsumer) computeWidth() {
 	width := 0
-	for n := range l.presenters {
-		if len(n) > width {
-			width = len(n)
+	for _, p := range l.presenters {
+		if len(p.name) > width {
+			width = len(p.name)
 		}
 	}
 	l.width = width + 1
@@ -106,12 +105,11 @@ type logConsumer struct {
 }
 
 type presenter struct {
-	colors    colorFunc
-	service   string
-	container string
-	prefix    string
+	colors colorFunc
+	name   string
+	prefix string
 }
 
 func (p *presenter) setPrefix(width int) {
-	p.prefix = p.colors(fmt.Sprintf("%-"+strconv.Itoa(width)+"s |", p.container))
+	p.prefix = p.colors(fmt.Sprintf("%-"+strconv.Itoa(width)+"s |", p.name))
 }
diff --git a/local/compose/attach.go b/local/compose/attach.go
index 3dd1c09a5..a07f2e5ee 100644
--- a/local/compose/attach.go
+++ b/local/compose/attach.go
@@ -48,7 +48,8 @@ func (s *composeService) attach(ctx context.Context, project *types.Project, con
 	for _, container := range containers {
 		consumer(compose.ContainerEvent{
 			Type:    compose.ContainerEventAttach,
-			Source:  getContainerNameWithoutProject(container),
+			Source:  container.ID,
+			Name:    getContainerNameWithoutProject(container),
 			Service: container.Labels[serviceLabel],
 		})
 		err := s.attachContainer(ctx, container, consumer, project)
@@ -61,7 +62,7 @@ func (s *composeService) attach(ctx context.Context, project *types.Project, con
 
 func (s *composeService) attachContainer(ctx context.Context, container moby.Container, consumer compose.ContainerEventListener, project *types.Project) error {
 	serviceName := container.Labels[serviceLabel]
-	w := getWriter(serviceName, getContainerNameWithoutProject(container), consumer)
+	w := getWriter(getContainerNameWithoutProject(container), serviceName, container.ID, consumer)
 
 	service, err := project.GetService(serviceName)
 	if err != nil {
diff --git a/local/compose/logs.go b/local/compose/logs.go
index 6835d9681..1949ad609 100644
--- a/local/compose/logs.go
+++ b/local/compose/logs.go
@@ -88,14 +88,16 @@ func (s *composeService) Logs(ctx context.Context, projectName string, consumer
 }
 
 type splitBuffer struct {
-	service   string
+	name      string
 	container string
 	consumer  compose.ContainerEventListener
+	service   string
 }
 
 // getWriter creates a io.Writer that will actually split by line and format by LogConsumer
-func getWriter(service, container string, events compose.ContainerEventListener) io.Writer {
+func getWriter(name, service, container string, events compose.ContainerEventListener) io.Writer {
 	return splitBuffer{
+		name:      name,
 		service:   service,
 		container: container,
 		consumer:  events,
@@ -108,6 +110,7 @@ func (s splitBuffer) Write(b []byte) (n int, err error) {
 		if len(line) != 0 {
 			s.consumer(compose.ContainerEvent{
 				Type:    compose.ContainerEventLog,
+				Name:    s.name,
 				Service: s.service,
 				Source:  s.container,
 				Line:    string(line),
diff --git a/local/compose/start.go b/local/compose/start.go
index c1aa7cc9c..8222e4344 100644
--- a/local/compose/start.go
+++ b/local/compose/start.go
@@ -55,7 +55,8 @@ func (s *composeService) Start(ctx context.Context, project *types.Project, opti
 			case status := <-statusC:
 				options.Attach(compose.ContainerEvent{
 					Type:     compose.ContainerEventExit,
-					Source:   getCanonicalContainerName(c),
+					Source:   c.ID,
+					Name:     getCanonicalContainerName(c),
 					Service:  c.Labels[serviceLabel],
 					ExitCode: int(status.StatusCode),
 				})