pkg/scheduler/backend/cache/cache.go - master - githubmirror/kubernetes

pkg/scheduler/backend/cache/cache.go
1 420 строк50 KB
Tetiana Khotiainsteva
scheduler: optimize evaluation of kubernetes.io/hostname-scoped inter-pod affinity
25 июл 2026, 15:25
25 июл 2026, 15:25fb74900
/*
Copyright 2015 The Kubernetes 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 cache

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	v1 "k8s.io/api/core/v1"
	schedulingv1alpha3 "k8s.io/api/scheduling/v1alpha3"
	schedulingv1beta1 "k8s.io/api/scheduling/v1beta1"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/klog/v2"
	fwk "k8s.io/kube-scheduler/framework"
	"k8s.io/kubernetes/pkg/scheduler/framework"
	apicalls "k8s.io/kubernetes/pkg/scheduler/framework/api_calls"
	"k8s.io/kubernetes/pkg/scheduler/metrics"
)

var (
	updateMetricsPeriod = 1 * time.Second
)

// New returns a Cache implementation.
// It automatically starts a go routine that exports cache metrics.
// "ctx" is the context that would close the background goroutine.
func New(ctx context.Context, apiDispatcher fwk.APIDispatcher, genericWorkloadEnabled, compositePodGroupEnabled bool) Cache {
	cache := newCache(ctx, updateMetricsPeriod, apiDispatcher, genericWorkloadEnabled, compositePodGroupEnabled)
	cache.run()
	return cache
}

// nodeInfoListItem holds a NodeInfo pointer and acts as an item in a doubly
// linked list. When a NodeInfo is updated, it goes to the head of the list.
// The items closer to the head are the most recently updated items.
type nodeInfoListItem struct {
	info *framework.NodeInfo
	next *nodeInfoListItem
	prev *nodeInfoListItem
}

type cacheImpl struct {
	stop   <-chan struct{}
	period time.Duration

	// This mutex guards all fields within this cache struct.
	mu sync.RWMutex
	// a set of assumed pod keys.
	// The key could further be used to get an entry in podStates.
	assumedPods sets.Set[string]
	// a map from pod key to podState.
	podStates map[string]*podState
	nodes     map[string]*nodeInfoListItem
	// headNode points to the most recently updated NodeInfo in "nodes". It is the
	// head of the linked list.
	headNode *nodeInfoListItem
	nodeTree *nodeTree
	// A map from image name to its ImageStateSummary.
	imageStates map[string]*fwk.ImageStateSummary
	// podGroupStates stores the runtime state for each known pod group (only if GenericWorkload feature gate is enabled).
	podGroupStates map[fwk.EntityKey]*podGroupState
	// compositePodGroupStates stores the runtime state for each known composite pod group (only if CompositePodGroup feature gate is enabled).
	compositePodGroupStates map[fwk.EntityKey]*compositePodGroupState
	// genericWorkloadEnabled stores the GenericWorkload feature gate value.
	genericWorkloadEnabled bool
	// compositePodGroupEnabled stores the CompositePodGroup feature gate value.
	compositePodGroupEnabled bool
	// apiDispatcher is used for the methods that are expected to send API calls.
	// It's non-nil only if the SchedulerAsyncAPICalls feature gate is enabled.
	apiDispatcher fwk.APIDispatcher

	// pvcRefCountsDelta contains the delta of changes to PVCRefCounts since the last snapshot.
	// Keys are in the format "namespace/name". This data struct serves as an optimization for avoiding
	// burdensome PVC ref count aggregations during scheduler cycles. PVCRefCountsDelta holds the incoming
	// deltas from events handlers while within the scheduler cycle, we only apply and reset the delta to
	// avoid global re-calculation.
	pvcRefCountsDelta map[string]int
}

type podState struct {
	pod *v1.Pod
}

func newCache(ctx context.Context, period time.Duration, apiDispatcher fwk.APIDispatcher, genericWorkloadEnabled, compositePodGroupEnabled bool) *cacheImpl {
	logger := klog.FromContext(ctx)
	return &cacheImpl{
		period: period,
		stop:   ctx.Done(),

		nodes:                    make(map[string]*nodeInfoListItem),
		nodeTree:                 newNodeTree(logger, nil),
		assumedPods:              sets.New[string](),
		podStates:                make(map[string]*podState),
		imageStates:              make(map[string]*fwk.ImageStateSummary),
		podGroupStates:           make(map[fwk.EntityKey]*podGroupState),
		compositePodGroupStates:  make(map[fwk.EntityKey]*compositePodGroupState),
		genericWorkloadEnabled:   genericWorkloadEnabled,
		compositePodGroupEnabled: compositePodGroupEnabled,
		apiDispatcher:            apiDispatcher,
		pvcRefCountsDelta:        make(map[string]int),
	}
}

// newNodeInfoListItem initializes a new nodeInfoListItem.
func newNodeInfoListItem(ni *framework.NodeInfo) *nodeInfoListItem {
	return &nodeInfoListItem{
		info: ni,
	}
}

// moveNodeInfoToHead moves a NodeInfo to the head of "cache.nodes" doubly
// linked list. The head is the most recently updated NodeInfo.
// We assume cache lock is already acquired.
func (cache *cacheImpl) moveNodeInfoToHead(logger klog.Logger, name string) {
	ni, ok := cache.nodes[name]
	if !ok {
		utilruntime.HandleErrorWithLogger(logger, nil, "No node info with given name found in the cache", "node", klog.KRef("", name))
		return
	}
	// if the node info list item is already at the head, we are done.
	if ni == cache.headNode {
		return
	}

	if ni.prev != nil {
		ni.prev.next = ni.next
	}
	if ni.next != nil {
		ni.next.prev = ni.prev
	}
	if cache.headNode != nil {
		cache.headNode.prev = ni
	}
	ni.next = cache.headNode
	ni.prev = nil
	cache.headNode = ni
}

// removeNodeInfoFromList removes a NodeInfo from the "cache.nodes" doubly
// linked list.
// We assume cache lock is already acquired.
func (cache *cacheImpl) removeNodeInfoFromList(logger klog.Logger, name string) {
	ni, ok := cache.nodes[name]
	if !ok {
		utilruntime.HandleErrorWithLogger(logger, nil, "No node info with given name found in the cache", "node", klog.KRef("", name))
		return
	}

	if ni.prev != nil {
		ni.prev.next = ni.next
	}
	if ni.next != nil {
		ni.next.prev = ni.prev
	}
	// if the removed item was at the head, we must update the head.
	if ni == cache.headNode {
		cache.headNode = ni.next
	}
	delete(cache.nodes, name)
}

// Dump produces a dump of the current scheduler cache. This is used for
// debugging purposes only and shouldn't be confused with UpdateSnapshot
// function.
// This method is expensive, and should be only used in non-critical path.
func (cache *cacheImpl) Dump() *Dump {
	cache.mu.RLock()
	defer cache.mu.RUnlock()

	nodes := make(map[string]*framework.NodeInfo, len(cache.nodes))
	for k, v := range cache.nodes {
		nodes[k] = v.info.SnapshotConcrete()
	}

	return &Dump{
		Nodes:       nodes,
		AssumedPods: cache.assumedPods.Union(nil),
	}
}

// UpdateSnapshot takes a snapshot of cached NodeInfo map. This is called at
// beginning of every scheduling cycle.
// The snapshot only includes Nodes that are not deleted at the time this function is called.
// nodeInfo.Node() is guaranteed to be not nil for all the nodes in the snapshot.
// This function tracks generation number of NodeInfo and updates only the
// entries of an existing snapshot that have changed after the snapshot was taken.
func (cache *cacheImpl) UpdateSnapshot(logger klog.Logger, nodeSnapshot *Snapshot) error {
	cache.mu.Lock()
	defer cache.mu.Unlock()

	if nodeSnapshot.placementNodes != nil {
		logger.Error(nil, "UpdateSnapshot called with assumed placement. This is unexpected. Placement will be cleared.")
		nodeSnapshot.ForgetPlacement()
	}

	// Get the last generation of the snapshot.
	snapshotGeneration := nodeSnapshot.generation

	// NodeInfoList and HavePodsWithAffinityNodeInfoList must be re-created if a node was added
	// or removed from the cache.
	updateAllLists := false
	// HavePodsWithAffinityNodeInfoList must be re-created if a node changed its
	// status from having pods with affinity to NOT having pods with affinity or the other
	// way around.
	updateNodesHavePodsWithAffinity := false
	// HavePodsWithRequiredAntiAffinityNodeInfoList must be re-created if a node changed its
	// status from having pods with required anti-affinity to NOT having pods with required
	// anti-affinity or the other way around.
	updateNodesHavePodsWithRequiredAntiAffinity := false

	// Forget all assumed pods from a previous snapshot version.
	// This is a safety check in case any pod wasn't forgotten in the previous scheduling cycle.
	nodeSnapshot.forgetAllAssumedPods(logger)

	// HavePodsWithRequiredNonHostScopedAntiAffinityNodeInfoList must be re-created if a node changed its status
	// from having pods with required anti-affinity with topologyKey other than kubernetes.io/hostname to NOT
	// having such pods, or the other way around.
	updateNodesHavePodsWithRequiredNonHostScopedAntiAffinity := false

	// Start from the head of the NodeInfo doubly linked list and update snapshot
	// of NodeInfos updated after the last snapshot.
	for node := cache.headNode; node != nil; node = node.next {
		if node.info.Generation <= snapshotGeneration {
			// all the nodes are updated before the existing snapshot. We are done.
			break
		}
		if np := node.info.Node(); np != nil {
			existing, ok := nodeSnapshot.nodeInfoMap[np.Name]
			if !ok {
				updateAllLists = true
				existing = &framework.NodeInfo{}
				nodeSnapshot.nodeInfoMap[np.Name] = existing
			}
			clone := node.info.SnapshotConcrete()
			// We track nodes that have pods with affinity, here we check if this node changed its
			// status from having pods with affinity to NOT having pods with affinity or the other
			// way around.
			if (len(existing.PodsWithAffinity) > 0) != (len(clone.PodsWithAffinity) > 0) {
				updateNodesHavePodsWithAffinity = true
			}
			if (len(existing.PodsWithRequiredAntiAffinity) > 0) != (len(clone.PodsWithRequiredAntiAffinity) > 0) {
				updateNodesHavePodsWithRequiredAntiAffinity = true
			}
			if (len(existing.PodsWithRequiredNonHostScopedAntiAffinity) > 0) != (len(clone.PodsWithRequiredNonHostScopedAntiAffinity) > 0) {
				updateNodesHavePodsWithRequiredNonHostScopedAntiAffinity = true
			}
			// We need to preserve the original pointer of the NodeInfo struct since it
			// is used in the NodeInfoList, which we may not update.
			*existing = *clone
		}
	}
	// Update the snapshot generation with the latest NodeInfo generation.
	if cache.headNode != nil {
		nodeSnapshot.generation = cache.headNode.info.Generation
	}

	// Comparing to pods in nodeTree.
	// Deleted nodes get removed from the tree, but they might remain in the nodes map
	// if they still have non-deleted Pods.
	if len(nodeSnapshot.nodeInfoMap) > cache.nodeTree.numNodes {
		cache.removeDeletedNodesFromSnapshot(nodeSnapshot)
		updateAllLists = true
	}

	// Apply the deltas for PVC reference count to the snapshot.
	// This no-op if the snapshot is built afresh i.e. updateAllLists=true
	if !updateAllLists {
		if err := cache.applyPVCRefCountDelta(nodeSnapshot); err != nil {
			logger.Error(err, "rebuilding node snapshot due to unexpected error from refreshing PVC ref counts")
			updateAllLists = true
		}
	}

	if updateAllLists || updateNodesHavePodsWithAffinity || updateNodesHavePodsWithRequiredAntiAffinity || updateNodesHavePodsWithRequiredNonHostScopedAntiAffinity {
		cache.updateNodeInfoSnapshotList(logger, nodeSnapshot, updateAllLists)
	}

	if len(nodeSnapshot.nodeInfoList) != cache.nodeTree.numNodes {
		errMsg := fmt.Sprintf("snapshot state is not consistent, length of NodeInfoList=%v not equal to length of nodes in tree=%v "+
			", length of NodeInfoMap=%v, length of nodes in cache=%v"+
			", trying to recover",
			len(nodeSnapshot.nodeInfoList), cache.nodeTree.numNodes,
			len(nodeSnapshot.nodeInfoMap), len(cache.nodes))
		logger.Error(nil, errMsg)
		// We will try to recover by re-creating the lists for the next scheduling cycle, but still return an
		// error to surface the problem, the error will likely cause a failure to the current scheduling cycle.
		cache.updateNodeInfoSnapshotList(logger, nodeSnapshot, true)
		return errors.New(errMsg)
	}

	if cache.genericWorkloadEnabled {
		// Take a snapshot of pod group states for this scheduling cycle.
		cache.updatePodGroupStateSnapshot(nodeSnapshot)
		if cache.compositePodGroupEnabled {
			// Take a snapshot of composite pod group states for this scheduling cycle.
			cache.updateCompositePodGroupStateSnapshot(nodeSnapshot)
		}
	}

	return nil
}

// updatePodGroupStateSnapshot updates the pod group state portion of the given snapshot.
// It assumes that the cache lock is already held.
// It removes entries that no longer exist in the live cache
// and clones entries whose generation has advanced since the last snapshot.
func (cache *cacheImpl) updatePodGroupStateSnapshot(snapshot *Snapshot) {
	// Remove pod group states from snapshot that no longer exist in cache.
	for key := range snapshot.podGroupStates {
		if _, exists := cache.podGroupStates[key]; !exists {
			delete(snapshot.podGroupStates, key)
		}
	}
	// Clone only pod group states that changed since the last snapshot.
	for key, podGroupState := range cache.podGroupStates {
		if existing, ok := snapshot.podGroupStates[key]; ok && existing.generation == podGroupState.generation {
			continue
		}
		snapshot.podGroupStates[key] = podGroupState.snapshot()
	}
}

// updateCompositePodGroupStateSnapshot updates the composite pod group state portion of the given snapshot.
// It assumes that the cache lock is already held.
// It removes entries that no longer exist in the live cache
// and clones entries whose generation has advanced since the last snapshot.
func (cache *cacheImpl) updateCompositePodGroupStateSnapshot(snapshot *Snapshot) {
	// Remove pod group states from snapshot that no longer exist in cache.
	for key := range snapshot.compositePodGroupStates {
		if _, exists := cache.compositePodGroupStates[key]; !exists {
			delete(snapshot.compositePodGroupStates, key)
		}
	}
	// Clone only pod group states that changed since the last snapshot.
	for key, cpgState := range cache.compositePodGroupStates {
		if existing, ok := snapshot.compositePodGroupStates[key]; ok && existing.generation == cpgState.generation {
			continue
		}
		snapshot.compositePodGroupStates[key] = cpgState.snapshot()
	}
}

func (cache *cacheImpl) updateNodeInfoSnapshotList(logger klog.Logger, snapshot *Snapshot, updateAll bool) {
	snapshot.havePodsWithAffinityNodeInfoList = make([]fwk.NodeInfo, 0, cache.nodeTree.numNodes)
	snapshot.havePodsWithRequiredAntiAffinityNodeInfoList = make([]fwk.NodeInfo, 0, cache.nodeTree.numNodes)
	snapshot.havePodsWithRequiredNonHostScopedAntiAffinityNodeInfoList = make([]fwk.NodeInfo, 0, cache.nodeTree.numNodes)
	if updateAll {
		snapshot.usedPVCRefCounts = make(map[string]int)
		// Take a snapshot of the nodes order in the tree
		snapshot.nodeInfoList = make([]fwk.NodeInfo, 0, cache.nodeTree.numNodes)
		nodesList, err := cache.nodeTree.list()
		if err != nil {
			utilruntime.HandleErrorWithLogger(logger, err, "Error occurred while retrieving the list of names of the nodes from node tree")
		}
		for _, nodeName := range nodesList {
			if nodeInfo := snapshot.nodeInfoMap[nodeName]; nodeInfo != nil {
				snapshot.nodeInfoList = append(snapshot.nodeInfoList, nodeInfo)
				if len(nodeInfo.PodsWithAffinity) > 0 {
					snapshot.havePodsWithAffinityNodeInfoList = append(snapshot.havePodsWithAffinityNodeInfoList, nodeInfo)
				}
				if len(nodeInfo.PodsWithRequiredAntiAffinity) > 0 {
					snapshot.havePodsWithRequiredAntiAffinityNodeInfoList = append(snapshot.havePodsWithRequiredAntiAffinityNodeInfoList, nodeInfo)
				}
				if len(nodeInfo.PodsWithRequiredNonHostScopedAntiAffinity) > 0 {
					snapshot.havePodsWithRequiredNonHostScopedAntiAffinityNodeInfoList = append(snapshot.havePodsWithRequiredNonHostScopedAntiAffinityNodeInfoList, nodeInfo)
				}
				for key, value := range nodeInfo.PVCRefCounts {
					snapshot.usedPVCRefCounts[key] += value
				}
			} else {
				utilruntime.HandleErrorWithLogger(logger, nil, "Node exists in nodeTree but not in NodeInfoMap, this should not happen", "node", klog.KRef("", nodeName))
			}
		}
		// reset the deltas if update all
		cache.pvcRefCountsDelta = map[string]int{}
	} else {
		for _, nodeInfo := range snapshot.nodeInfoList {
			if len(nodeInfo.GetPodsWithAffinity()) > 0 {
				snapshot.havePodsWithAffinityNodeInfoList = append(snapshot.havePodsWithAffinityNodeInfoList, nodeInfo)
			}
			if len(nodeInfo.GetPodsWithRequiredAntiAffinity()) > 0 {
				snapshot.havePodsWithRequiredAntiAffinityNodeInfoList = append(snapshot.havePodsWithRequiredAntiAffinityNodeInfoList, nodeInfo)
			}
			if len(nodeInfo.GetPodsWithRequiredNonHostScopedAntiAffinity()) > 0 {
				snapshot.havePodsWithRequiredNonHostScopedAntiAffinityNodeInfoList = append(snapshot.havePodsWithRequiredNonHostScopedAntiAffinityNodeInfoList, nodeInfo)
			}
		}
	}
}

// If certain nodes were deleted after the last snapshot was taken, we should remove them from the snapshot.
func (cache *cacheImpl) removeDeletedNodesFromSnapshot(snapshot *Snapshot) {
	toDelete := len(snapshot.nodeInfoMap) - cache.nodeTree.numNodes
	for name := range snapshot.nodeInfoMap {
		if toDelete <= 0 {
			break
		}
		if n, ok := cache.nodes[name]; !ok || n.info.Node() == nil {
			delete(snapshot.nodeInfoMap, name)
			toDelete--
		}
	}
}

// NodeCount returns the number of nodes in the cache.
// DO NOT use outside of tests.
func (cache *cacheImpl) NodeCount() int {
	cache.mu.RLock()
	defer cache.mu.RUnlock()
	return len(cache.nodes)
}

// PodCount returns the number of pods in the cache (including those from deleted nodes).
// DO NOT use outside of tests.
func (cache *cacheImpl) PodCount() (int, error) {
	cache.mu.RLock()
	defer cache.mu.RUnlock()
	// podFilter is expected to return true for most or all of the pods. We
	// can avoid expensive array growth without wasting too much memory by
	// pre-allocating capacity.
	count := 0
	for _, n := range cache.nodes {
		count += len(n.info.Pods)
	}
	return count, nil
}

func (cache *cacheImpl) AssumePod(logger klog.Logger, pod *v1.Pod) error {
	key, err := framework.GetPodKey(pod)
	if err != nil {
		return err
	}

	cache.mu.Lock()
	defer cache.mu.Unlock()
	if _, ok := cache.podStates[key]; ok {
		return fmt.Errorf("pod %v(%v) is in the cache, so can't be assumed", key, klog.KObj(pod))
	}

	return cache.addPod(logger, pod, true)
}

// validateAssumedPod checks that the given pod is currently assumed.
// Assumes that lock is already acquired.
func (cache *cacheImpl) validateAssumedPod(pod *v1.Pod, key string, currState *podState) error {
	if currState.pod.Spec.NodeName != pod.Spec.NodeName {
		return fmt.Errorf("pod %v(%v) was assumed on %v but assigned to %v", key, klog.KObj(pod), pod.Spec.NodeName, currState.pod.Spec.NodeName)
	}
	if !cache.assumedPods.Has(key) {
		return fmt.Errorf("pod %v(%v) is not assumed, so it cannot be removed or forgotten", key, klog.KObj(pod))
	}
	return nil
}

// ForgetPod forgets an assumed pod from the cache. It should be called when the pod
// still exists, as an undo operation for AssumePod.
// If the pod is a pod group member, it is moved from assumed to unscheduled pods of that pod group state in cache.
func (cache *cacheImpl) ForgetPod(logger klog.Logger, pod *v1.Pod) error {
	key, err := framework.GetPodKey(pod)
	if err != nil {
		return err
	}

	cache.mu.Lock()
	defer cache.mu.Unlock()

	currState, ok := cache.podStates[key]
	if !ok {
		// Pod does not exist in the cache anymore.
		return nil
	}
	if err := cache.validateAssumedPod(pod, key, currState); err != nil {
		return err
	}
	return cache.removePod(logger, pod, true)
}

// RemoveAssumedPod removes an assumed pod from the cache. It should be called when the assumed
// pod was removed from the cluster to correctly clean up internal state.
// It differs from ForgetPod in how it handles pod group members, as it removes the pod from
// the pod group state in the cache.
func (cache *cacheImpl) RemoveAssumedPod(logger klog.Logger, pod *v1.Pod) error {
	key, err := framework.GetPodKey(pod)
	if err != nil {
		return err
	}

	cache.mu.Lock()
	defer cache.mu.Unlock()

	currState, ok := cache.podStates[key]
	if !ok {
		// Pod does not exist in the cache anymore.
		return nil
	}
	if err := cache.validateAssumedPod(pod, key, currState); err != nil {
		return err
	}
	return cache.removePod(logger, pod, false)
}

// Assumes that lock is already acquired.
func (cache *cacheImpl) addPod(logger klog.Logger, pod *v1.Pod, assumePod bool) error {
	key, err := framework.GetPodKey(pod)
	if err != nil {
		return err
	}
	n, ok := cache.nodes[pod.Spec.NodeName]
	if !ok {
		n = newNodeInfoListItem(framework.NewNodeInfo())
		cache.nodes[pod.Spec.NodeName] = n
	}

	// new_delta = old_delta + (PVCRefCounts_after − PVCRefCounts_before)
	cache.refreshPVCRefCountsDelta(n.info, -1)
	defer cache.refreshPVCRefCountsDelta(n.info, 1)

	n.info.AddPod(pod)
	cache.moveNodeInfoToHead(logger, pod.Spec.NodeName)
	ps := &podState{
		pod: pod,
	}
	cache.podStates[key] = ps
	if assumePod {
		cache.assumedPods.Insert(key)
	}

	if !cache.isPodGroupMember(pod) {
		return nil
	}
	if assumePod {
		cache.assumePodGroupMember(pod)
	} else {
		cache.addPodGroupMember(pod)
	}
	return nil
}

// Assumes that lock is already acquired.
func (cache *cacheImpl) updatePod(logger klog.Logger, oldPod, newPod *v1.Pod) error {
	if err := cache.removePod(logger, oldPod, false); err != nil {
		return err
	}
	return cache.addPod(logger, newPod, false)
}

// Assumes that lock is already acquired.
// Removes a pod from the cached node info. If the node information was already
// removed and there are no more pods left in the node, cleans up the node from
// the cache.
func (cache *cacheImpl) removePod(logger klog.Logger, pod *v1.Pod, forgetPod bool) error {
	key, err := framework.GetPodKey(pod)
	if err != nil {
		return err
	}

	n, ok := cache.nodes[pod.Spec.NodeName]
	if !ok {
		utilruntime.HandleErrorWithLogger(logger, nil, "Node not found when trying to remove pod", "node", klog.KRef("", pod.Spec.NodeName), "podKey", key, "pod", klog.KObj(pod))
	} else {
		// new_delta = old_delta + (PVCRefCounts_after − PVCRefCounts_before)
		cache.refreshPVCRefCountsDelta(n.info, -1)
		defer cache.refreshPVCRefCountsDelta(n.info, 1)

		if err := n.info.RemovePod(logger, pod); err != nil {
			return err
		}
		if len(n.info.Pods) == 0 && n.info.Node() == nil {
			cache.removeNodeInfoFromList(logger, pod.Spec.NodeName)
		} else {
			cache.moveNodeInfoToHead(logger, pod.Spec.NodeName)
		}
	}

	delete(cache.podStates, key)
	delete(cache.assumedPods, key)

	if !cache.isPodGroupMember(pod) {
		return nil
	}
	if forgetPod {
		cache.forgetPodGroupMember(logger, pod)
	} else {
		cache.removePodGroupMember(pod)
	}

	return nil
}

func (cache *cacheImpl) AddPod(logger klog.Logger, pod *v1.Pod) error {
	key, err := framework.GetPodKey(pod)
	if err != nil {
		return err
	}

	cache.mu.Lock()
	defer cache.mu.Unlock()

	currState, ok := cache.podStates[key]
	switch {
	case ok && cache.assumedPods.Has(key):
		// When assuming, we've already added the Pod to cache,
		// Just update here to make sure the Pod's status is up-to-date.
		if err = cache.updatePod(logger, currState.pod, pod); err != nil {
			utilruntime.HandleErrorWithLogger(logger, err, "Error occurred while updating pod")
		}
		if currState.pod.Spec.NodeName != pod.Spec.NodeName {
			// The pod was added to a different node than it was assumed to.
			logger.Info("Pod was added to a different node than it was assumed", "podKey", key, "pod", klog.KObj(pod), "assumedNode", klog.KRef("", pod.Spec.NodeName), "currentNode", klog.KRef("", currState.pod.Spec.NodeName))
			return nil
		}
	case !ok:
		if err = cache.addPod(logger, pod, false); err != nil {
			utilruntime.HandleErrorWithLogger(logger, err, "Error occurred while adding pod")
		}
	default:
		return fmt.Errorf("pod %v(%v) was already in added state", key, klog.KObj(pod))
	}
	return nil
}

func (cache *cacheImpl) UpdatePod(logger klog.Logger, oldPod, newPod *v1.Pod) error {
	key, err := framework.GetPodKey(oldPod)
	if err != nil {
		return err
	}

	cache.mu.Lock()
	defer cache.mu.Unlock()

	currState, ok := cache.podStates[key]
	if !ok {
		return fmt.Errorf("pod %v(%v) is not added to scheduler cache, so cannot be updated", key, klog.KObj(oldPod))
	}

	// An assumed pod won't have Update/Remove event. It needs to have Add event
	// before Update event, in which case the state would change from Assumed to Added.
	if cache.assumedPods.Has(key) {
		return fmt.Errorf("assumed pod %v(%v) should not be updated", key, klog.KObj(oldPod))
	}

	if currState.pod.Spec.NodeName != newPod.Spec.NodeName {
		utilruntime.HandleErrorWithLogger(logger, nil, "Pod updated on a different node than previously added to. Scheduler cache is corrupted and can badly affect scheduling decisions", "podKey", key, "pod", klog.KObj(oldPod))
		klog.FlushAndExit(klog.ExitFlushTimeout, 1)
	}
	return cache.updatePod(logger, oldPod, newPod)
}

func (cache *cacheImpl) RemovePod(logger klog.Logger, pod *v1.Pod) error {
	key, err := framework.GetPodKey(pod)
	if err != nil {
		return err
	}

	cache.mu.Lock()
	defer cache.mu.Unlock()

	currState, ok := cache.podStates[key]
	if !ok {
		return fmt.Errorf("pod %v(%v) is not found in scheduler cache, so cannot be removed from it", key, klog.KObj(pod))
	}
	if currState.pod.Spec.NodeName != pod.Spec.NodeName {
		utilruntime.HandleErrorWithLogger(logger, nil, "Pod was added to a different node than it was assumed", "podKey", key, "pod", klog.KObj(pod), "assumedNode", klog.KRef("", pod.Spec.NodeName), "currentNode", klog.KRef("", currState.pod.Spec.NodeName))
		if pod.Spec.NodeName != "" {
			// An empty NodeName is possible when the scheduler misses a Delete
			// event and it gets the last known state from the informer cache.
			utilruntime.HandleErrorWithLogger(logger, nil, "Scheduler cache is corrupted and can badly affect scheduling decisions")
			klog.FlushAndExit(klog.ExitFlushTimeout, 1)
		}
	}
	return cache.removePod(logger, currState.pod, false)
}

func (cache *cacheImpl) IsAssumedPod(pod *v1.Pod) (bool, error) {
	key, err := framework.GetPodKey(pod)
	if err != nil {
		return false, err
	}

	cache.mu.RLock()
	defer cache.mu.RUnlock()

	return cache.assumedPods.Has(key), nil
}

// GetPod might return a pod for which its node has already been deleted from
// the main cache. This is useful to properly process pod update events.
func (cache *cacheImpl) GetPod(pod *v1.Pod) (*v1.Pod, error) {
	key, err := framework.GetPodKey(pod)
	if err != nil {
		return nil, err
	}

	cache.mu.RLock()
	defer cache.mu.RUnlock()

	podState, ok := cache.podStates[key]
	if !ok {
		return nil, fmt.Errorf("pod %v(%v) does not exist in scheduler cache", key, klog.KObj(pod))
	}

	return podState.pod, nil
}

func (cache *cacheImpl) AddNode(logger klog.Logger, node *v1.Node) {
	cache.mu.Lock()
	defer cache.mu.Unlock()

	n, ok := cache.nodes[node.Name]
	if !ok {
		n = newNodeInfoListItem(framework.NewNodeInfo())
		cache.nodes[node.Name] = n
	} else {
		cache.removeNodeImageStates(n.info.Node())
	}
	cache.moveNodeInfoToHead(logger, node.Name)

	cache.nodeTree.addNode(logger, node)
	cache.addNodeImageStates(node, n.info)
	n.info.SetNode(node)
}

func (cache *cacheImpl) UpdateNode(logger klog.Logger, oldNode, newNode *v1.Node) {
	cache.mu.Lock()
	defer cache.mu.Unlock()
	n, ok := cache.nodes[newNode.Name]
	if !ok {
		n = newNodeInfoListItem(framework.NewNodeInfo())
		cache.nodes[newNode.Name] = n
		cache.nodeTree.addNode(logger, newNode)
	} else {
		cache.removeNodeImageStates(n.info.Node())
	}
	cache.moveNodeInfoToHead(logger, newNode.Name)

	cache.nodeTree.updateNode(logger, oldNode, newNode)
	cache.addNodeImageStates(newNode, n.info)
	n.info.SetNode(newNode)
}

// RemoveNode removes a node from the cache's tree.
// The node might still have pods because their deletion events didn't arrive
// yet. Those pods are considered removed from the cache, being the node tree
// the source of truth.
// However, we keep a ghost node with the list of pods until all pod deletion
// events have arrived. A ghost node is skipped from snapshots.
func (cache *cacheImpl) RemoveNode(logger klog.Logger, node *v1.Node) error {
	cache.mu.Lock()
	defer cache.mu.Unlock()

	n, ok := cache.nodes[node.Name]
	if !ok {
		return fmt.Errorf("node %v is not found", node.Name)
	}

	// only subtract the PVCRefCount into the delta map
	cache.refreshPVCRefCountsDelta(n.info, -1)

	n.info.RemoveNode()
	// We remove NodeInfo for this node only if there aren't any pods on this node.
	// We can't do it unconditionally, because notifications about pods are delivered
	// in a different watch, and thus can potentially be observed later, even though
	// they happened before node removal.
	if len(n.info.Pods) == 0 {
		cache.removeNodeInfoFromList(logger, node.Name)
	} else {
		cache.moveNodeInfoToHead(logger, node.Name)
	}
	if err := cache.nodeTree.removeNode(logger, node); err != nil {
		return err
	}
	cache.removeNodeImageStates(node)
	return nil
}

// GetNode returns the copy of node stored in the cache.
// DO NOT use outside of tests.
func (cache *cacheImpl) GetNode(name string) (*framework.NodeInfo, error) {
	cache.mu.RLock()
	defer cache.mu.RUnlock()

	n, ok := cache.nodes[name]
	if !ok {
		return nil, fmt.Errorf("node %v does not exist in scheduler cache", name)
	}
	return n.info.SnapshotConcrete(), nil
}

// addNodeImageStates adds states of the images on given node to the given nodeInfo and update the imageStates in
// scheduler cache. This function assumes the lock to scheduler cache has been acquired.
func (cache *cacheImpl) addNodeImageStates(node *v1.Node, nodeInfo *framework.NodeInfo) {
	newSum := make(map[string]*fwk.ImageStateSummary)

	for _, image := range node.Status.Images {
		for _, name := range image.Names {
			// update the entry in imageStates
			state, ok := cache.imageStates[name]
			if !ok {
				state = &fwk.ImageStateSummary{
					Size:  image.SizeBytes,
					Nodes: sets.New(node.Name),
				}
				cache.imageStates[name] = state
			} else {
				state.Nodes.Insert(node.Name)
			}
			// create the ImageStateSummary for this image
			if _, ok := newSum[name]; !ok {
				newSum[name] = state
			}
		}
	}
	nodeInfo.ImageStates = newSum
}

// removeNodeImageStates removes the given node record from image entries having the node
// in imageStates cache. After the removal, if any image becomes free, i.e., the image
// is no longer available on any node, the image entry will be removed from imageStates.
func (cache *cacheImpl) removeNodeImageStates(node *v1.Node) {
	if node == nil {
		return
	}

	for _, image := range node.Status.Images {
		for _, name := range image.Names {
			state, ok := cache.imageStates[name]
			if ok {
				state.Nodes.Delete(node.Name)
				if state.Nodes.Len() == 0 {
					// Remove the unused image to make sure the length of
					// imageStates represents the total number of different
					// images on all nodes
					delete(cache.imageStates, name)
				}
			}
		}
	}
}

func (cache *cacheImpl) run() {
	go wait.Until(cache.updateMetrics, cache.period, cache.stop)
}

// updateMetrics updates cache size metric values for pods, assumed pods, and nodes
func (cache *cacheImpl) updateMetrics() {
	cache.mu.RLock()
	defer cache.mu.RUnlock()

	metrics.CacheSize.WithLabelValues("assumed_pods").Set(float64(len(cache.assumedPods)))
	metrics.CacheSize.WithLabelValues("pods").Set(float64(len(cache.podStates)))
	metrics.CacheSize.WithLabelValues("nodes").Set(float64(len(cache.nodes)))
}

// isPodGroupMember returns true if the pod belongs to a pod group,
// provided that GenericWorkload feature gate is enabled.
func (cache *cacheImpl) isPodGroupMember(pod *v1.Pod) bool {
	return cache.genericWorkloadEnabled && pod.Spec.SchedulingGroup != nil
}

// AddPodGroupMember adds not assigned and not assumed pod to its pod group state in the cache.
func (cache *cacheImpl) AddPodGroupMember(pod *v1.Pod) {
	if !cache.isPodGroupMember(pod) {
		return
	}
	cache.mu.Lock()
	defer cache.mu.Unlock()

	cache.addPodGroupMember(pod)
}

// UpdatePodGroupMember updates a pod's entry inside its pod group state in the cache.
func (cache *cacheImpl) UpdatePodGroupMember(logger klog.Logger, oldPod, newPod *v1.Pod) {
	if !cache.isPodGroupMember(newPod) {
		return
	}
	cache.mu.Lock()
	defer cache.mu.Unlock()

	cache.updatePodGroupMember(logger, oldPod, newPod)
}

// RemovePodGroupMember removes the pod from its pod group state in the cache.
func (cache *cacheImpl) RemovePodGroupMember(pod *v1.Pod) {
	if !cache.isPodGroupMember(pod) {
		return
	}
	cache.mu.Lock()
	defer cache.mu.Unlock()

	cache.removePodGroupMember(pod)
}

// addPodGroupMember adds the pod to its pod group state, creating the group entry if it doesn't exist yet.
// Assumes that the cache lock is already held.
func (cache *cacheImpl) addPodGroupMember(pod *v1.Pod) {
	key := fwk.PodGroupKey(pod.Namespace, *pod.Spec.SchedulingGroup.PodGroupName)
	podGroupState, exists := cache.podGroupStates[key]
	if !exists {
		podGroupState = newPodGroupState()
		cache.podGroupStates[key] = podGroupState
	}

	podGroupState.addPod(pod)
}

// updatePodGroupMember updates the pod entry inside its pod group state.
// Assumes that the cache lock is already held.
func (cache *cacheImpl) updatePodGroupMember(logger klog.Logger, oldPod, newPod *v1.Pod) {
	key := fwk.PodGroupKey(newPod.Namespace, *newPod.Spec.SchedulingGroup.PodGroupName)
	podGroupState, exists := cache.podGroupStates[key]
	if !exists {
		// This should not happen: the pod group state should have been already created by a prior pod add action.
		utilruntime.HandleErrorWithLogger(logger, nil, "Pod group state not found for update, this indicates a missed add event", "pod", klog.KObj(newPod), "fwk.EntityKey", key)
		return
	}

	podGroupState.updatePod(oldPod, newPod)
}

// removePodGroupMember removes the pod from its pod group state, deleting the group entry when empty.
// Assumes that the cache lock is already held.
func (cache *cacheImpl) removePodGroupMember(pod *v1.Pod) {
	key := fwk.PodGroupKey(pod.Namespace, *pod.Spec.SchedulingGroup.PodGroupName)
	podGroupState, exists := cache.podGroupStates[key]
	if !exists {
		return
	}
	podGroupState.deletePod(pod.UID)
	if podGroupState.empty() {
		// podGroupState can exist without the member pods, but when the PodGroup object exists.
		// Only when there are no member pods and the PodGroup object is removed, the podGroupState can be removed.
		delete(cache.podGroupStates, key)
	}
}

// assumePodGroupMember marks the pod as assumed in its pod group state.
// Assumes that the cache lock is already held.
func (cache *cacheImpl) assumePodGroupMember(pod *v1.Pod) {
	key := fwk.PodGroupKey(pod.Namespace, *pod.Spec.SchedulingGroup.PodGroupName)
	podGroupState, exists := cache.podGroupStates[key]
	if !exists {
		podGroupState = newPodGroupState()
		podGroupState.allPods[pod.UID] = pod
		cache.podGroupStates[key] = podGroupState
	}
	podGroupState.assumePod(pod)
}

// forgetPodGroupMember moves the pod back from assumed to unscheduled in its pod group state.
// Assumes that the cache lock is already held.
func (cache *cacheImpl) forgetPodGroupMember(logger klog.Logger, pod *v1.Pod) {
	key := fwk.PodGroupKey(pod.Namespace, *pod.Spec.SchedulingGroup.PodGroupName)
	pgs, exists := cache.podGroupStates[key]
	if !exists {
		// This should not happen: the pod group state should have been already created by a prior pod add or assume action.
		utilruntime.HandleErrorWithLogger(logger, nil, "Pod group state not found for forget, this indicates a missed add or assume event", "pod", klog.KObj(pod), "fwk.EntityKey", key)
		return
	}
	pgs.forgetPod(pod.UID)
}

// PodGroupStates returns the PodGroupStateLister for this cache.
func (cache *cacheImpl) PodGroupStates() fwk.PodGroupStateLister {
	return &podGroupStateListerImpl{cache: cache}
}

type podGroupStateListerImpl struct {
	cache *cacheImpl
}

func (l *podGroupStateListerImpl) Get(namespace, podGroupName string) (fwk.PodGroupState, error) {
	if !l.cache.genericWorkloadEnabled {
		return nil, fmt.Errorf("generic workload feature gate is disabled")
	}
	l.cache.mu.RLock()
	defer l.cache.mu.RUnlock()

	key := fwk.PodGroupKey(namespace, podGroupName)
	podGroupState, exists := l.cache.podGroupStates[key]
	if !exists {
		return nil, fmt.Errorf("pod group state not found for pod group %s/%s", namespace, podGroupName)
	}
	return podGroupState, nil
}

// PodGroups returns the PodGroupLister for this cache.
func (cache *cacheImpl) PodGroups() fwk.PodGroupLister {
	return &podGroupListerImpl{cache: cache}
}

type podGroupListerImpl struct {
	cache *cacheImpl
}

// Get returns the cached pod group object.
func (l *podGroupListerImpl) Get(namespace, name string) (*schedulingv1beta1.PodGroup, error) {
	if !l.cache.genericWorkloadEnabled {
		return nil, fmt.Errorf("generic workload feature gate is disabled")
	}
	l.cache.mu.RLock()
	defer l.cache.mu.RUnlock()

	key := fwk.PodGroupKey(namespace, name)
	pgs, exists := l.cache.podGroupStates[key]
	if !exists {
		return nil, fmt.Errorf("pod group state not found for pod group %s", key)
	}
	pg := pgs.podGroup
	if pg == nil {
		return nil, fmt.Errorf("pod group object not found for pod group %s", key)
	}
	return pg, nil
}

// CompositePodGroupStates returns the CompositePodGroupStateLister for this cache.
func (cache *cacheImpl) CompositePodGroupStates() fwk.CompositePodGroupStateLister {
	return &compositePodGroupStateListerImpl{cache: cache}
}

type compositePodGroupStateListerImpl struct {
	cache *cacheImpl
}

func (l *compositePodGroupStateListerImpl) Get(namespace string, podGroupName string) (fwk.CompositePodGroupState, error) {
	if !l.cache.compositePodGroupEnabled {
		return nil, fmt.Errorf("composite pod group feature gate is disabled")
	}
	l.cache.mu.RLock()
	defer l.cache.mu.RUnlock()

	key := fwk.CompositePodGroupKey(namespace, podGroupName)
	cpgs, exists := l.cache.compositePodGroupStates[key]
	if !exists {
		return nil, fmt.Errorf("composite pod group state not found for composite pod group %s", key)
	}
	return cpgs, nil
}

// CompositePodGroups returns the CompositePodGroupLister for this cache.
func (cache *cacheImpl) CompositePodGroups() fwk.CompositePodGroupLister {
	return &compositePodGroupListerImpl{cache: cache}
}

type compositePodGroupListerImpl struct {
	cache *cacheImpl
}

// Get returns the cached pod group object.
func (l *compositePodGroupListerImpl) Get(namespace, name string) (*schedulingv1alpha3.CompositePodGroup, error) {
	if !l.cache.compositePodGroupEnabled {
		return nil, fmt.Errorf("composite pod group feature gate is disabled")
	}
	l.cache.mu.RLock()
	defer l.cache.mu.RUnlock()

	key := fwk.CompositePodGroupKey(namespace, name)
	cpgs, exists := l.cache.compositePodGroupStates[key]
	if !exists {
		return nil, fmt.Errorf("composite pod group state not found for composite pod group %s", key)
	}
	cpg := cpgs.compositePodGroup
	if cpg == nil {
		return nil, fmt.Errorf("composite pod group object not found for composite pod group %s", key)
	}
	return cpg, nil
}

// BindPod handles the pod binding by adding a bind API call to the dispatcher.
// This method should be used only if the SchedulerAsyncAPICalls feature gate is enabled.
func (cache *cacheImpl) BindPod(binding *v1.Binding) (<-chan error, error) {
	// Don't store anything in the cache, as the pod is already assumed, and in case of a binding failure, it will be forgotten.
	onFinish := make(chan error, 1)
	err := cache.apiDispatcher.Add(apicalls.Implementations.PodBinding(binding), fwk.APICallOptions{
		OnFinish: onFinish,
	})
	if fwk.IsUnexpectedError(err) {
		return onFinish, err
	}
	return onFinish, nil
}

// refreshPVCRefCountsDelta accumulates the given node's PVC reference counts
// into cache.pvcRefCountsDelta, which is later applied to the snapshot during
// UpdateSnapshot. sign should be +1 to add the node's contribution or -1 to
// remove it.
func (cache *cacheImpl) refreshPVCRefCountsDelta(nodeInfo *framework.NodeInfo, sign int) {
	for key, count := range nodeInfo.PVCRefCounts {
		cache.pvcRefCountsDelta[key] += sign * count
	}
}

// applyPVCRefCountDelta merges cache.pvcRefCountsDelta into the snapshot's
// PVC ref counts, removes entries that reach zero, and clears the delta.
func (cache *cacheImpl) applyPVCRefCountDelta(snapshot *Snapshot) error {
	for key, delta := range cache.pvcRefCountsDelta {
		snapshot.usedPVCRefCounts[key] += delta
		if refCount := snapshot.usedPVCRefCounts[key]; refCount <= 0 {
			if refCount < 0 {
				delete(snapshot.usedPVCRefCounts, key)
				return fmt.Errorf("PVC %s had negative ref count %v", key, refCount)
			}
			delete(snapshot.usedPVCRefCounts, key)
		}
	}
	cache.pvcRefCountsDelta = map[string]int{}
	return nil
}

// AddPodGroup adds a pod group object to the cache.
func (cache *cacheImpl) AddPodGroup(podGroup *schedulingv1beta1.PodGroup) {
	if !cache.genericWorkloadEnabled {
		return
	}
	cache.mu.Lock()
	defer cache.mu.Unlock()

	key := fwk.PodGroupKey(podGroup.Namespace, podGroup.Name)
	pgs, exists := cache.podGroupStates[key]
	if !exists {
		pgs = newPodGroupState()
		cache.podGroupStates[key] = pgs
	}
	pgs.setPodGroup(podGroup)

	if !cache.compositePodGroupEnabled || podGroup.Spec.ParentCompositePodGroupName == nil {
		return
	}

	parentKey := fwk.CompositePodGroupKey(podGroup.Namespace, *podGroup.Spec.ParentCompositePodGroupName)

	// Even if the parent composite pod group has not been observed in the cache yet, we can assume that it
	// will be observed eventually, and we should create an entry (placeholder) representing it in the cache.
	parent, parentExists := cache.compositePodGroupStates[parentKey]
	if !parentExists {
		parent = newCompositePodGroupState()
		cache.compositePodGroupStates[parentKey] = parent
	}
	parent.addChild(key)
}

// UpdatePodGroup updates a pod group object in the cache.
func (cache *cacheImpl) UpdatePodGroup(logger klog.Logger, oldPodGroup, newPodGroup *schedulingv1beta1.PodGroup) {
	if !cache.genericWorkloadEnabled {
		return
	}
	cache.mu.Lock()
	defer cache.mu.Unlock()

	key := fwk.PodGroupKey(newPodGroup.Namespace, newPodGroup.Name)
	pgs, exists := cache.podGroupStates[key]
	if !exists {
		// This should not happen: the pod group state should have been already created by a prior pod group add action.
		utilruntime.HandleErrorWithLogger(logger, nil, "Pod group state not found for update, this indicates a missed add event", "podGroup", klog.KObj(newPodGroup))
		return
	}
	pgs.setPodGroup(newPodGroup)
}

// RemovePodGroup removes a pod group object from the cache.
func (cache *cacheImpl) RemovePodGroup(logger klog.Logger, podGroup *schedulingv1beta1.PodGroup) {
	if !cache.genericWorkloadEnabled {
		return
	}
	cache.mu.Lock()
	defer cache.mu.Unlock()

	key := fwk.PodGroupKey(podGroup.Namespace, podGroup.Name)
	pgs, exists := cache.podGroupStates[key]
	if !exists {
		// This should not happen: the pod group state should have been already created by a prior pod group add action.
		utilruntime.HandleErrorWithLogger(logger, nil, "Pod group state not found for removal", "podGroup", klog.KObj(podGroup))
		return
	}

	if cache.compositePodGroupEnabled && podGroup.Spec.ParentCompositePodGroupName != nil {
		parentKey := fwk.CompositePodGroupKey(podGroup.Namespace, *podGroup.Spec.ParentCompositePodGroupName)
		if parent, exists := cache.compositePodGroupStates[parentKey]; exists {
			parent.removeChild(key)
			if parent.empty() {
				// podGroupState can exist without the PodGroup object, but when any of the member pods exists.
				delete(cache.compositePodGroupStates, parentKey)
			}
		}
	}

	pgs.removePodGroup()
	if pgs.empty() {
		delete(cache.podGroupStates, key)
	}
}

// AddCompositePodGroup adds a composite pod group to the cache.
func (cache *cacheImpl) AddCompositePodGroup(logger klog.Logger, cpg *schedulingv1alpha3.CompositePodGroup) {
	if !cache.compositePodGroupEnabled {
		return
	}
	cache.mu.Lock()
	defer cache.mu.Unlock()

	key := fwk.CompositePodGroupKey(cpg.Namespace, cpg.Name)
	cpgState, exists := cache.compositePodGroupStates[key]
	if !exists {
		cpgState = newCompositePodGroupState()
		cache.compositePodGroupStates[key] = cpgState
	}
	cpgState.setCompositePodGroup(cpg)

	if cpg.Spec.ParentCompositePodGroupName == nil {
		return
	}

	parentKey := fwk.CompositePodGroupKey(cpg.Namespace, *cpg.Spec.ParentCompositePodGroupName)

	parent, parentExists := cache.compositePodGroupStates[parentKey]
	if !parentExists {
		parent = newCompositePodGroupState()
		cache.compositePodGroupStates[parentKey] = parent
	}
	parent.addChild(key)
}

// UpdateCompositePodGroup updates a composite pod group in the cache.
func (cache *cacheImpl) UpdateCompositePodGroup(logger klog.Logger, oldCPG, newCPG *schedulingv1alpha3.CompositePodGroup) {
	if !cache.compositePodGroupEnabled {
		return
	}
	cache.mu.Lock()
	defer cache.mu.Unlock()
	key := fwk.CompositePodGroupKey(newCPG.Namespace, newCPG.Name)
	pgs, exists := cache.compositePodGroupStates[key]
	if !exists {
		// This should not happen: the pod group state should have been already created by a prior pod group add action.
		utilruntime.HandleErrorWithLogger(logger, nil, "Pod group state not found for update, this indicates a missed add event", "podGroup", klog.KObj(newCPG))
		return
	}
	pgs.setCompositePodGroup(newCPG)
}

// RemoveCompositePodGroup removes a composite pod group from the cache.
func (cache *cacheImpl) RemoveCompositePodGroup(logger klog.Logger, cpg *schedulingv1alpha3.CompositePodGroup) {
	if !cache.compositePodGroupEnabled {
		return
	}
	cache.mu.Lock()
	defer cache.mu.Unlock()

	key := fwk.CompositePodGroupKey(cpg.Namespace, cpg.Name)
	cpgs, exists := cache.compositePodGroupStates[key]
	if !exists {
		// This should not happen: the pod group state should have been already created by a prior pod group add action.
		utilruntime.HandleErrorWithLogger(logger, nil, "Composite pod group state not found for removal", "compositePodGroup", klog.KObj(cpg))
		return
	}

	if cpgs.compositePodGroup != nil && cpgs.compositePodGroup.Spec.ParentCompositePodGroupName != nil {
		parentKey := fwk.CompositePodGroupKey(cpg.Namespace, *cpgs.compositePodGroup.Spec.ParentCompositePodGroupName)
		parent, exists := cache.compositePodGroupStates[parentKey]
		if !exists {
			return
		}
		parent.removeChild(key)
		if parent.empty() {
			delete(cache.compositePodGroupStates, parentKey)
		}
	}

	cpgs.removeCompositePodGroup()
	if cpgs.empty() {
		delete(cache.compositePodGroupStates, key)
	}
}

// BuildHierarchySnapshotFromPod returns a snapshot of the pod group hierarchy for the given pod.
func (cache *cacheImpl) BuildHierarchySnapshotFromPod(pod *v1.Pod) (fwk.PodGroupManager, error) {
	if pod.Spec.SchedulingGroup == nil {
		return nil, fmt.Errorf("pod has no scheduling group")
	}

	cache.mu.RLock()
	defer cache.mu.RUnlock()

	// 1. Find root CPG/PG by traversing upwards
	currentKey := fwk.PodGroupKey(pod.Namespace, *pod.Spec.SchedulingGroup.PodGroupName)
	pgs, exists := cache.podGroupStates[currentKey]
	if !exists {
		return nil, fmt.Errorf("pod group state not found for %s", currentKey.String())
	}

	pg := pgs.podGroup
	if pg == nil {
		return nil, fmt.Errorf("pod group object not found in state for %s", currentKey.String())
	}

	if cache.compositePodGroupEnabled && pg.Spec.ParentCompositePodGroupName != nil {
		currentKey = fwk.CompositePodGroupKey(pod.Namespace, *pg.Spec.ParentCompositePodGroupName)
		for range schedulingv1alpha3.WorkloadMaxTreeDepth - 1 {
			cpgs, exists := cache.compositePodGroupStates[currentKey]
			if !exists {
				return nil, fmt.Errorf("parent composite pod group state not found for %s", currentKey.String())
			}
			cpg := cpgs.compositePodGroup
			if cpg == nil {
				return nil, fmt.Errorf("composite pod group object not found in state for %s", currentKey.String())
			}
			if cpg.Spec.ParentCompositePodGroupName == nil {
				break
			}
			currentKey = fwk.CompositePodGroupKey(pod.Namespace, *cpg.Spec.ParentCompositePodGroupName)
		}
	}

	// 2. We have the root key. Now traverse downwards and update the snapshot.
	snapshot := NewEmptySnapshot()
	visited := sets.New[fwk.EntityKey]()
	err := cache.buildPodGroupStateSnapshotTree(currentKey, snapshot, visited)
	if err != nil {
		return nil, err
	}
	return snapshot, nil
}

// buildPodGroupStateSnapshotTree recursively builds a snapshot of the pod group state tree starting from the given key.
// It assumes that the cache lock is held by the caller.
func (cache *cacheImpl) buildPodGroupStateSnapshotTree(key fwk.EntityKey, snapshot *Snapshot, visited sets.Set[fwk.EntityKey]) error {
	if visited.Has(key) {
		return fmt.Errorf("cycle detected in composite pod group hierarchy: %s", key.String())
	}
	visited.Insert(key)

	switch key.Type {
	case fwk.PodGroupKeyType:
		pgs, exists := cache.podGroupStates[key]
		if !exists {
			return fmt.Errorf("pod group state not found for %s", key.String())
		}
		snapshot.podGroupStates[key] = &podGroupStateSnapshot{podGroupStateData: pgs.podGroupStateData.clone()}

	case fwk.CompositePodGroupKeyType:
		cpgs, exists := cache.compositePodGroupStates[key]
		if !exists {
			return fmt.Errorf("composite pod group state not found for %s", key.String())
		}
		snapshot.compositePodGroupStates[key] = &compositePodGroupStateSnapshot{compositePodGroupStateData: cpgs.compositePodGroupStateData.clone()}

		children := cpgs.children.Clone()
		for childKey := range children {
			if err := cache.buildPodGroupStateSnapshotTree(childKey, snapshot, visited); err != nil {
				return err
			}
		}
	}

	return nil
}

// GetRootKeyForGroup returns the root key of the given EntityKey.
// The key must be of PodGroupKey or CompositePodGroupKey type.
func (cache *cacheImpl) GetRootKeyForGroup(key fwk.EntityKey) (fwk.EntityKey, bool, error) {
	cache.mu.RLock()
	defer cache.mu.RUnlock()

	currentKey := key
	visited := sets.New[fwk.EntityKey]()
	for {
		if visited.Has(currentKey) {
			return fwk.EntityKey{}, false, fmt.Errorf("cycle detected in the hierarchy: %v", visited.UnsortedList())
		}
		visited.Insert(currentKey)

		switch currentKey.Type {
		case fwk.PodGroupKeyType:
			pgs, exists := cache.podGroupStates[currentKey]
			if !exists {
				return fwk.EntityKey{}, false, nil
			}
			pg := pgs.podGroup
			if pg == nil {
				return fwk.EntityKey{}, false, nil
			}
			if !cache.compositePodGroupEnabled || pg.Spec.ParentCompositePodGroupName == nil {
				return currentKey, true, nil
			}
			currentKey = fwk.CompositePodGroupKey(pg.Namespace, *pg.Spec.ParentCompositePodGroupName)
		case fwk.CompositePodGroupKeyType:
			cpgs, exists := cache.compositePodGroupStates[currentKey]
			if !exists {
				return fwk.EntityKey{}, false, nil
			}
			cpg := cpgs.compositePodGroup
			if cpg == nil {
				return fwk.EntityKey{}, false, nil
			}
			if cpg.Spec.ParentCompositePodGroupName == nil {
				return currentKey, true, nil
			}
			currentKey = fwk.CompositePodGroupKey(cpg.Namespace, *cpg.Spec.ParentCompositePodGroupName)
		case fwk.PodKeyType:
			return fwk.EntityKey{}, false, fmt.Errorf("pod key type not supported in GetRootKeyForGroup for %s", currentKey.String())
		}
	}
}