Externalize the Progress Bar output from the Download Logic

The progress bar output was embedded in the download logic, preventing a full reuse of the download logic itself.
It is now outsourced and the download logic just has lifecycle handles to notify an external actor of the progress.

Author: lorenzo-biava

commands/commands.go

@@ -31,6 +31,8 @@ package commands
 
 import (
 	"github.com/sirupsen/logrus"
+	"github.com/bcmi-labs/arduino-cli/common/releases"
+	"github.com/bcmi-labs/arduino-cli/common/formatter"
 )
 
 // Error codes to be used for os.Exit().
@@ -53,3 +55,11 @@ var GlobalFlags struct {
 	Debug  bool   // If true, dump debug output to stderr.
 	Format string // The Output format (e.g. text, json).
 }
+
+// FIXME: Move away? Where should the display logic reside; in the formatter? That causes an import cycle BTW...
+func GenerateDownloadProgressFormatter() releases.ParallelDownloadProgressHandler {
+	if formatter.IsCurrentFormat("text") {
+		return &ProgressBarFormatter{}
+	}
+	return nil
+}

commands/core/download.go

@@ -96,7 +96,7 @@ func downloadToolsArchives(tools []*cores.ToolRelease, results *output.CoreProce
 		})
 	}
 	logrus.Info("Downloading tools")
-	releases.ParallelDownload(downloads, false, "Downloaded", &results.Tools, "tool")
+	releases.ParallelDownload(downloads, false, "Downloaded", &results.Tools, "tool", commands.GenerateDownloadProgressFormatter())
 }
 
 func downloadPlatformArchives(platforms []*cores.PlatformRelease, results *output.CoreProcessResults) {
@@ -108,5 +108,5 @@ func downloadPlatformArchives(platforms []*cores.PlatformRelease, results *outpu
 		})
 	}
 	logrus.Info("Downloading cores")
-	releases.ParallelDownload(downloads, false, "Downloaded", &results.Cores, "core")
+	releases.ParallelDownload(downloads, false, "Downloaded", &results.Cores, "core", commands.GenerateDownloadProgressFormatter())
 }

commands/lib/download.go

@@ -78,7 +78,7 @@ func runDownloadCommand(cmd *cobra.Command, args []string) {
 		libs[i] = releases.DownloadItem(libsToDownload[i])
 	}
 	logrus.Info("Downloading")
-	releases.ParallelDownload(libs, false, "Downloaded", &outputResults.Libraries, "library")
+	releases.ParallelDownload(libs, false, "Downloaded", &outputResults.Libraries, "library", commands.GenerateDownloadProgressFormatter())
 	logrus.Info("Download finished")
 	formatter.Print(outputResults)
 	logrus.Info("Done")

commands/lib/install.go

@@ -80,7 +80,7 @@ func runInstallCommand(cmd *cobra.Command, args []string) {
 	}
 
 	logrus.Info("Downloading")
-	releases.ParallelDownload(libs, false, "Installed", &outputResults.Libraries, "library")
+	releases.ParallelDownload(libs, false, "Installed", &outputResults.Libraries, "library", commands.GenerateDownloadProgressFormatter())
 	logrus.Info("Download finished")
 
 	logrus.Info("Installing")

commands/progress_formatter.go

@@ -0,0 +1,75 @@
+package commands
+
+import (
+	"gopkg.in/cheggaaa/pb.v1"
+	"fmt"
+
+	"github.com/sirupsen/logrus"
+	"github.com/bcmi-labs/arduino-cli/common/releases"
+)
+
+// ProgressBarFormatter implements the visualization of the progress bars
+// to display the progress of a ParallelDownload task set (i.e. one bar per file)
+// WARNING: The implementation library is experimental and unstable; Do not print to terminal while the bars are active.
+type ProgressBarFormatter struct {
+	progressBars     map[string]*pb.ProgressBar
+	progressBarsPool *pb.Pool
+}
+
+// Implement interface releases.ParallelDownloadProgressHandler
+
+func (pbf *ProgressBarFormatter) OnNewDownloadTask(fileName string, fileSize int64) releases.FileDownloadFilter {
+	// Initialize progress bars and a new one for each the new task
+	if pbf.progressBars == nil {
+		pbf.progressBars = map[string]*pb.ProgressBar{}
+	}
+
+	logrus.Debug(fmt.Sprintf("Initializing progress bar for file '%s'", fileName))
+
+	// Initialization is in bytes, to display full information about the file (not only the percentage)
+	progressBar := pb.New64(fileSize).SetUnits(pb.U_BYTES).Prefix(fmt.Sprintf("%-20s", fileName)).Start()
+	pbf.progressBars[fileName] = progressBar
+
+	// TODO: this was the legacy way to run the progress bar update; since the logic has been outsourced
+	// and the OnProgressChanged callback is now available, it can be safely removed.
+	/*return func(source io.Reader, initialData *os.File, initialSize int) (io.Reader, error) {
+		logrus.Info(fmt.Sprintf("Initialized progress bar for file '%s'", fileName))
+
+		progressBar.Add(int(initialSize))
+		return progressBar.NewProxyReader(source), nil
+	}*/
+	return nil
+}
+
+func (pbf *ProgressBarFormatter) OnProgressChanged(fileName string, fileSize int64, downloadedSoFar int64) {
+	// Update a specific file's progress bar
+	progressBar := pbf.progressBars[fileName]
+
+	if progressBar != nil {
+		progressBar.Set(int(downloadedSoFar))
+	} else {
+		logrus.Debug(fmt.Sprintf("Progress bar for file '%s' not found", fileName))
+	}
+}
+
+func (pbf *ProgressBarFormatter) OnDownloadStarted() {
+	// WARNING!!
+	// (experimental and unstable) Do not print to terminal while pool is active.
+	// See https://github.com/cheggaaa/pb#multiple-progress-bars-experimental-and-unstable
+
+	// Start the progress bar pool
+	progressBarsAsSlice := []*pb.ProgressBar{}
+	for _, value := range pbf.progressBars {
+		progressBarsAsSlice = append(progressBarsAsSlice, value)
+	}
+	pbf.progressBarsPool, _ = pb.StartPool(progressBarsAsSlice...)
+}
+
+func (pbf *ProgressBarFormatter) OnDownloadFinished() {
+	// Stop the progress bar pool
+	if pbf.progressBarsPool != nil {
+		pbf.progressBarsPool.Stop()
+	}
+}
+
+// END -- Implement interface releases.ParallelDownloadProgressHandler

