Create atomic string to allow for simpler logic

environment/docker/environment.go

@@ -8,6 +8,7 @@ import (
 	"github.com/pterodactyl/wings/api"
 	"github.com/pterodactyl/wings/environment"
 	"github.com/pterodactyl/wings/events"
+	"github.com/pterodactyl/wings/system"
 	"io"
 	"sync"
 )
@@ -47,8 +48,7 @@ type Environment struct {
 	emitter *events.EventBus
 
 	// Tracks the environment state.
-	st   string
+	State system.AtomicString
-	stMu sync.RWMutex
 }
 
 // Creates a new base Docker environment. The ID passed through will be the ID that is used to
@@ -65,9 +65,10 @@ func New(id string, m *Metadata, c *environment.Configuration) (*Environment, er
 		Configuration: c,
 		meta:          m,
 		client:        cli,
-		st:            environment.ProcessOfflineState,
 	}
 
+	e.State.Store(environment.ProcessOfflineState)
+
 	return e, nil
 }
 

environment/docker/power.go

@@ -136,7 +136,7 @@ func (e *Environment) Stop() error {
 
 	// If the process is already offline don't switch it back to stopping. Just leave it how
 	// it is and continue through to the stop handling for the process.
-	if e.State() != environment.ProcessOfflineState {
+	if e.State.Load() != environment.ProcessOfflineState {
 		e.setState(environment.ProcessStoppingState)
 	}
 
@@ -217,7 +217,7 @@ func (e *Environment) Terminate(signal os.Signal) error {
 		// If the container is not running but we're not already in a stopped state go ahead
 		// and update things to indicate we should be completely stopped now. Set to stopping
 		// first so crash detection is not triggered.
-		if e.State() != environment.ProcessOfflineState {
+		if e.State.Load() != environment.ProcessOfflineState {
 			e.setState(environment.ProcessStoppingState)
 			e.setState(environment.ProcessOfflineState)
 		}

environment/docker/state.go

@@ -6,14 +6,6 @@ import (
 	"github.com/pterodactyl/wings/environment"
 )
 
-// Returns the current environment state.
-func (e *Environment) State() string {
-	e.stMu.RLock()
-	defer e.stMu.RUnlock()
-
-	return e.st
-}
-
 // Sets the state of the environment. This emits an event that server's can hook into to
 // take their own actions and track their own state based on the environment.
 func (e *Environment) setState(state string) error {
@@ -25,16 +17,13 @@ func (e *Environment) setState(state string) error {
 	}
 
 	// Get the current state of the environment before changing it.
-	prevState := e.State()
+	prevState := e.State.Load()
 
 	// Emit the event to any listeners that are currently registered.
 	if prevState != state {
 		// If the state changed make sure we update the internal tracking to note that.
-		e.stMu.Lock()
+		e.State.Store(state)
-		e.st = state
+		e.Events().Publish(environment.StateChangeEvent, state)
-		e.stMu.Unlock()
-
-		e.Events().Publish(environment.StateChangeEvent, e.State())
 	}
 
 	return nil

environment/docker/stats.go

@@ -20,7 +20,7 @@ func (e *Environment) pollResources(ctx context.Context) error {
 	l.Debug("starting resource polling for container")
 	defer l.Debug("stopped resource polling for container")
 
-	if e.State() == environment.ProcessOfflineState {
+	if e.State.Load() == environment.ProcessOfflineState {
 		return errors.New("cannot enable resource polling on a stopped server")
 	}
 
@@ -50,7 +50,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.State() == environment.ProcessOfflineState {
+			if e.State.Load() == environment.ProcessOfflineState {
 				l.Debug("process in offline state while resource polling is still active; stopping poll")
 				return nil
 			}

system/bool.go

@@ -1,20 +0,0 @@
-package system
-
-import "sync/atomic"
-
-type AtomicBool struct {
-	flag uint32
-}
-
-func (ab *AtomicBool) Set(v bool) {
-	i := 0
-	if v {
-		i = 1
-	}
-
-	atomic.StoreUint32(&ab.flag, uint32(i))
-}
-
-func (ab *AtomicBool) Get() bool {
-	return atomic.LoadUint32(&ab.flag) == 1
-}

system/utils.go

@@ -0,0 +1,51 @@
+package system
+
+import (
+	"sync/atomic"
+)
+
+type AtomicBool struct {
+	flag uint32
+}
+
+func (ab *AtomicBool) Set(v bool) {
+	i := 0
+	if v {
+		i = 1
+	}
+
+	atomic.StoreUint32(&ab.flag, uint32(i))
+}
+
+func (ab *AtomicBool) Get() bool {
+	return atomic.LoadUint32(&ab.flag) == 1
+}
+
+// AtomicString allows for reading/writing to a given struct field without having to worry
+// about a potential race condition scenario. Under the hood it uses a simple sync.RWMutex
+// to control access to the value.
+type AtomicString struct {
+	v atomic.Value
+}
+
+// Returns a new instance of an AtomicString.
+func NewAtomicString(v string) *AtomicString {
+	as := &AtomicString{}
+	if v != "" {
+		as.Store(v)
+	}
+	return as
+}
+
+// Stores the string value passed atomically.
+func (as *AtomicString) Store(v string) {
+	as.v.Store(v)
+}
+
+// Loads the string value and returns it.
+func (as *AtomicString) Load() string {
+	if v := as.v.Load(); v != nil {
+		return v.(string)
+	}
+	return ""
+}