chore(qrm): refactor cpu/memory read-only state

pkg/agent/qrm-plugins/cpu/dynamicpolicy/policy.go

@@ -71,23 +71,6 @@ const (
 	healthCheckTolerationTimes = 3
 )
 
-var (
-	readonlyStateLock sync.RWMutex
-	readonlyState     state.ReadonlyState
-)
-
-// GetReadonlyState returns state.ReadonlyState to provides a way
-// to obtain the running states of the plugin
-func GetReadonlyState() (state.ReadonlyState, error) {
-	readonlyStateLock.RLock()
-	defer readonlyStateLock.RUnlock()
-
-	if readonlyState == nil {
-		return nil, fmt.Errorf("readonlyState isn't set")
-	}
-	return readonlyState, nil
-}
-
 // DynamicPolicy is the policy that's used by default;
 // it will consider the dynamic running information to calculate
 // and adjust resource requirements and configurations
@@ -150,9 +133,7 @@ func NewDynamicPolicy(agentCtx *agent.GenericContext, conf *config.Configuration
 		return false, agent.ComponentStub{}, fmt.Errorf("NewCheckpointState failed with error: %v", stateErr)
 	}
 
-	readonlyStateLock.Lock()
-	readonlyState = stateImpl
-	readonlyStateLock.Unlock()
+	state.SetReadonlyState(stateImpl)
 
 	wrappedEmitter := agentCtx.EmitterPool.GetDefaultMetricsEmitter().WithTags(agentName, metrics.MetricTag{
 		Key: util.QRMPluginPolicyTagName,

pkg/agent/qrm-plugins/cpu/dynamicpolicy/policy_test.go

@@ -4614,15 +4614,6 @@ func entriesMatch(entries1, entries2 state.PodEntries) (bool, error) {
 	return true, nil
 }
 
-func TestGetReadonlyState(t *testing.T) {
-	t.Parallel()
-
-	as := require.New(t)
-	readonlyState, err := GetReadonlyState()
-	as.NotNil(err)
-	as.Nil(readonlyState)
-}
-
 func TestClearResidualState(t *testing.T) {
 	t.Parallel()
 

pkg/agent/qrm-plugins/cpu/dynamicpolicy/state/state.go

@@ -19,6 +19,7 @@ package state
 import (
 	"encoding/json"
 	"fmt"
+	"sync"
 
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/util/sets"
@@ -562,3 +563,29 @@ type ReadonlyState interface {
 }
 
 type GenerateMachineStateFromPodEntriesFunc func(topology *machine.CPUTopology, podEntries PodEntries) (NUMANodeMap, error)
+
+var (
+	readonlyStateLock sync.RWMutex
+	readonlyState     ReadonlyState
+)
+
+// GetReadonlyState returns state.ReadonlyState to provides a way
+// to obtain the running states of the plugin
+func GetReadonlyState() (ReadonlyState, error) {
+	readonlyStateLock.RLock()
+	defer readonlyStateLock.RUnlock()
+
+	if readonlyState == nil {
+		return nil, fmt.Errorf("readonlyState isn't set")
+	}
+	return readonlyState, nil
+}
+
+// SetReadonlyState sets state.ReadonlyState to be used by the plugin
+// this function should only be called once during initialization
+func SetReadonlyState(state ReadonlyState) {
+	readonlyStateLock.Lock()
+	defer readonlyStateLock.Unlock()
+
+	readonlyState = state
+}

pkg/agent/qrm-plugins/cpu/dynamicpolicy/state/state_test.go

@@ -3233,3 +3233,12 @@ func TestGetAvailableCPUQuantity(t *testing.T) {
 	cpuQuantity := int(math.Ceil(nodeState.GetAvailableCPUQuantity(machine.NewCPUSet())))
 	require.Equal(t, 15, cpuQuantity)
 }
+
+func TestGetReadonlyState(t *testing.T) {
+	t.Parallel()
+
+	as := require.New(t)
+	state, err := GetReadonlyState()
+	as.NotNil(err)
+	as.Nil(state)
+}

pkg/agent/qrm-plugins/cpu/nativepolicy/policy.go

@@ -55,23 +55,6 @@ const (
 	maxResidualTime  = 5 * time.Minute
 )
 
-var (
-	readonlyStateLock sync.RWMutex
-	readonlyState     state.ReadonlyState
-)
-
-// GetReadonlyState returns state.ReadonlyState to provides a way
-// to obtain the running states of the plugin
-func GetReadonlyState() (state.ReadonlyState, error) {
-	readonlyStateLock.RLock()
-	defer readonlyStateLock.RUnlock()
-
-	if readonlyState == nil {
-		return nil, fmt.Errorf("readonlyState isn't setted")
-	}
-	return readonlyState, nil
-}
-
 // NativePolicy is a policy compatible with Kubernetes native semantics and is used in topology-aware scheduling scenarios.
 type NativePolicy struct {
 	sync.RWMutex
@@ -115,9 +98,7 @@ func NewNativePolicy(agentCtx *agent.GenericContext, conf *config.Configuration,
 		return false, agent.ComponentStub{}, fmt.Errorf("NewCheckpointState failed with error: %v", stateErr)
 	}
 
-	readonlyStateLock.Lock()
-	readonlyState = stateImpl
-	readonlyStateLock.Unlock()
+	state.SetReadonlyState(stateImpl)
 
 	wrappedEmitter := agentCtx.EmitterPool.GetDefaultMetricsEmitter().WithTags(agentName, metrics.MetricTag{
 		Key: util.QRMPluginPolicyTagName,

pkg/agent/qrm-plugins/cpu/nativepolicy/policy_test.go

@@ -390,15 +390,6 @@ func TestAllocateForPod(t *testing.T) {
 	_ = os.RemoveAll(tmpDir)
 }
 
-func TestGetReadonlyState(t *testing.T) {
-	t.Parallel()
-
-	as := require.New(t)
-	readonlyState, err := GetReadonlyState()
-	as.NotNil(err)
-	as.Nil(readonlyState)
-}
-
 func TestClearResidualState(t *testing.T) {
 	t.Parallel()
 

pkg/agent/qrm-plugins/memory/dynamicpolicy/policy.go

@@ -90,23 +90,6 @@ const (
 	movePagesWorkLimit    = 2
 )
 
-var (
-	readonlyStateLock sync.RWMutex
-	readonlyState     state.ReadonlyState
-)
-
-// GetReadonlyState returns state.ReadonlyState to provides a way
-// to obtain the running states of the plugin
-func GetReadonlyState() (state.ReadonlyState, error) {
-	readonlyStateLock.RLock()
-	defer readonlyStateLock.RUnlock()
-
-	if readonlyState == nil {
-		return nil, fmt.Errorf("readonlyState isn't setted")
-	}
-	return readonlyState, nil
-}
-
 type DynamicPolicy struct {
 	sync.RWMutex
 
@@ -193,9 +176,7 @@ func NewDynamicPolicy(agentCtx *agent.GenericContext, conf *config.Configuration
 		general.Infof("empty ExtraControlKnobConfigFile, initialize empty extraControlKnobConfigs")
 	}
 
-	readonlyStateLock.Lock()
-	readonlyState = stateImpl
-	readonlyStateLock.Unlock()
+	state.SetReadonlyState(stateImpl)
 
 	wrappedEmitter := agentCtx.EmitterPool.GetDefaultMetricsEmitter().WithTags(agentName, metrics.MetricTag{
 		Key: util.QRMPluginPolicyTagName,

pkg/agent/qrm-plugins/memory/dynamicpolicy/policy_test.go

@@ -1953,16 +1953,6 @@ func TestGetResourcesAllocation(t *testing.T) {
 	}, resp5.PodResources[req.PodUid].ContainerResources[testName].ResourceAllocation[string(v1.ResourceMemory)])
 }
 
-func TestGetReadonlyState(t *testing.T) {
-	t.Parallel()
-
-	as := require.New(t)
-	readonlyState, err := GetReadonlyState()
-	if readonlyState == nil {
-		as.NotNil(err)
-	}
-}
-
 func TestGenerateResourcesMachineStateFromPodEntries(t *testing.T) {
 	t.Parallel()
 

pkg/agent/qrm-plugins/memory/dynamicpolicy/state/state.go

@@ -19,6 +19,7 @@ package state
 import (
 	"encoding/json"
 	"fmt"
+	"sync"
 
 	info "github.com/google/cadvisor/info/v1"
 	v1 "k8s.io/api/core/v1"
@@ -384,3 +385,29 @@ type State interface {
 	writer
 	ReadonlyState
 }
+
+var (
+	readonlyStateLock sync.RWMutex
+	readonlyState     ReadonlyState
+)
+
+// GetReadonlyState returns state.ReadonlyState to provides a way
+// to obtain the running states of the plugin
+func GetReadonlyState() (ReadonlyState, error) {
+	readonlyStateLock.RLock()
+	defer readonlyStateLock.RUnlock()
+
+	if readonlyState == nil {
+		return nil, fmt.Errorf("readonlyState isn't setted")
+	}
+	return readonlyState, nil
+}
+
+// SetReadonlyState sets state.ReadonlyState to be used by the plugin
+// this function should only be called once during initialization
+func SetReadonlyState(state ReadonlyState) {
+	readonlyStateLock.Lock()
+	defer readonlyStateLock.Unlock()
+
+	readonlyState = state
+}

pkg/agent/qrm-plugins/memory/dynamicpolicy/state/state_test.go

@@ -0,0 +1,33 @@
+/*
+Copyright 2022 The Katalyst Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package state
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestGetReadonlyState(t *testing.T) {
+	t.Parallel()
+
+	as := require.New(t)
+	state, err := GetReadonlyState()
+	if state == nil {
+		as.NotNil(err)
+	}
+}