feat: refactor logic into new file

feat: refactor logic into new file

feat: refactor logic into new file

Author: JustinChengLZ

cmd/katalyst-agent/app/options/qrm/statedirectory/statedirectory_base.go

@@ -26,13 +26,11 @@ type StateDirectoryOptions struct {
 	StateFileDirectory         string
 	InMemoryStateFileDirectory string
 	EnableInMemoryState        bool
-	HasPreStop                 bool
 }
 
 func NewStateDirectoryOptions() *StateDirectoryOptions {
 	return &StateDirectoryOptions{
-		StateFileDirectory:         "/var/lib/katalyst/qrm_advisor",
-		InMemoryStateFileDirectory: "/dev/shm/qrm/state",
+		StateFileDirectory: "/var/lib/katalyst/qrm_advisor",
 	}
 }
 
@@ -44,13 +42,11 @@ func (o *StateDirectoryOptions) AddFlags(fss *cliflag.NamedFlagSets) {
 		o.InMemoryStateFileDirectory, "The in memory directory to store the state file.")
 	fs.BoolVar(&o.EnableInMemoryState, "qrm-enable-in-memory-state",
 		o.EnableInMemoryState, "if set true, the state will be stored in the in-memory directory.")
-	fs.BoolVar(&o.HasPreStop, "qrm-has-pre-stop", o.HasPreStop, "if set true, we will be able to migrate state files")
 }
 
 func (o *StateDirectoryOptions) ApplyTo(conf *statedirectory.StateDirectoryConfiguration) error {
 	conf.StateFileDirectory = o.StateFileDirectory
 	conf.InMemoryStateFileDirectory = o.InMemoryStateFileDirectory
 	conf.EnableInMemoryState = o.EnableInMemoryState
-	conf.HasPreStop = o.HasPreStop
 	return nil
 }

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

@@ -27,8 +27,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/kubewharf/katalyst-core/pkg/config/agent/qrm/statedirectory"
-
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
@@ -59,6 +57,7 @@ import (
 	"github.com/kubewharf/katalyst-core/pkg/agent/utilcomponent/featuregatenegotiation"
 	"github.com/kubewharf/katalyst-core/pkg/config"
 	"github.com/kubewharf/katalyst-core/pkg/config/agent/dynamic"
+	"github.com/kubewharf/katalyst-core/pkg/config/agent/qrm/statedirectory"
 	"github.com/kubewharf/katalyst-core/pkg/config/generic"
 	"github.com/kubewharf/katalyst-core/pkg/metaserver"
 	"github.com/kubewharf/katalyst-core/pkg/metaserver/agent"
@@ -95,7 +94,9 @@ func generateSharedNumaBindingPoolAllocationMeta(poolName string) commonstate.Al
 	return meta
 }
 
