diff --git a/pkg/scheduler/framework/cyclestate.go b/pkg/scheduler/framework/cyclestate.go new file mode 100644 index 000000000..07a10a421 --- /dev/null +++ b/pkg/scheduler/framework/cyclestate.go @@ -0,0 +1,58 @@ +/* +Copyright (c) Microsoft Corporation. +Licensed under the MIT license. +*/ + +package framework + +import ( + "fmt" + "sync" +) + +// StateKey is the key for a state value stored in a CycleState. +type StateKey string + +// StateValue is the value stored in a CycleState under a specific key. +type StateValue interface{} + +// CycleStatePluginReadWriter is an interface through which plugins can store and retrieve data. +type CycleStatePluginReadWriter interface { + Read(key StateKey) (StateValue, error) + Write(key StateKey, val StateValue) + Delete(key StateKey) +} + +// CycleState is, similar to its namesake in kube-scheduler, provides a way for plugins to +// store and retrieve arbitrary data during a scheduling cycle. The scheduler also uses +// this struct to keep some global states during a scheduling cycle; note that these +// state are only accessible to the scheduler itself, not to plugins. +// +// It uses a sync.Map for concurrency-safe storage. +type CycleState struct { + // store is a concurrency-safe store (a map). + store sync.Map +} + +// Read retrieves a value from CycleState by a key. +func (c *CycleState) Read(key StateKey) (StateValue, error) { + if v, ok := c.store.Load(key); ok { + return v, nil + } + return nil, fmt.Errorf("key %s is not found", key) +} + +// Write stores a value in CycleState under a key. +func (c *CycleState) Write(key StateKey, val StateValue) { + c.store.Store(key, val) +} + +// Delete deletes a key from CycleState. +func (c *CycleState) Delete(key StateKey) { + c.store.Delete(key) +} + +// NewCycleState creates a CycleState. +func NewCycleState() *CycleState { + return &CycleState{} +} diff --git a/pkg/scheduler/framework/cyclestate_test.go b/pkg/scheduler/framework/cyclestate_test.go new file mode 100644 index 000000000..b39e659ed --- /dev/null +++ b/pkg/scheduler/framework/cyclestate_test.go @@ -0,0 +1,23 @@ +/* +Copyright (c) Microsoft Corporation. +Licensed under the MIT license. +*/ + +package framework + +import "testing" + +// TestCycleStateBasicOps tests the basic ops (Read, Write, and Delete) of a CycleState. +func TestCycleStateBasicOps(t *testing.T) { + cs := NewCycleState() + + k, v := "key", "value" + cs.Write(StateKey(k), StateValue(v)) + if out, err := cs.Read("key"); out != "value" || err != nil { + t.Fatalf("Read(%v) = %v, %v, want %v, nil", k, out, err, v) + } + cs.Delete(StateKey(k)) + if out, err := cs.Read("key"); out != nil || err == nil { + t.Fatalf("Read(%v) = %v, %v, want nil, not found error", k, out, err) + } +}