common/net_functions.go

@@ -75,8 +75,32 @@ func DownloadIndex(indexPath pathutils.Path, URL string) error {
 	return nil
 }
 
-// DownloadPackage downloads a package from arduino repository.
-func DownloadPackage(URL string, initialData *os.File, totalSize int64, handleResultFunc func(io.Reader, *os.File, int) error) error {
+// HandleResultFunc defines a function able to handle the content of the
+// download stream of the package (DownloadPackage), filling the File with the content
+// of the Reader, starting from the initial position
+type HandleDownloadPackageResultFunc func(io.Reader, *os.File, int) error
+
+// DefaultDownloadHandlerFunc is the default HandleDownloadPackageResultFunc, which
+// simply copies the content of the Reader into the File, starting from the initialSize
+func DefaultDownloadHandlerFunc(source io.Reader, initialData *os.File, initialSize int) error {
+	// Copy the file content
+	_, err := io.Copy(initialData, source)
+	return err
+}
+
+// DownloadPackageProgressChangedHandler defines a function able to handle the update
+// of the progress of the current download
+type DownloadPackageProgressChangedHandler func(fileSize int64, downloadedSoFar int64)
+
+// DownloadPackage downloads a package from Arduino repository.
+// Besides the download information (URL, initialData and totalSize), two external handlers are available for:
+//  - (handleResultFunc) handling the result of the download (i.e. decide how to copy the download to the file
+// 	  or do something weird during the operation)
+//  - (progressChangedHandler) handling the download progress change (and perhaps display it somehow)
+// None of the handlers is mandatory; they won't be used if nil.
+func DownloadPackage(URL string, initialData *os.File, totalSize int64, handleResultFunc HandleDownloadPackageResultFunc,
+	progressChangedHandler DownloadPackageProgressChangedHandler) error {
+
 	if initialData == nil {
 		return errors.New("Cannot fill a nil file pointer")
 	}
@@ -120,13 +144,51 @@ func DownloadPackage(URL string, initialData *os.File, totalSize int64, handleRe
 	}
 	defer response.Body.Close()
 
+	// Handle the progress update handler, by creating a ProgressProxyReader;
+	// only if it's needed (i.e. we actually have an external progressChangedHandler)
+	progressProxyReader := response.Body
+	downloadedSoFar := initialSize
+	if progressChangedHandler != nil {
+		progressProxyReader = ProgressProxyReader{response.Body, func(progressDelta int64) {
+			// WARNING: This is using a closure on downloadedSoFar!
+			downloadedSoFar += progressDelta
+			progressChangedHandler(totalSize, downloadedSoFar)
+		},
+		}
+	}
+
+	// Use the external handleResultFunc, if available, or the default one otherwise
 	if handleResultFunc == nil {
-		_, err = io.Copy(initialData, response.Body)
-	} else {
-		err = handleResultFunc(response.Body, initialData, int(initialSize))
+		handleResultFunc = DefaultDownloadHandlerFunc
 	}
+
+	err = handleResultFunc(progressProxyReader, initialData, int(initialSize))
 	if err != nil {
 		return fmt.Errorf("Cannot read response body from %s : %s", URL, err)
 	}
 	return nil
 }
+
+// FIXME: Move outside? perhaps in commons?
+// HandleProgressUpdateFunc defines a function able to handle the progressDelta, in bytes
+type HandleProgressUpdateFunc func(progressDelta int64)
+
+// It's proxy reader, intercepting reads to post progress updates, implement io.Reader
+type ProgressProxyReader struct {
+	io.Reader
+	handleProgressUpdateFunc HandleProgressUpdateFunc
+}
+
+func (r ProgressProxyReader) Read(p []byte) (n int, err error) {
+	n, err = r.Reader.Read(p)
+	r.handleProgressUpdateFunc(int64(n))
+	return
+}
+
+// Close the reader when it implements io.Closer
+func (r ProgressProxyReader) Close() (err error) {
+	if closer, ok := r.Reader.(io.Closer); ok {
+		return closer.Close()
+	}
+	return
+}