Streaming Transfers (#153)

internal/cron/cron.go

@@ -5,7 +5,7 @@ import (
 	"time"
 
 	"emperror.dev/errors"
-	log2 "github.com/apex/log"
+	"github.com/apex/log"
 	"github.com/go-co-op/gocron"
 
 	"github.com/pterodactyl/wings/config"
@@ -24,7 +24,7 @@ func Scheduler(ctx context.Context, m *server.Manager) (*gocron.Scheduler, error
 	if !o.SwapIf(true) {
 		return nil, errors.New("cron: cannot call scheduler more than once in application lifecycle")
 	}
-	l, err := time.LoadLocation(config.Get().System.Timezone)
+	location, err := time.LoadLocation(config.Get().System.Timezone)
 	if err != nil {
 		return nil, errors.Wrap(err, "cron: failed to parse configured system timezone")
 	}
@@ -41,30 +41,30 @@ func Scheduler(ctx context.Context, m *server.Manager) (*gocron.Scheduler, error
 		max:     config.Get().System.ActivitySendCount,
 	}
 
-	s := gocron.NewScheduler(l)
-	log := log2.WithField("subsystem", "cron")
+	s := gocron.NewScheduler(location)
+	l := log.WithField("subsystem", "cron")
 
 	interval := time.Duration(config.Get().System.ActivitySendInterval) * time.Second
-	log.WithField("interval", interval).Info("configuring system crons")
+	l.WithField("interval", interval).Info("configuring system crons")
 
 	_, _ = s.Tag("activity").Every(interval).Do(func() {
-		log.WithField("cron", "activity").Debug("sending internal activity events to Panel")
+		l.WithField("cron", "activity").Debug("sending internal activity events to Panel")
 		if err := activity.Run(ctx); err != nil {
 			if errors.Is(err, ErrCronRunning) {
-				log.WithField("cron", "activity").Warn("activity process is already running, skipping...")
+				l.WithField("cron", "activity").Warn("activity process is already running, skipping...")
 			} else {
-				log.WithField("cron", "activity").WithField("error", err).Error("activity process failed to execute")
+				l.WithField("cron", "activity").WithField("error", err).Error("activity process failed to execute")
 			}
 		}
 	})
 
 	_, _ = s.Tag("sftp").Every(interval).Do(func() {
-		log.WithField("cron", "sftp").Debug("sending sftp events to Panel")
+		l.WithField("cron", "sftp").Debug("sending sftp events to Panel")
 		if err := sftp.Run(ctx); err != nil {
 			if errors.Is(err, ErrCronRunning) {
-				log.WithField("cron", "sftp").Warn("sftp events process already running, skipping...")
+				l.WithField("cron", "sftp").Warn("sftp events process already running, skipping...")
 			} else {
-				log.WithField("cron", "sftp").WithField("error", err).Error("sftp events process failed to execute")
+				l.WithField("cron", "sftp").WithField("error", err).Error("sftp events process failed to execute")
 			}
 		}
 	})

internal/progress/progress.go

@@ -0,0 +1,90 @@
+package progress
+
+import (
+	"io"
+	"strings"
+	"sync/atomic"
+
+	"github.com/pterodactyl/wings/system"
+)
+
+// 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 uint64
+	// Total is the total size of the archive in bytes.
+	total uint64
+
+	// Writer .
+	Writer io.Writer
+}
+
+// NewProgress returns a new progress tracker for the given total size.
+func NewProgress(total uint64) *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() uint64 {
+	return atomic.LoadUint64(&p.written)
+}
+
+// Total returns the total size in bytes.
+func (p *Progress) Total() uint64 {
+	return atomic.LoadUint64(&p.total)
+}
+
+// SetTotal sets the total size of the archive in bytes. This function is safe
+// to call concurrently and can be used to update the total size if it changes,
+// such as when the total size is simultaneously being calculated as data is
+// being written through the progress writer.
+func (p *Progress) SetTotal(total uint64) {
+	atomic.StoreUint64(&p.total, 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.AddUint64(&p.written, uint64(n))
+	if p.Writer != nil {
+		return p.Writer.Write(v)
+	}
+	return n, nil
+}
+
+// Progress returns a formatted progress string for the current progress.
+func (p *Progress) Progress(width int) string {
+	// current = 100 (Progress, dynamic)
+	// total = 1000 (Content-Length, dynamic)
+	// width = 25 (Number of ticks to display, static)
+	// widthPercentage = 100 / width (What percentage does each tick represent, static)
+	//
+	// percentageDecimal = current / total = 0.1
+	// percentage = percentageDecimal * 100 = 10%
+	// ticks = percentage / widthPercentage = 2.5
+	//
+	// ticks is a float64, so we cast it to an int which rounds it down to 2.
+
+	// Values are cast to floats to prevent integer division.
+	current := p.Written()
+	total := p.Total()
+	// width := is passed as a parameter
+	widthPercentage := float64(100) / float64(width)
+	percentageDecimal := float64(current) / float64(total)
+	percentage := percentageDecimal * 100
+	ticks := int(percentage / widthPercentage)
+
+	// Ensure that we never get a negative number of ticks, this will prevent strings#Repeat
+	// from panicking.  A negative number of ticks is likely to happen when the total size is
+	// inaccurate, such as when we are going off of rough disk usage calculation.
+	if ticks < 0 {
+		ticks = 0
+	} else if ticks > width {
+		ticks = width
+	}
+
+	bar := strings.Repeat("=", ticks) + strings.Repeat(" ", width-ticks)
+	return "[" + bar + "] " + system.FormatBytes(current) + " / " + system.FormatBytes(total)
+}

internal/progress/progress_test.go

@@ -0,0 +1,50 @@
+package progress_test
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/franela/goblin"
+
+	"github.com/pterodactyl/wings/internal/progress"
+)
+
+func TestProgress(t *testing.T) {
+	g := goblin.Goblin(t)
+
+	g.Describe("Progress", func() {
+		g.It("properly initializes", func() {
+			total := uint64(1000)
+			p := progress.NewProgress(total)
+			g.Assert(p).IsNotNil()
+			g.Assert(p.Total()).Equal(total)
+			g.Assert(p.Written()).Equal(uint64(0))
+		})
+
+		g.It("increments written when Write is called", func() {
+			v := []byte("hello")
+			p := progress.NewProgress(1000)
+			_, err := p.Write(v)
+			g.Assert(err).IsNil()
+			g.Assert(p.Written()).Equal(uint64(len(v)))
+		})
+
+		g.It("renders a progress bar", func() {
+			v := bytes.Repeat([]byte{' '}, 100)
+			p := progress.NewProgress(1000)
+			_, err := p.Write(v)
+			g.Assert(err).IsNil()
+			g.Assert(p.Written()).Equal(uint64(len(v)))
+			g.Assert(p.Progress(25)).Equal("[==                       ] 100 B / 1000 B")
+		})
+
+		g.It("renders a progress bar when written exceeds total", func() {
+			v := bytes.Repeat([]byte{' '}, 1001)
+			p := progress.NewProgress(1000)
+			_, err := p.Write(v)
+			g.Assert(err).IsNil()
+			g.Assert(p.Written()).Equal(uint64(len(v)))
+			g.Assert(p.Progress(25)).Equal("[=========================] 1001 B / 1000 B")
+		})
+	})
+}