diff --git a/environment/docker/environment.go b/environment/docker/environment.go index 5f11736..1ab4b49 100644 --- a/environment/docker/environment.go +++ b/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 - stMu sync.RWMutex + State system.AtomicString } // 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 } diff --git a/environment/docker/power.go b/environment/docker/power.go index 3a4a071..56d26ec 100644 --- a/environment/docker/power.go +++ b/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) } diff --git a/environment/docker/state.go b/environment/docker/state.go index 51fbf8f..50d4c2f 100644 --- a/environment/docker/state.go +++ b/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.st = state - e.stMu.Unlock() - - e.Events().Publish(environment.StateChangeEvent, e.State()) + e.State.Store(state) + e.Events().Publish(environment.StateChangeEvent, state) } return nil diff --git a/environment/docker/stats.go b/environment/docker/stats.go index 84b0034..01fbd3c 100644 --- a/environment/docker/stats.go +++ b/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 } diff --git a/system/bool.go b/system/bool.go deleted file mode 100644 index 5f30367..0000000 --- a/system/bool.go +++ /dev/null @@ -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 -} diff --git a/system/utils.go b/system/utils.go new file mode 100644 index 0000000..8a9008e --- /dev/null +++ b/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 "" +}