fix: Add detailed output for compile stage (#131)

* fix: Add detailed output for compile stage

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* fix: Fix test cases

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* fix: Fix lint

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* fix: Fix lint

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

pkg/editor/vscode/vsocde.go

@@ -30,7 +30,7 @@ import (
 )
 
 type Client interface {
-	DownloadOrCache(plugin Plugin) error
+	DownloadOrCache(plugin Plugin) (bool, error)
 	PluginPath(p Plugin) string
 }
 
@@ -50,7 +50,9 @@ func unzipPath(p Plugin) string {
 		p.Publisher, p.Extension, p.Version)
 }
 
-func (c generalClient) DownloadOrCache(p Plugin) error {
+// DownloadOrCache downloads or cache the plugin.
+// If the plugin is already downloaded, it returns true.
+func (c generalClient) DownloadOrCache(p Plugin) (bool, error) {
 	url := fmt.Sprintf(vscodePackageURLTemplate,
 		p.Publisher, p.Publisher, p.Extension, p.Version)
 
@@ -64,36 +66,36 @@ func (c generalClient) DownloadOrCache(p Plugin) error {
 		"file":      filename,
 	})
 	if ok, err := fileutil.FileExists(filename); err != nil {
-		return err
+		return false, err
 	} else if ok {
 		logger.Debug("vscode plugin is cached")
-		return nil
+		return true, nil
 	}
 	out, err := os.Create(filename)
 
 	if err != nil {
-		return err
+		return false, err
 	}
 	defer out.Close()
 
 	resp, err := http.Get(url)
 	if err != nil {
-		return err
+		return false, err
 	}
 	logger.Debugf("downloading vscode plugin")
 
 	defer resp.Body.Close()
 	_, err = io.Copy(out, resp.Body)
 	if err != nil {
-		return err
+		return false, err
 	}
 
 	_, err = unzip.Unzip(filename, unzipPath(p))
 	if err != nil {
-		return errors.Wrap(err, "failed to unzip")
+		return false, errors.Wrap(err, "failed to unzip")
 	}
 
-	return nil
+	return false, nil
 }
 
 func ParsePlugin(p string) (*Plugin, error) {

pkg/lang/ir/graph.go

@@ -233,8 +233,11 @@ func (g Graph) compileVSCode() (*llb.State, error) {
 	inputs := []llb.State{}
 	for _, p := range g.VSCodePlugins {
 		vscodeClient := vscode.NewClient()
-		if err := vscodeClient.DownloadOrCache(p); err != nil {
+		g.Writer.LogVSCodePlugin(p, compileui.ActionStart, false)
+		if cached, err := vscodeClient.DownloadOrCache(p); err != nil {
 			return nil, err
+		} else {
+			g.Writer.LogVSCodePlugin(p, compileui.ActionEnd, cached)
 		}
 		ext := llb.Scratch().File(llb.Copy(llb.Local(flag.FlagCacheDir),
 			vscodeClient.PluginPath(p),
@@ -278,8 +281,11 @@ func (g Graph) compilePyPIMirror(root llb.State) llb.State {
 func (g Graph) compileZSH(root llb.State) (llb.State, error) {
 	installPath := "/root/install.sh"
 	m := shell.NewManager()
-	if err := m.DownloadOrCache(); err != nil {
+	g.Writer.LogZSH(compileui.ActionStart, false)
+	if cached, err := m.DownloadOrCache(); err != nil {
 		return llb.State{}, errors.Wrap(err, "failed to download oh-my-zsh")
+	} else {
+		g.Writer.LogZSH(compileui.ActionEnd, cached)
 	}
 	zshStage := root.
 		File(llb.Copy(llb.Local(flag.FlagCacheDir), "oh-my-zsh", "/root/.oh-my-zsh",

pkg/progress/compileui/display.go

@@ -17,25 +17,30 @@ package compileui
 import (
 	"context"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/cockroachdb/errors"
 	"github.com/containerd/console"
 	"github.com/morikuni/aec"
 	"github.com/sirupsen/logrus"
+	"github.com/tensorchord/MIDI/pkg/editor/vscode"
 )
 
 type Writer interface {
-	Print(s string)
+	LogVSCodePlugin(p vscode.Plugin, action Action, cached bool)
+	LogZSH(action Action, cached bool)
 	Finish()
 }
 
 type generalWriter struct {
-	console console.Console
-	phase   string
-	trace   *trace
-	doneCh  chan bool
-	repeatd bool
+	console   console.Console
+	phase     string
+	trace     *trace
+	doneCh    chan bool
+	repeatd   bool
+	result    *Result
+	lineCount int
 }
 
 func New(ctx context.Context, out console.File, mode string) (Writer, error) {
@@ -61,15 +66,48 @@ func New(ctx context.Context, out console.File, mode string) (Writer, error) {
 		trace:   t,
 		doneCh:  make(chan bool),
 		repeatd: false,
+		result: &Result{
+			plugins: make(map[string]*PluginInfo),
+		},
+		lineCount: 0,
 	}
 	// TODO(gaocegege): Have a result chan
 	//nolint
 	go w.run(ctx)
 	return w, nil
 }
 
-func (w generalWriter) Print(s string) {
-	fmt.Fprintln(w.console, s)
+func (w *generalWriter) LogVSCodePlugin(p vscode.Plugin, action Action, cached bool) {
+	switch action {
+	case ActionStart:
+		c := time.Now()
+		w.result.plugins[p.String()] = &PluginInfo{
+			Plugin:    p,
+			startTime: &c,
+			cached:    cached,
+		}
+	case ActionEnd:
+		c := time.Now()
+		w.result.plugins[p.String()].endTime = &c
+		w.result.plugins[p.String()].cached = cached
+	}
+
+}
+
+func (w *generalWriter) LogZSH(action Action, cached bool) {
+	switch action {
+	case ActionStart:
+		c := time.Now()
+		w.result.ZSHInfo = &ZSHInfo{
+			OHMYZSH:   "oh-my-zsh",
+			startTime: &c,
+			cached:    cached,
+		}
+	case ActionEnd:
+		c := time.Now()
+		w.result.ZSHInfo.endTime = &c
+		w.result.ZSHInfo.cached = cached
+	}
 }
 
 func (w generalWriter) Finish() {
@@ -79,33 +117,93 @@ func (w generalWriter) Finish() {
 func (w *generalWriter) run(ctx context.Context) error {
 	displayTimeout := 100 * time.Millisecond
 	ticker := time.NewTicker(displayTimeout)
-	width, height := w.getSize()
-	logger := logrus.WithFields(logrus.Fields{
-		"console-height": height,
-		"console-width":  width,
-	})
-	logger.Debug("print compile progress")
 	for {
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
 		case <-w.doneCh:
+			w.output(true)
 			return nil
 		case <-ticker.C:
-			b := aec.EmptyBuilder.Up(1)
-			if !w.repeatd {
-				b = b.Down(1)
-			}
-			w.repeatd = true
-			fmt.Fprint(w.console, b.Column(0).ANSI)
-			fmt.Fprint(w.console, aec.Hide)
-			defer fmt.Fprint(w.console, aec.Show)
-			s := fmt.Sprintf("[+] ⌚ %s %.1fs\n", w.phase, time.Since(*w.trace.startTime).Seconds())
-			fmt.Fprint(w.console, s)
+			w.output(false)
 		}
 	}
 }
 
+func (w *generalWriter) output(finished bool) {
+	width, _ := w.getSize()
+	b := aec.EmptyBuilder.Up(uint(1 + w.lineCount))
+	if !w.repeatd {
+		b = b.Down(1)
+	}
+	w.repeatd = true
+	if finished {
+		fmt.Fprint(w.console, colorRun)
+	}
+	fmt.Fprint(w.console, b.Column(0).ANSI)
+	fmt.Fprint(w.console, aec.Hide)
+	defer fmt.Fprint(w.console, aec.Show)
+
+	statusStr := ""
+	if finished {
+		statusStr = "✅ (finished)"
+	}
+	s := fmt.Sprintf("[+] ⌚ %s %.1fs %s \n",
+		w.phase, time.Since(*w.trace.startTime).Seconds(), statusStr)
+	fmt.Fprint(w.console, s)
+	loc := 0
+
+	// output shell info.
+	if w.result.ZSHInfo != nil {
+		timer := time.Since(*w.result.ZSHInfo.startTime).Seconds()
+		if w.result.ZSHInfo.endTime != nil {
+			timer = w.result.ZSHInfo.endTime.Sub(*w.result.ZSHInfo.startTime).Seconds()
+		}
+		template := " => download %s"
+		if w.result.ZSHInfo.cached {
+			template = " => 💽 (cached) download %s"
+		}
+		timerStr := fmt.Sprintf(" %.1fs\n", timer)
+		out := fmt.Sprintf(template, w.result.ZSHInfo.OHMYZSH)
+		out = align(out, timerStr, width)
+		fmt.Fprint(w.console, out)
+		loc++
+	}
+
+	// output vscode plugins.
+	for _, p := range w.result.plugins {
+		if p.startTime == nil {
+			continue
+		}
+		timer := time.Since(*p.startTime).Seconds()
+		if p.endTime != nil {
+			timer = p.endTime.Sub(*p.startTime).Seconds()
+		}
+		timerStr := fmt.Sprintf(" %.1fs\n", timer)
+		template := " => download %s"
+		if p.cached {
+			template = " => 💽 (cached) download %s"
+		}
+		out := fmt.Sprintf(template, p.Plugin)
+		out = align(out, timerStr, width)
+		fmt.Fprint(w.console, out)
+		loc++
+	}
+
+	// override previous content
+	if diff := w.lineCount - loc; diff > 0 {
+		logrus.WithFields(logrus.Fields{
+			"diff":    diff,
+			"plugins": len(w.result.plugins),
+		}).Debug("override previous content", diff)
+		for i := 0; i < diff; i++ {
+			fmt.Fprintln(w.console, strings.Repeat(" ", width))
+		}
+		fmt.Fprint(w.console, aec.EmptyBuilder.Up(uint(diff)).Column(0).ANSI)
+	}
+	w.lineCount = loc
+}
+
 func (w generalWriter) getSize() (int, int) {
 	width := 80
 	height := 10
@@ -118,3 +216,7 @@ func (w generalWriter) getSize() (int, int) {
 	}
 	return width, height
 }
+
+func align(l, r string, w int) string {
+	return fmt.Sprintf("%-[2]*[1]s %[3]s", l, w-len(r)-1, r)
+}

pkg/progress/compileui/term.go

@@ -0,0 +1,27 @@
+// Copyright 2022 The MIDI Authors
+// Copyright 2022 The buildkit Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build !windows
+// +build !windows
+
+package compileui
+
+import "github.com/morikuni/aec"
+
+var (
+	colorRun = aec.BlueF
+	// colorCancel = aec.YellowF
+	// colorError  = aec.RedF
+)

pkg/progress/compileui/term_windows.go

@@ -0,0 +1,27 @@
+// Copyright 2022 The MIDI Authors
+// Copyright 2022 The buildkit Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build windows
+// +build windows
+
+package compileui
+
+import "github.com/morikuni/aec"
+
+var (
+	colorRun    = aec.CyanF
+	colorCancel = aec.YellowF
+	colorError  = aec.RedF
+)