Less obtuse logic for polling resource usage when attaching a container

environment/docker/container.go

@@ -30,6 +30,10 @@ type imagePullStatus struct {
 // of the process stream. This should not be used for reading console data as you *will*
 // miss important output at the beginning because of the time delay with attaching to the
 // output.
+//
+// Calling this function will poll resources for the container in the background until the
+// provided context is canceled by the caller. Failure to cancel said context will cause
+// background memory leaks as the goroutine will not exit.
 func (e *Environment) Attach() error {
 	if e.IsAttached() {
 		return nil
@@ -53,38 +57,43 @@ func (e *Environment) Attach() error {
 		e.SetStream(&st)
 	}
 
-	c := new(Console)
+	go func() {
-	go func(console *Console) {
 		ctx, cancel := context.WithCancel(context.Background())
-
 		defer cancel()
-		defer e.stream.Close()
 		defer func() {
+			e.stream.Close()
 			e.SetState(environment.ProcessOfflineState)
 			e.SetStream(nil)
 		}()
 
-		// Poll resources in a separate thread since this will block the copy call below
+		go func() {
-		// from being reached until it is completed if not run in a separate process. However,
-		// we still want it to be stopped when the copy operation below is finished running which
-		// indicates that the container is no longer running.
-		go func(ctx context.Context) {
 			if err := e.pollResources(ctx); err != nil {
-				l := log.WithField("environment_id", e.Id)
 				if !errors.Is(err, context.Canceled) {
-					l.WithField("error", err).Error("error during environment resource polling")
+					e.log().WithField("error", err).Error("error during environment resource polling")
 				} else {
-					l.Warn("stopping server resource polling: context canceled")
+					e.log().Warn("stopping server resource polling: context canceled")
 				}
 			}
-		}(ctx)
+		}()
 
-		// Stream the reader output to the console which will then fire off events and handle console
+		ok, err := e.client.ContainerWait(ctx, e.Id, container.WaitConditionNotRunning)
-		// throttling and sending the output to the user.
+		select {
-		if _, err := io.Copy(console, e.stream.Reader); err != nil {
+		case <-ctx.Done():
-			log.WithField("environment_id", e.Id).WithField("error", err).Error("error while copying environment output to console")
+			// Do nothing, the context was canceled by a different process, there is no error
+			// to report at this point.
+			e.log().Debug("terminating ContainerWait blocking process, context canceled")
+			return
+		case _ = <-err:
+			// An error occurred with the ContainerWait call, report it here and then hope
+			// for the fucking best I guess?
+			e.log().WithField("error", err).Error("error while blocking using ContainerWait")
+			return
+		case <-ok:
+			// Do nothing, everything is running as expected. This will allow us to keep
+			// blocking the termination of this function until the container stops at which
+			// point all of our deferred functions can run.
 		}
-	}(c)
+	}()
 
 	return nil
 }
@@ -280,7 +289,6 @@ func (e *Environment) followOutput() error {
 		if err != nil {
 			return err
 		}
-
 		return errors.New(fmt.Sprintf("no such container: %s", e.Id))
 	}
 

environment/docker/environment.go

@@ -2,6 +2,7 @@ package docker
 
 import (
 	"context"
+	"github.com/apex/log"
 	"github.com/docker/docker/api/types"
 	"github.com/docker/docker/client"
 	"github.com/pterodactyl/wings/api"
@@ -70,6 +71,10 @@ func New(id string, m *Metadata, c *environment.Configuration) (*Environment, er
 	return e, nil
 }
 
+func (e *Environment) log() *log.Entry {
+	return log.WithField("environment", e.Type()).WithField("container_id", e.Id)
+}
+
 func (e *Environment) Type() string {
 	return "docker"
 }

environment/docker/stats.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"emperror.dev/errors"
 	"encoding/json"
-	"github.com/apex/log"
 	"github.com/docker/docker/api/types"
 	"github.com/pterodactyl/wings/environment"
 	"io"
@@ -19,11 +18,10 @@ func (e *Environment) pollResources(ctx context.Context) error {
 		return errors.New("cannot enable resource polling on a stopped server")
 	}
 
-	l := log.WithField("container_id", e.Id)
+	e.log().Info("starting resource polling for container")
-	l.Debug("starting resource polling for container")
+	defer e.log().Debug("stopped resource polling for container")
-	defer l.Debug("stopped resource polling for container")
 
-	stats, err := e.client.ContainerStats(context.Background(), e.Id, true)
+	stats, err := e.client.ContainerStats(ctx, e.Id, true)
 	if err != nil {
 		return err
 	}
@@ -39,10 +37,10 @@ func (e *Environment) pollResources(ctx context.Context) error {
 			var v *types.StatsJSON
 
 			if err := dec.Decode(&v); err != nil {
-				if err != io.EOF {
+				if err != io.EOF && !errors.Is(err, context.Canceled) {
-					l.WithField("error", err).Warn("error while processing Docker stats output for container")
+					e.log().WithField("error", err).Warn("error while processing Docker stats output for container")
 				} else {
-					l.Debug("io.EOF encountered during stats decode, stopping polling...")
+					e.log().Debug("io.EOF encountered during stats decode, stopping polling...")
 				}
 
 				return nil
@@ -50,7 +48,7 @@ func (e *Environment) pollResources(ctx context.Context) error {
 
 			// Disable collection if the server is in an offline state and this process is still running.
 			if e.st.Load() == environment.ProcessOfflineState {
-				l.Debug("process in offline state while resource polling is still active; stopping poll")
+				e.log().Debug("process in offline state while resource polling is still active; stopping poll")
 				return nil
 			}
 
@@ -75,7 +73,7 @@ func (e *Environment) pollResources(ctx context.Context) error {
 			}
 
 			if b, err := json.Marshal(st); err != nil {
-				l.WithField("error", err).Warn("error while marshaling stats object for environment")
+				e.log().WithField("error", err).Warn("error while marshaling stats object for environment")
 			} else {
 				e.Events().Publish(environment.ResourceEvent, string(b))
 			}