-func getTestDynamicPolicyWithInitialization(topology *machine.CPUTopology, stateFileDirectory string) (*DynamicPolicy, error) {
+func getTestDynamicPolicyWithInitialization(
+	topology *machine.CPUTopology, stateFileDirectory string,
+) (*DynamicPolicy, error) {
 	dynamicPolicy, err := getTestDynamicPolicyWithoutInitialization(topology, stateFileDirectory)
 	if err != nil {
 		return nil, err
@@ -115,7 +116,9 @@ func getTestDynamicPolicyWithInitialization(topology *machine.CPUTopology, state
 	return dynamicPolicy, nil
 }
 
-func getTestDynamicPolicyWithoutInitialization(topology *machine.CPUTopology, stateFileDirectory string) (*DynamicPolicy, error) {
+func getTestDynamicPolicyWithoutInitialization(
+	topology *machine.CPUTopology, stateFileDirectory string,
+) (*DynamicPolicy, error) {
 	stateDirectoryConfig := &statedirectory.StateDirectoryConfiguration{
 		StateFileDirectory: stateFileDirectory,
 	}
@@ -6212,7 +6215,9 @@ type mockCPUAdvisor struct {
 	advisorapi.CPUAdvisorServer
 }
 
-func (m *mockCPUAdvisor) GetAdvice(ctx context.Context, in *advisorapi.GetAdviceRequest) (*advisorapi.GetAdviceResponse, error) {
+func (m *mockCPUAdvisor) GetAdvice(
+	ctx context.Context, in *advisorapi.GetAdviceRequest,
+) (*advisorapi.GetAdviceResponse, error) {
 	args := m.Called(ctx, in)
 	return args.Get(0).(*advisorapi.GetAdviceResponse), args.Error(1)
 }
@@ -6222,11 +6227,15 @@ func (m *mockCPUAdvisor) ListAndWatch(in *advisorsvc.Empty, srv advisorapi.CPUAd
 	return args.Error(0)
 }
 
-func (m *mockCPUAdvisor) AddContainer(ctx context.Context, req *advisorsvc.ContainerMetadata) (*advisorsvc.AddContainerResponse, error) {
+func (m *mockCPUAdvisor) AddContainer(
+	ctx context.Context, req *advisorsvc.ContainerMetadata,
+) (*advisorsvc.AddContainerResponse, error) {
 	return &advisorsvc.AddContainerResponse{}, nil
 }
 
-func (m *mockCPUAdvisor) RemovePod(ctx context.Context, req *advisorsvc.RemovePodRequest) (*advisorsvc.RemovePodResponse, error) {
+func (m *mockCPUAdvisor) RemovePod(
+	ctx context.Context, req *advisorsvc.RemovePodRequest,
+) (*advisorsvc.RemovePodResponse, error) {
 	return &advisorsvc.RemovePodResponse{}, nil
 }
 

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

@@ -19,20 +19,18 @@ package state
 import (
 	"fmt"
 	"path"
-	"path/filepath"
 	"reflect"
 	"sync"
 	"time"
 
 	"k8s.io/klog/v2"
-	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager"
 	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager/errors"
 
 	"github.com/kubewharf/katalyst-core/pkg/config/agent/qrm/statedirectory"
 	"github.com/kubewharf/katalyst-core/pkg/metrics"
-	"github.com/kubewharf/katalyst-core/pkg/util/file"
 	"github.com/kubewharf/katalyst-core/pkg/util/general"
 	"github.com/kubewharf/katalyst-core/pkg/util/machine"
+	"github.com/kubewharf/katalyst-core/pkg/util/qrmcheckpointmanager"
 )
 
 const (
@@ -44,15 +42,16 @@ const (
 // go to in-memory State, and then go to disk State, i.e. in write-back mode
 type stateCheckpoint struct {
 	sync.RWMutex
-	cache             *cpuPluginState
-	policyName        string
-	checkpointManager checkpointmanager.CheckpointManager
-	checkpointName    string
+	cache                *cpuPluginState
+	policyName           string
+	qrmCheckpointManager *qrmcheckpointmanager.QRMCheckpointManager
+	checkpointName       string
 	// when we add new properties to checkpoint,
 	// it will cause checkpoint corruption, and we should skip it
 	skipStateCorruption                bool
 	GenerateMachineStateFromPodEntries GenerateMachineStateFromPodEntriesFunc
 	emitter                            metrics.MetricEmitter
+	hasPreStop                         bool
 }
 
 var _ State = &stateCheckpoint{}
@@ -64,23 +63,24 @@ func NewCheckpointState(
 	emitter metrics.MetricEmitter,
 ) (State, error) {
 	currentStateDir, otherStateDir := stateDirectoryConfig.GetCurrentAndOtherStateFileDirectory()
-	hasPreStop := stateDirectoryConfig.HasPreStop
-	checkpointManager, err := checkpointmanager.NewCheckpointManager(currentStateDir)
+	// If there is an empty otherStateDir, this means that there is no pre-stop script in place
+	hasPreStop := otherStateDir != ""
+	qrmCheckpointManager, err := qrmcheckpointmanager.NewQRMCheckpointManager(currentStateDir, otherStateDir, checkpointName, "cpu_plugin")
 	if err != nil {
 		return nil, fmt.Errorf("failed to initialize checkpoint manager: %v", err)
 	}
-
 	sc := &stateCheckpoint{
 		cache:                              NewCPUPluginState(topology),
 		policyName:                         policyName,
-		checkpointManager:                  checkpointManager,
+		qrmCheckpointManager:               qrmCheckpointManager,
 		checkpointName:                     checkpointName,
 		skipStateCorruption:                skipStateCorruption,
 		GenerateMachineStateFromPodEntries: generateMachineStateFunc,
 		emitter:                            emitter,
+		hasPreStop:                         hasPreStop,
 	}
 
-	if err := sc.restoreState(currentStateDir, otherStateDir, hasPreStop, topology); err != nil {
+	if err := sc.restoreState(topology); err != nil {
 		return nil, fmt.Errorf("could not restore state from checkpoint: %v, please drain this node and delete "+
 			"the cpu plugin checkpoint file %q before restarting Kubelet", err, path.Join(currentStateDir, checkpointName))
 	}
@@ -90,19 +90,17 @@ func NewCheckpointState(
 
 // restoreState is first done by searching the current directory for the state file.
 // If it does not exist, we search the other directory for the state file and try to migrate the state file over to the current directory.
-func (sc *stateCheckpoint) restoreState(
-	currentStateDir, otherStateDir string, hasPreStop bool, topology *machine.CPUTopology,
-) error {
+func (sc *stateCheckpoint) restoreState(topology *machine.CPUTopology) error {
 	sc.Lock()
 	defer sc.Unlock()
 	var err error
 	var foundAndSkippedStateCorruption bool
 
 	checkpoint := NewCPUPluginCheckpoint()
-	if err = sc.checkpointManager.GetCheckpoint(sc.checkpointName, checkpoint); err != nil {
+	if err = sc.qrmCheckpointManager.GetCurrentCheckpoint(sc.checkpointName, checkpoint, true); err != nil {
 		if err == errors.ErrCheckpointNotFound {
 			// We cannot find checkpoint, so it is possible that previous checkpoint was stored in either disk or memory
-			return sc.tryMigrateState(topology, currentStateDir, otherStateDir, hasPreStop, checkpoint)
+			return sc.tryMigrateState(topology, checkpoint)
 		} else if err == errors.ErrCorruptCheckpoint {
 			if !sc.skipStateCorruption {
 				return err
@@ -155,27 +153,20 @@ func (sc *stateCheckpoint) populateCacheAndState(
 // tryMigrateState tries to migrate the state file from the other directory to current directory.
 // If the other directory does not have a state file, then we build a new checkpoint.
 func (sc *stateCheckpoint) tryMigrateState(
-	topology *machine.CPUTopology, currentStateDir, otherStateDir string, hasPreStop bool,
-	checkpoint *CPUPluginCheckpoint,
+	topology *machine.CPUTopology, checkpoint *CPUPluginCheckpoint,
 ) error {
 	var foundAndSkippedStateCorruption bool
 	klog.Infof("[cpu_plugin] trying to migrate state")
 
 	// Do not migrate and build new checkpoint if there is no pre-stop script
-	if !hasPreStop {
+	if !sc.hasPreStop {
 		return sc.storeState()
 	}
 
-	// Get the old checkpoint using the provided file directory
-	oldCheckpointManager, err := checkpointmanager.NewCheckpointManager(otherStateDir)
-	if err != nil {
-		return fmt.Errorf("[cpu_plugin] failed to initialize old checkpoint manager for migration: %v", err)
-	}
-
-	if err = oldCheckpointManager.GetCheckpoint(sc.checkpointName, checkpoint); err != nil {
+	if err := sc.qrmCheckpointManager.GetPreviousCheckpoint(sc.checkpointName, checkpoint); err != nil {
 		if err == errors.ErrCheckpointNotFound {
 			// Old checkpoint file is not found, so we just store state in new checkpoint
-			general.Infof("[cpu_plugin] checkpoint %v doesn't exist in dir %v, create it", sc.checkpointName, otherStateDir)
+			general.Infof("[cpu_plugin] checkpoint %v doesn't exist, create it", sc.checkpointName)
 			return sc.storeState()
 		} else if err == errors.ErrCorruptCheckpoint {
 			if !sc.skipStateCorruption {
@@ -188,41 +179,23 @@ func (sc *stateCheckpoint) tryMigrateState(
 		}
 	}
 
-	if err = sc.populateCacheAndState(topology, checkpoint, foundAndSkippedStateCorruption); err != nil {
+	if err := sc.populateCacheAndState(topology, checkpoint, foundAndSkippedStateCorruption); err != nil {
 		return fmt.Errorf("[cpu_plugin] failed to populate checkpoint state during state migration: %v", err)
 	}
 
 	// always store state after migrating to new checkpoint
-	if err = sc.storeState(); err != nil {
+	if err := sc.storeState(); err != nil {
 		return fmt.Errorf("[cpu_plugin] failed to store checkpoint state during end of migration: %v", err)
 	}
 
-	// validate that the two files are equal
-	equal, err := sc.checkpointFilesEqual(currentStateDir, otherStateDir)
-	if err != nil {
-		return fmt.Errorf("[cpu_plugin] failed to compare checkpoint files: %v", err)
-	}
-	if !equal {
-		klog.Infof("[cpu_plugin] checkpoint files are not equal, migration failed, fall back to old checkpoint")
-		sc.checkpointManager = oldCheckpointManager
-		return nil
-	}
-
-	// remove old checkpoint file
-	if err = oldCheckpointManager.RemoveCheckpoint(sc.checkpointName); err != nil {
-		return fmt.Errorf("[cpu_plugin] failed to remove old checkpoint: %v", err)
+	if err := sc.qrmCheckpointManager.ValidateCheckpointFilesMigration(); err != nil {
+		return fmt.Errorf("[cpu_plugin] ValidateCheckpointFilesMigration failed with error: %v", err)
 	}
 
 	klog.Infof("[cpu_plugin] migrate checkpoint succeeded")
 	return nil
 }
 
-func (sc *stateCheckpoint) checkpointFilesEqual(currentStateDir, otherStateDir string) (bool, error) {
-	currentFilePath := filepath.Join(currentStateDir, sc.checkpointName)
-	otherFilePath := filepath.Join(otherStateDir, sc.checkpointName)
-	return file.FilesEqual(currentFilePath, otherFilePath)
-}
-
 func (sc *stateCheckpoint) StoreState() error {
 	sc.Lock()
 	defer sc.Unlock()
@@ -244,7 +217,7 @@ func (sc *stateCheckpoint) storeState() error {
 	checkpoint.PodEntries = sc.cache.GetPodEntries()
 	checkpoint.AllowSharedCoresOverlapReclaimedCores = sc.cache.GetAllowSharedCoresOverlapReclaimedCores()
 
-	err := sc.checkpointManager.CreateCheckpoint(sc.checkpointName, checkpoint)
+	err := sc.qrmCheckpointManager.CreateCheckpoint(sc.checkpointName, checkpoint)
 	if err != nil {
 		klog.ErrorS(err, "Could not save checkpoint")
 		return err

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

@@ -43,6 +43,7 @@ import (
 	"github.com/kubewharf/katalyst-core/pkg/config/agent/qrm/statedirectory"
 	"github.com/kubewharf/katalyst-core/pkg/metrics"
 	"github.com/kubewharf/katalyst-core/pkg/util/machine"
+	"github.com/kubewharf/katalyst-core/pkg/util/qrmcheckpointmanager"
 )
 
 const (
@@ -3352,14 +3353,15 @@ func TestTryMigrateState(t *testing.T) {
 				skipStateCorruption:                false,
 				GenerateMachineStateFromPodEntries: generateTestMachineStateFromPodEntries,
 				emitter:                            metrics.DummyMetrics{},
+				hasPreStop:                         tt.preStop,
 			}
 
 			// current checkpoint is pointing to the in memory directory
-			sc.checkpointManager, err = checkpointmanager.NewCheckpointManager(inMemoryTmpDir)
+			sc.qrmCheckpointManager, err = qrmcheckpointmanager.NewQRMCheckpointManager(inMemoryTmpDir, stateDir, checkpointName, "cpu_plugin")
 			assert.NoError(t, err)
 
 			newCheckpoint := NewCPUPluginCheckpoint()
-			err = sc.tryMigrateState(cpuTopology, inMemoryTmpDir, stateDir, tt.preStop, newCheckpoint)
+			err = sc.tryMigrateState(cpuTopology, newCheckpoint)
 
 			if tt.corruptFile {
 				assert.Error(t, err)
@@ -3368,7 +3370,7 @@ func TestTryMigrateState(t *testing.T) {
 			assert.NoError(t, err)
 
 			// check if new checkpoint is created and verify equality
-			err = sc.checkpointManager.GetCheckpoint(checkpointName, newCheckpoint)
+			err = sc.qrmCheckpointManager.GetCurrentCheckpoint(checkpointName, newCheckpoint, false)
 			assert.NoError(t, err)
 
 			// verify old checkpoint file existence

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

@@ -32,8 +32,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/kubewharf/katalyst-core/pkg/config/agent/qrm/statedirectory"
-
 	"github.com/cilium/ebpf"
 	"github.com/cilium/ebpf/rlimit"
 	info "github.com/google/cadvisor/info/v1"
@@ -70,6 +68,7 @@ import (
 	"github.com/kubewharf/katalyst-core/pkg/config/agent/dynamic"
 	"github.com/kubewharf/katalyst-core/pkg/config/agent/global"
 	qrmconfig "github.com/kubewharf/katalyst-core/pkg/config/agent/qrm"
+	"github.com/kubewharf/katalyst-core/pkg/config/agent/qrm/statedirectory"
 	"github.com/kubewharf/katalyst-core/pkg/config/generic"
 	coreconsts "github.com/kubewharf/katalyst-core/pkg/consts"
 	"github.com/kubewharf/katalyst-core/pkg/metaserver"
@@ -112,7 +111,9 @@ var fakeConf = &config.Configuration{
 	},
 }
 
-func getTestDynamicPolicyWithInitialization(topology *machine.CPUTopology, machineInfo *info.MachineInfo, stateFileDirectory string) (*DynamicPolicy, error) {
+func getTestDynamicPolicyWithInitialization(
+	topology *machine.CPUTopology, machineInfo *info.MachineInfo, stateFileDirectory string,
+) (*DynamicPolicy, error) {
 	reservedMemory, err := getReservedMemory(fakeConf, &metaserver.MetaServer{}, machineInfo)
 	if err != nil {
 		return nil, err
@@ -4871,7 +4872,9 @@ type mockMemoryAdvisor struct {
 	advisorsvc.AdvisorServiceServer
 }
 
-func (m *mockMemoryAdvisor) GetAdvice(ctx context.Context, in *advisorsvc.GetAdviceRequest) (*advisorsvc.GetAdviceResponse, error) {
+func (m *mockMemoryAdvisor) GetAdvice(
+	ctx context.Context, in *advisorsvc.GetAdviceRequest,
+) (*advisorsvc.GetAdviceResponse, error) {
 	args := m.Called(ctx, in)
 	return args.Get(0).(*advisorsvc.GetAdviceResponse), args.Error(1)
 }