Create atomic string to allow for simpler logic

This commit is contained in:
Dane Everitt 2020-11-06 21:14:29 -08:00
parent a74be8f4eb
commit 3fce1b98d5
No known key found for this signature in database
GPG Key ID: EEA66103B3D71F53
6 changed files with 62 additions and 41 deletions

View File

@ -8,6 +8,7 @@ import (
"github.com/pterodactyl/wings/api" "github.com/pterodactyl/wings/api"
"github.com/pterodactyl/wings/environment" "github.com/pterodactyl/wings/environment"
"github.com/pterodactyl/wings/events" "github.com/pterodactyl/wings/events"
"github.com/pterodactyl/wings/system"
"io" "io"
"sync" "sync"
) )
@ -47,8 +48,7 @@ type Environment struct {
emitter *events.EventBus emitter *events.EventBus
// Tracks the environment state. // 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 // 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, Configuration: c,
meta: m, meta: m,
client: cli, client: cli,
st: environment.ProcessOfflineState,
} }
e.State.Store(environment.ProcessOfflineState)
return e, nil return e, nil
} }

View File

@ -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 // 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. // 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) 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 // 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 // and update things to indicate we should be completely stopped now. Set to stopping
// first so crash detection is not triggered. // 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.ProcessStoppingState)
e.setState(environment.ProcessOfflineState) e.setState(environment.ProcessOfflineState)
} }

View File

@ -6,14 +6,6 @@ import (
"github.com/pterodactyl/wings/environment" "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 // 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. // take their own actions and track their own state based on the environment.
func (e *Environment) setState(state string) error { 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. // 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. // Emit the event to any listeners that are currently registered.
if prevState != state { if prevState != state {
// If the state changed make sure we update the internal tracking to note that. // 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 return nil

View File

@ -20,7 +20,7 @@ func (e *Environment) pollResources(ctx context.Context) error {
l.Debug("starting resource polling for container") l.Debug("starting resource polling for container")
defer l.Debug("stopped 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") 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. // 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") l.Debug("process in offline state while resource polling is still active; stopping poll")
return nil return nil
} }

View File

@ -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
}

51
system/utils.go Normal file
View File

@ -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 ""
}