server(transfers): track progress of archive creation and extraction (#143)

server/filesystem/archive.go

@@ -8,6 +8,7 @@ import (
 	"path/filepath"
 	"strings"
 	"sync"
+	"sync/atomic"
 
 	"emperror.dev/errors"
 	"github.com/apex/log"
@@ -17,6 +18,7 @@ import (
 	ignore "github.com/sabhiram/go-gitignore"
 
 	"github.com/pterodactyl/wings/config"
+	"github.com/pterodactyl/wings/system"
 )
 
 const memory = 4 * 1024
@@ -28,6 +30,62 @@ var pool = sync.Pool{
 	},
 }
 
+// Progress is used to track the progress of any I/O operation that are being
+// performed.
+type Progress struct {
+	// written is the total size of the files that have been written to the writer.
+	written int64
+	// Total is the total size of the archive in bytes.
+	total int64
+	// w .
+	w io.Writer
+}
+
+// NewProgress .
+func NewProgress(total int64) *Progress {
+	return &Progress{total: total}
+}
+
+// Written returns the total number of bytes written.
+// This function should be used when the progress is tracking data being written.
+func (p *Progress) Written() int64 {
+	return atomic.LoadInt64(&p.written)
+}
+
+// Total returns the total size in bytes.
+func (p *Progress) Total() int64 {
+	return atomic.LoadInt64(&p.total)
+}
+
+// Write totals the number of bytes that have been written to the writer.
+func (p *Progress) Write(v []byte) (int, error) {
+	n := len(v)
+	atomic.AddInt64(&p.written, int64(n))
+	if p.w != nil {
+		return p.w.Write(v)
+	}
+	return n, nil
+}
+
+// Progress returns a formatted progress string for the current progress.
+func (p *Progress) Progress(width int) string {
+	current := p.Written()
+	total := p.Total()
+
+	// v = 100 (Progress)
+	// size = 1000 (Content-Length)
+	// p / size = 0.1
+	// * 100 = 10% (Multiply by 100 to get a percentage of the download)
+	// 10% / tickPercentage = (10% / (100 / 25)) (Divide by tick percentage to get the number of ticks)
+	// 2.5 (Number of ticks as a float64)
+	// 2 (convert to an integer)
+
+	// We have to cast these numbers to float in order to get a float result from the division.
+	ticks := ((float64(current) / float64(total)) * 100) / (float64(100) / float64(width))
+	bar := strings.Repeat("=", int(ticks)) + strings.Repeat(" ", width-int(ticks))
+	return "[" + bar + "] " + system.FormatBytes(current) + " / " + system.FormatBytes(total)
+}
+
 type Archive struct {
 	// BasePath is the absolute path to create the archive from where Files and Ignore are
 	// relative to.
@@ -40,10 +98,13 @@ type Archive struct {
 	// Files specifies the files to archive, this takes priority over the Ignore option, if
 	// unspecified, all files in the BasePath will be archived unless Ignore is set.
 	Files []string
+
+	// Progress wraps the writer of the archive to pass through the progress tracker.
+	Progress *Progress
 }
 
-// Create creates an archive at dst with all of the files defined in the
-// included files struct.
+// Create creates an archive at dst with all the files defined in the
+// included Files array.
 func (a *Archive) Create(dst string) error {
 	f, err := os.OpenFile(dst, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600)
 	if err != nil {
@@ -62,26 +123,34 @@ func (a *Archive) Create(dst string) error {
 		writer = f
 	}
 
-	// The default compression level is BestSpeed
-	var cl = pgzip.BestSpeed
-
 	// Choose which compression level to use based on the compression_level configuration option
+	var compressionLevel int
 	switch config.Get().System.Backups.CompressionLevel {
 	case "none":
-		cl = pgzip.NoCompression
-	case "best_speed":
-		cl = pgzip.BestSpeed
+		compressionLevel = pgzip.NoCompression
 	case "best_compression":
-		cl = pgzip.BestCompression
+		compressionLevel = pgzip.BestCompression
+	case "best_speed":
+		fallthrough
+	default:
+		compressionLevel = pgzip.BestSpeed
 	}
 
 	// Create a new gzip writer around the file.
-	gw, _ := pgzip.NewWriterLevel(writer, cl)
+	gw, _ := pgzip.NewWriterLevel(writer, compressionLevel)
 	_ = gw.SetConcurrency(1<<20, 1)
 	defer gw.Close()
 
+	var pw io.Writer
+	if a.Progress != nil {
+		a.Progress.w = gw
+		pw = a.Progress
+	} else {
+		pw = gw
+	}
+
 	// Create a new tar writer around the gzip writer.
-	tw := tar.NewWriter(gw)
+	tw := tar.NewWriter(pw)
 	defer tw.Close()
 
 	// Configure godirwalk.
@@ -116,7 +185,7 @@ func (a *Archive) Create(dst string) error {
 // being generated.
 func (a *Archive) callback(tw *tar.Writer, opts ...func(path string, relative string) error) func(path string, de *godirwalk.Dirent) error {
 	return func(path string, de *godirwalk.Dirent) error {
-		// Skip directories because we walking them recursively.
+		// Skip directories because we are walking them recursively.
 		if de.IsDir() {
 			return nil
 		}