make color assignment predictable and option to disable prefixes/colors

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

Author: ndeloof

api/compose/api.go

@@ -184,6 +184,7 @@ type Stack struct {
 type LogConsumer interface {
 	Log(service, container, message string)
 	Status(service, container, msg string)
+	Register(service string, source string)
 }
 
 // ContainerEventListener is a callback to process ContainerEvent from services
@@ -201,6 +202,8 @@ type ContainerEvent struct {
 const (
 	// ContainerEventLog is a ContainerEvent of type log. Line is set
 	ContainerEventLog = iota
+	// ContainerEventAttach is a ContainerEvent of type attach. First event sent about a container
+	ContainerEventAttach
 	// ContainerEventExit is a ContainerEvent of type exit. ExitCode is set
 	ContainerEventExit
 )

cli/cmd/compose/logs.go

@@ -31,8 +31,10 @@ import (
 type logsOptions struct {
 	*projectOptions
 	composeOptions
-	follow bool
-	tail   string
+	follow   bool
+	tail     string
+	noColor  bool
+	noPrefix bool
 }
 
 func logsCommand(p *projectOptions, contextType string) *cobra.Command {
@@ -46,9 +48,13 @@ func logsCommand(p *projectOptions, contextType string) *cobra.Command {
 			return runLogs(cmd.Context(), opts, args)
 		},
 	}
-	logsCmd.Flags().BoolVar(&opts.follow, "follow", false, "Follow log output.")
+	flags := logsCmd.Flags()
+	flags.BoolVar(&opts.follow, "follow", false, "Follow log output.")
+	flags.BoolVar(&opts.noColor, "no-color", false, "Produce monochrome output.")
+	flags.BoolVar(&opts.noPrefix, "no-log-prefix", false, "Don't print prefix in logs.")
+
 	if contextType == store.DefaultContextType {
-		logsCmd.Flags().StringVar(&opts.tail, "tail", "all", "Number of lines to show from the end of the logs for each container.")
+		flags.StringVar(&opts.tail, "tail", "all", "Number of lines to show from the end of the logs for each container.")
 	}
 	return logsCmd
 }
@@ -63,7 +69,7 @@ func runLogs(ctx context.Context, opts logsOptions, services []string) error {
 	if err != nil {
 		return err
 	}
-	consumer := formatter.NewLogConsumer(ctx, os.Stdout)
+	consumer := formatter.NewLogConsumer(ctx, os.Stdout, !opts.noColor, !opts.noPrefix)
 	return c.ComposeService().Logs(ctx, projectName, consumer, compose.LogOptions{
 		Services: services,
 		Follow:   opts.follow,

cli/cmd/compose/start.go

@@ -78,7 +78,7 @@ func runStart(ctx context.Context, opts startOptions, services []string) error {
 		return err
 	}
 
-	_, err = printer.run(ctx, false, "", func() error {
+	_, err = printer.run(ctx, false, "", nil, func() error {
 		ctx := context.Background()
 		_, err := progress.Run(ctx, func(ctx context.Context) (string, error) {
 			return "", c.ComposeService().Stop(ctx, project)

cli/cmd/compose/up.go

@@ -19,6 +19,7 @@ package compose
 import (
 	"context"
 	"fmt"
+	"golang.org/x/sync/errgroup"
 	"os"
 	"os/signal"
 	"path/filepath"
@@ -57,6 +58,8 @@ type upOptions struct {
 	cascadeStop   bool
 	exitCodeFrom  string
 	scale         []string
+	noColor       bool
+	noPrefix      bool
 }
 
 func (o upOptions) recreateStrategy() string {
@@ -102,6 +105,8 @@ func upCommand(p *projectOptions, contextType string) *cobra.Command {
 	flags.BoolVar(&opts.Build, "build", false, "Build images before starting containers.")
 	flags.BoolVar(&opts.removeOrphans, "remove-orphans", false, "Remove containers for services not defined in the Compose file.")
 	flags.StringArrayVar(&opts.scale, "scale", []string{}, "Scale SERVICE to NUM instances. Overrides the `scale` setting in the Compose file if present.")
+	flags.BoolVar(&opts.noColor, "no-color", false, "Produce monochrome output.")
+	flags.BoolVar(&opts.noPrefix, "no-log-prefix", false, "Don't print prefix in logs.")
 
 	switch contextType {
 	case store.AciContextType:
@@ -199,6 +204,16 @@ func runCreateStart(ctx context.Context, opts upOptions, services []string) erro
 		stopFunc() // nolint:errcheck
 	}()
 
+	consumer := formatter.NewLogConsumer(ctx, os.Stdout, !opts.noColor, !opts.noPrefix)
+
+	var exitCode int
+	eg, ctx := errgroup.WithContext(ctx)
+	eg.Go(func() error {
+		code, err := printer.run(ctx, opts.cascadeStop, opts.exitCodeFrom, consumer, stopFunc)
+		exitCode = code
+		return err
+	})
+
 	err = c.ComposeService().Start(ctx, project, compose.StartOptions{
 		Attach: func(event compose.ContainerEvent) {
 			queue <- event
@@ -208,7 +223,7 @@ func runCreateStart(ctx context.Context, opts upOptions, services []string) erro
 		return err
 	}
 
-	exitCode, err := printer.run(ctx, opts.cascadeStop, opts.exitCodeFrom, stopFunc)
+	eg.Wait()
 	if exitCode != 0 {
 		return cmd.ExitCodeError{ExitCode: exitCode}
 	}
@@ -298,12 +313,13 @@ type printer struct {
 	queue chan compose.ContainerEvent
 }
 
-func (p printer) run(ctx context.Context, cascadeStop bool, exitCodeFrom string, stopFn func() error) (int, error) { //nolint:unparam
-	consumer := formatter.NewLogConsumer(ctx, os.Stdout)
+func (p printer) run(ctx context.Context, cascadeStop bool, exitCodeFrom string, consumer compose.LogConsumer, stopFn func() error) (int, error) { //nolint:unparam
 	var aborting bool
 	for {
 		event := <-p.queue
 		switch event.Type {
+		case compose.ContainerEventAttach:
+			consumer.Register(event.Service, event.Source)
 		case compose.ContainerEventExit:
 			if !aborting {
 				consumer.Status(event.Service, event.Source, fmt.Sprintf("exited with code %d", event.ExitCode))

cli/formatter/colors.go

@@ -35,6 +35,10 @@ var names = []string{
 // 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 ansiColor(code, s string) string {
 	return fmt.Sprintf("%s%s%s", ansi(code), s, ansi("0"))
 }

cli/formatter/logs.go

@@ -17,7 +17,6 @@
 package formatter
 
 import (
-	"bytes"
 	"context"
 	"fmt"
 	"io"
@@ -28,59 +27,91 @@ import (
 )
 
 // NewLogConsumer creates a new LogConsumer
-func NewLogConsumer(ctx context.Context, w io.Writer) compose.LogConsumer {
+func NewLogConsumer(ctx context.Context, w io.Writer, color bool, prefix bool) compose.LogConsumer {
 	return &logConsumer{
-		ctx:    ctx,
-		colors: map[string]colorFunc{},
-		width:  0,
-		writer: w,
+		ctx:        ctx,
+		presenters: map[string]*Presenter{},
+		width:      0,
+		writer:     w,
+		color:      color,
+		prefix:     prefix,
 	}
 }
 
+func (l *logConsumer) Register(service string, source string) {
+	l.register(service, source)
+}
+
+func (l *logConsumer) register(service string, source string) *Presenter {
+	cf := monochrome
+	if l.color {
+		cf = <-loop
+	}
+	p := &Presenter{
+		colors:    cf,
+		service:   service,
+		container: source,
+	}
+	l.presenters[source] = p
+	if l.prefix {
+		l.computeWidth()
+		for _, p := range l.presenters {
+			p.setPrefix(l.width)
+		}
+	}
+	return p
+}
+
 // Log formats a log message as received from service/container
 func (l *logConsumer) Log(service, container, message string) {
 	if l.ctx.Err() != nil {
 		return
 	}
-	cf := l.getColorFunc(service)
-	prefix := fmt.Sprintf("%-"+strconv.Itoa(l.width)+"s |", container)
-
+	p, ok := l.presenters[container]
+	if !ok { // should have been registered, but ¯\_(ツ)_/¯
+		p = l.register(service, container)
+	}
 	for _, line := range strings.Split(message, "\n") {
-		buf := bytes.NewBufferString(fmt.Sprintf("%s %s\n", cf(prefix), line))
-		l.writer.Write(buf.Bytes()) // nolint:errcheck
+		fmt.Fprintf(l.writer, "%s %s\n", p.prefix, line) // nolint:errcheck
 	}
 }
 
 func (l *logConsumer) Status(service, container, msg string) {
-	cf := l.getColorFunc(service)
-	buf := bytes.NewBufferString(cf(fmt.Sprintf("%s %s\n", container, msg)))
-	l.writer.Write(buf.Bytes()) // nolint:errcheck
-}
-
-func (l *logConsumer) getColorFunc(service string) colorFunc {
-	cf, ok := l.colors[service]
+	p, ok := l.presenters[container]
 	if !ok {
-		cf = <-loop
-		l.colors[service] = cf
-		l.computeWidth()
+		p = l.register(service, container)
 	}
-	return cf
+	s := p.colors(fmt.Sprintf("%s %s\n", container, msg))
+	l.writer.Write([]byte(s)) // nolint:errcheck
 }
 
 func (l *logConsumer) computeWidth() {
 	width := 0
-	for n := range l.colors {
+	for n := range l.presenters {
 		if len(n) > width {
 			width = len(n)
 		}
 	}
-	l.width = width + 3
+	l.width = width + 1
 }
 
 // LogConsumer consume logs from services and format them
 type logConsumer struct {
-	ctx    context.Context
-	colors map[string]colorFunc
-	width  int
-	writer io.Writer
+	ctx        context.Context
+	presenters map[string]*Presenter
+	width      int
+	writer     io.Writer
+	color      bool
+	prefix     bool
+}
+
+type Presenter struct {
+	colors    colorFunc
+	service   string
+	container string
+	prefix    string
+}
+
+func (p *Presenter) setPrefix(width int) {
+	p.prefix = p.colors(fmt.Sprintf("%-"+strconv.Itoa(width)+"s |", p.container))
 }

local/compose/attach.go

@@ -36,13 +36,21 @@ func (s *composeService) attach(ctx context.Context, project *types.Project, con
 		return nil, err
 	}
 
+	containers.sorted() // This enforce predictable colors assignment
+
 	var names []string
 	for _, c := range containers {
 		names = append(names, getCanonicalContainerName(c))
 	}
+
 	fmt.Printf("Attaching to %s\n", strings.Join(names, ", "))
 
 	for _, container := range containers {
+		consumer(compose.ContainerEvent{
+			Type:    compose.ContainerEventAttach,
+			Source:  getContainerNameWithoutProject(container),
+			Service: container.Labels[serviceLabel],
+		})
 		err := s.attachContainer(ctx, container, consumer, project)
 		if err != nil {
 			return nil, err

local/compose/containers.go

@@ -18,6 +18,7 @@ package compose
 
 import (
 	"context"
+	"sort"
 
 	"github.com/compose-spec/compose-go/types"
 	moby "github.com/docker/docker/api/types"
@@ -83,3 +84,10 @@ func (containers Containers) forEach(fn func(moby.Container)) {
 		fn(c)
 	}
 }
+
+func (containers Containers) sorted() Containers {
+	sort.Slice(containers, func(i, j int) bool {
+		return getCanonicalContainerName(containers[i]) < getCanonicalContainerName(containers[j])
+	})
+	return containers
+}

utils/logconsumer.go

@@ -64,6 +64,12 @@ func (a *allowListLogConsumer) Status(service, container, message string) {
 	}
 }
 
+func (a *allowListLogConsumer) Register(service string, source string) {
+	if a.allowList[service] {
+		a.delegate.Register(service, source)
+	}
+}
+
 type splitBuffer struct {
 	service   string
 	container string