From 1937d0366debf78fc5d222215cc77584a6004a98 Mon Sep 17 00:00:00 2001
From: Dane Everitt <dane@daneeveritt.com>
Date: Wed, 6 Jan 2021 20:47:44 -0800
Subject: [PATCH] cleanup; fix environment stats not reporting network TX
 correctly

---
 environment/docker/console.go   | 20 --------------------
 environment/docker/container.go |  4 ++++
 environment/docker/stats.go     | 30 +++++++++---------------------
 environment/stats.go            | 10 ++++++----
 4 files changed, 19 insertions(+), 45 deletions(-)
 delete mode 100644 environment/docker/console.go

diff --git a/environment/docker/console.go b/environment/docker/console.go
deleted file mode 100644
index b5a4fcc..0000000
--- a/environment/docker/console.go
+++ /dev/null
@@ -1,20 +0,0 @@
-package docker
-
-import "io"
-
-type Console struct {
-	HandlerFunc *func(string)
-}
-
-var _ io.Writer = Console{}
-
-func (c Console) Write(b []byte) (int, error) {
-	if c.HandlerFunc != nil {
-		l := make([]byte, len(b))
-		copy(l, b)
-
-		(*c.HandlerFunc)(string(l))
-	}
-
-	return len(b), nil
-}
diff --git a/environment/docker/container.go b/environment/docker/container.go
index 3e3827c..58a728d 100644
--- a/environment/docker/container.go
+++ b/environment/docker/container.go
@@ -76,6 +76,10 @@ func (e *Environment) Attach() error {
 			}
 		}()
 
+		// Block the completion of this routine until the container is no longer running. This allows
+		// the pollResources function to run until it needs to be stopped. Because the container
+		// can be polled for resource usage, even when sropped, we need to have this logic present
+		// in order to cancel the context and therefore stop the routine that is spawned.
 		ok, err := e.client.ContainerWait(ctx, e.Id, container.WaitConditionNotRunning)
 		select {
 		case <-ctx.Done():
diff --git a/environment/docker/stats.go b/environment/docker/stats.go
index 231af99..e446975 100644
--- a/environment/docker/stats.go
+++ b/environment/docker/stats.go
@@ -8,7 +8,6 @@ import (
 	"github.com/pterodactyl/wings/environment"
 	"io"
 	"math"
-	"sync/atomic"
 )
 
 // Attach to the instance and then automatically emit an event whenever the resource usage for the
@@ -28,21 +27,18 @@ func (e *Environment) pollResources(ctx context.Context) error {
 	defer stats.Body.Close()
 
 	dec := json.NewDecoder(stats.Body)
-
 	for {
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
 		default:
-			var v *types.StatsJSON
-
+			var v types.StatsJSON
 			if err := dec.Decode(&v); err != nil {
 				if err != io.EOF && !errors.Is(err, context.Canceled) {
 					e.log().WithField("error", err).Warn("error while processing Docker stats output for container")
 				} else {
 					e.log().Debug("io.EOF encountered during stats decode, stopping polling...")
 				}
-
 				return nil
 			}
 
@@ -52,24 +48,16 @@ func (e *Environment) pollResources(ctx context.Context) error {
 				return nil
 			}
 
-			var rx uint64
-			var tx uint64
-			for _, nw := range v.Networks {
-				atomic.AddUint64(&rx, nw.RxBytes)
-				atomic.AddUint64(&tx, nw.RxBytes)
-			}
-
 			st := environment.Stats{
 				Memory:      calculateDockerMemory(v.MemoryStats),
 				MemoryLimit: v.MemoryStats.Limit,
-				CpuAbsolute: calculateDockerAbsoluteCpu(&v.PreCPUStats, &v.CPUStats),
-				Network: struct {
-					RxBytes uint64 `json:"rx_bytes"`
-					TxBytes uint64 `json:"tx_bytes"`
-				}{
-					RxBytes: rx,
-					TxBytes: tx,
-				},
+				CpuAbsolute: calculateDockerAbsoluteCpu(v.PreCPUStats, v.CPUStats),
+				Network:     environment.NetworkStats{},
+			}
+
+			for _, nw := range v.Networks {
+				st.Network.RxBytes += nw.RxBytes
+				st.Network.TxBytes += nw.TxBytes
 			}
 
 			if b, err := json.Marshal(st); err != nil {
@@ -106,7 +94,7 @@ func calculateDockerMemory(stats types.MemoryStats) uint64 {
 // by the defined CPU limits on the container.
 //
 // @see https://github.com/docker/cli/blob/aa097cf1aa19099da70930460250797c8920b709/cli/command/container/stats_helpers.go#L166
-func calculateDockerAbsoluteCpu(pStats *types.CPUStats, stats *types.CPUStats) float64 {
+func calculateDockerAbsoluteCpu(pStats types.CPUStats, stats types.CPUStats) float64 {
 	// Calculate the change in CPU usage between the current and previous reading.
 	cpuDelta := float64(stats.CPUUsage.TotalUsage) - float64(pStats.CPUUsage.TotalUsage)
 
diff --git a/environment/stats.go b/environment/stats.go
index 8dbff47..966c50a 100644
--- a/environment/stats.go
+++ b/environment/stats.go
@@ -24,8 +24,10 @@ type Stats struct {
 	// Disk int64 `json:"disk_bytes"`
 
 	// Current network transmit in & out for a container.
-	Network struct {
-		RxBytes uint64 `json:"rx_bytes"`
-		TxBytes uint64 `json:"tx_bytes"`
-	} `json:"network"`
+	Network NetworkStats `json:"network"`
+}
+
+type NetworkStats struct {
+	RxBytes uint64 `json:"rx_bytes"`
+	TxBytes uint64 `json:"tx_bytes"`
 }