koord-scheduler: refactor ReservationInfo as public

pkg/scheduler/frameworkext/reservation_info.go

@@ -0,0 +1,106 @@
+/*
+Copyright 2022 The Koordinator 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 frameworkext
+
+import (
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/types"
+	quotav1 "k8s.io/apiserver/pkg/quota/v1"
+	"k8s.io/kubernetes/pkg/api/v1/resource"
+
+	schedulingv1alpha1 "github.com/koordinator-sh/koordinator/apis/scheduling/v1alpha1"
+	reservationutil "github.com/koordinator-sh/koordinator/pkg/util/reservation"
+)
+
+type ReservationInfo struct {
+	Reservation   *schedulingv1alpha1.Reservation
+	ResourceNames []corev1.ResourceName
+	Allocatable   corev1.ResourceList
+	Allocated     corev1.ResourceList
+	Pods          map[types.UID]*PodRequirement
+}
+
+type PodRequirement struct {
+	Namespace string
+	Name      string
+	UID       types.UID
+	Requests  corev1.ResourceList
+}
+
+func NewReservationInfo(r *schedulingv1alpha1.Reservation) *ReservationInfo {
+	allocatable := reservationutil.ReservationRequests(r)
+	resourceNames := quotav1.ResourceNames(allocatable)
+
+	return &ReservationInfo{
+		Reservation:   r.DeepCopy(),
+		ResourceNames: resourceNames,
+		Allocatable:   allocatable,
+		Pods:          map[types.UID]*PodRequirement{},
+	}
+}
+
+func (ri *ReservationInfo) Clone() *ReservationInfo {
+	resourceNames := make([]corev1.ResourceName, 0, len(ri.ResourceNames))
+	for _, v := range ri.ResourceNames {
+		resourceNames = append(resourceNames, v)
+	}
+
+	pods := map[types.UID]*PodRequirement{}
+	for k, v := range ri.Pods {
+		pods[k] = &PodRequirement{
+			Namespace: v.Namespace,
+			Name:      v.Name,
+			UID:       v.UID,
+			Requests:  v.Requests.DeepCopy(),
+		}
+	}
+
+	return &ReservationInfo{
+		Reservation:   ri.Reservation.DeepCopy(),
+		ResourceNames: resourceNames,
+		Allocatable:   ri.Allocatable.DeepCopy(),
+		Allocated:     ri.Allocated.DeepCopy(),
+		Pods:          pods,
+	}
+}
+
+func (ri *ReservationInfo) UpdateReservation(r *schedulingv1alpha1.Reservation) {
+	ri.Reservation = r.DeepCopy()
+	ri.Allocatable = reservationutil.ReservationRequests(r)
+	ri.ResourceNames = quotav1.ResourceNames(ri.Allocatable)
+	ri.Allocated = quotav1.Mask(ri.Allocated, ri.ResourceNames)
+}
+
+func (ri *ReservationInfo) AddPod(pod *corev1.Pod) {
+	requests, _ := resource.PodRequestsAndLimits(pod)
+	ri.Allocated = quotav1.Add(ri.Allocated, quotav1.Mask(requests, ri.ResourceNames))
+	ri.Pods[pod.UID] = &PodRequirement{
+		Namespace: pod.Namespace,
+		Name:      pod.Name,
+		UID:       pod.UID,
+		Requests:  requests,
+	}
+}
+
+func (ri *ReservationInfo) RemovePod(pod *corev1.Pod) {
+	if requirement, ok := ri.Pods[pod.UID]; ok {
+		if len(requirement.Requests) > 0 {
+			ri.Allocated = quotav1.SubtractWithNonNegativeResult(ri.Allocated, quotav1.Mask(requirement.Requests, ri.ResourceNames))
+		}
+		delete(ri.Pods, pod.UID)
+	}
+}

pkg/scheduler/plugins/reservation/cache.go

@@ -24,19 +24,20 @@ import (
 
 	schedulingv1alpha1 "github.com/koordinator-sh/koordinator/apis/scheduling/v1alpha1"
 	schedulinglister "github.com/koordinator-sh/koordinator/pkg/client/listers/scheduling/v1alpha1"
+	"github.com/koordinator-sh/koordinator/pkg/scheduler/frameworkext"
 )
 
 type reservationCache struct {
 	reservationLister  schedulinglister.ReservationLister
 	lock               sync.Mutex
-	reservationInfos   map[types.UID]*reservationInfo
+	reservationInfos   map[types.UID]*frameworkext.ReservationInfo
 	reservationsOnNode map[string]map[types.UID]*schedulingv1alpha1.Reservation
 }
 
 func newReservationCache(reservationLister schedulinglister.ReservationLister) *reservationCache {
 	cache := &reservationCache{
 		reservationLister:  reservationLister,
-		reservationInfos:   map[types.UID]*reservationInfo{},
+		reservationInfos:   map[types.UID]*frameworkext.ReservationInfo{},
 		reservationsOnNode: map[string]map[types.UID]*schedulingv1alpha1.Reservation{},
 	}
 	return cache
@@ -79,10 +80,10 @@ func (cache *reservationCache) updateReservation(newR *schedulingv1alpha1.Reserv
 	defer cache.lock.Unlock()
 	rInfo := cache.reservationInfos[newR.UID]
 	if rInfo == nil {
-		rInfo = newReservationInfo(newR)
+		rInfo = frameworkext.NewReservationInfo(newR)
 		cache.reservationInfos[newR.UID] = rInfo
 	} else {
-		rInfo.updateReservation(newR)
+		rInfo.UpdateReservation(newR)
 	}
 	if newR.Status.NodeName != "" {
 		cache.updateReservationsOnNode(newR.Status.NodeName, newR)
@@ -110,7 +111,7 @@ func (cache *reservationCache) addPod(reservationUID types.UID, pod *corev1.Pod)
 
 	rInfo := cache.reservationInfos[reservationUID]
 	if rInfo != nil {
-		rInfo.addPod(pod)
+		rInfo.AddPod(pod)
 	}
 }
 
@@ -120,8 +121,8 @@ func (cache *reservationCache) updatePod(reservationUID types.UID, oldPod, newPo
 
 	rInfo := cache.reservationInfos[reservationUID]
 	if rInfo != nil {
-		rInfo.removePod(oldPod)
-		rInfo.addPod(newPod)
+		rInfo.RemovePod(oldPod)
+		rInfo.AddPod(newPod)
 	}
 }
 
@@ -131,19 +132,19 @@ func (cache *reservationCache) deletePod(reservationUID types.UID, pod *corev1.P
 
 	rInfo := cache.reservationInfos[reservationUID]
 	if rInfo != nil {
-		rInfo.removePod(pod)
+		rInfo.RemovePod(pod)
 	}
 }
 
-func (cache *reservationCache) getReservationInfo(name string) *reservationInfo {
+func (cache *reservationCache) getReservationInfo(name string) *frameworkext.ReservationInfo {
 	reservation, err := cache.reservationLister.Get(name)
 	if err != nil {
 		return nil
 	}
 	return cache.getReservationInfoByUID(reservation.UID)
 }
 
-func (cache *reservationCache) getReservationInfoByUID(uid types.UID) *reservationInfo {
+func (cache *reservationCache) getReservationInfoByUID(uid types.UID) *frameworkext.ReservationInfo {
 	cache.lock.Lock()
 	defer cache.lock.Unlock()
 	rInfo := cache.reservationInfos[uid]
@@ -153,14 +154,14 @@ func (cache *reservationCache) getReservationInfoByUID(uid types.UID) *reservati
 	return nil
 }
 
-func (cache *reservationCache) listReservationInfosOnNode(nodeName string) []*reservationInfo {
+func (cache *reservationCache) listReservationInfosOnNode(nodeName string) []*frameworkext.ReservationInfo {
 	cache.lock.Lock()
 	defer cache.lock.Unlock()
 	rOnNode := cache.reservationsOnNode[nodeName]
 	if len(rOnNode) == 0 {
 		return nil
 	}
-	result := make([]*reservationInfo, 0, len(rOnNode))
+	result := make([]*frameworkext.ReservationInfo, 0, len(rOnNode))
 	for _, v := range rOnNode {
 		rInfo := cache.reservationInfos[v.UID]
 		if rInfo != nil {

pkg/scheduler/plugins/reservation/cache_test.go

@@ -29,6 +29,7 @@ import (
 	quotav1 "k8s.io/apiserver/pkg/quota/v1"
 
 	schedulingv1alpha1 "github.com/koordinator-sh/koordinator/apis/scheduling/v1alpha1"
+	"github.com/koordinator-sh/koordinator/pkg/scheduler/frameworkext"
 )
 
 func TestCacheUpdateReservation(t *testing.T) {
@@ -70,31 +71,31 @@ func TestCacheUpdateReservation(t *testing.T) {
 	reservationInfos := cache.listReservationInfosOnNode(reservation.Status.NodeName)
 	assert.Len(t, reservationInfos, 1)
 	rInfo := reservationInfos[0]
-	expectReservationInfo := &reservationInfo{
-		reservation: reservation,
-		resourceNames: []corev1.ResourceName{
+	expectReservationInfo := &frameworkext.ReservationInfo{
+		Reservation: reservation,
+		ResourceNames: []corev1.ResourceName{
 			corev1.ResourceCPU,
 			corev1.ResourceMemory,
 		},
-		allocatable: corev1.ResourceList{
+		Allocatable: corev1.ResourceList{
 			corev1.ResourceCPU:    resource.MustParse("4"),
 			corev1.ResourceMemory: resource.MustParse("4Gi"),
 		},
-		allocated: nil,
-		pods:      map[types.UID]*podRequirement{},
+		Allocated: nil,
+		Pods:      map[types.UID]*frameworkext.PodRequirement{},
 	}
-	sort.Slice(rInfo.resourceNames, func(i, j int) bool {
-		return rInfo.resourceNames[i] < rInfo.resourceNames[j]
+	sort.Slice(rInfo.ResourceNames, func(i, j int) bool {
+		return rInfo.ResourceNames[i] < rInfo.ResourceNames[j]
 	})
 	assert.Equal(t, expectReservationInfo, rInfo)
 
 	cache.updateReservation(reservation)
 	reservationInfos = cache.listReservationInfosOnNode(reservation.Status.NodeName)
 	assert.Len(t, reservationInfos, 1)
 	rInfo = reservationInfos[0]
-	expectReservationInfo.allocated = corev1.ResourceList{}
-	sort.Slice(rInfo.resourceNames, func(i, j int) bool {
-		return rInfo.resourceNames[i] < rInfo.resourceNames[j]
+	expectReservationInfo.Allocated = corev1.ResourceList{}
+	sort.Slice(rInfo.ResourceNames, func(i, j int) bool {
+		return rInfo.ResourceNames[i] < rInfo.ResourceNames[j]
 	})
 	assert.Equal(t, expectReservationInfo, rInfo)
 }
@@ -139,21 +140,21 @@ func TestCacheDeleteReservation(t *testing.T) {
 	rInfo := cache.getReservationInfoByUID(reservation.UID)
 	assert.NotNil(t, rInfo)
 
-	expectReservationInfo := &reservationInfo{
-		reservation: reservation,
-		resourceNames: []corev1.ResourceName{
+	expectReservationInfo := &frameworkext.ReservationInfo{
+		Reservation: reservation,
+		ResourceNames: []corev1.ResourceName{
 			corev1.ResourceCPU,
 			corev1.ResourceMemory,
 		},
-		allocatable: corev1.ResourceList{
+		Allocatable: corev1.ResourceList{
 			corev1.ResourceCPU:    resource.MustParse("4"),
 			corev1.ResourceMemory: resource.MustParse("4Gi"),
 		},
-		allocated: nil,
-		pods:      map[types.UID]*podRequirement{},
+		Allocated: nil,
+		Pods:      map[types.UID]*frameworkext.PodRequirement{},
 	}
-	sort.Slice(rInfo.resourceNames, func(i, j int) bool {
-		return rInfo.resourceNames[i] < rInfo.resourceNames[j]
+	sort.Slice(rInfo.ResourceNames, func(i, j int) bool {
+		return rInfo.ResourceNames[i] < rInfo.ResourceNames[j]
 	})
 	assert.Equal(t, expectReservationInfo, rInfo)
 
@@ -225,29 +226,29 @@ func TestCacheAddOrUpdateOrDeletePod(t *testing.T) {
 	cache.addPod(reservation.UID, pod)
 
 	rInfo = cache.getReservationInfoByUID(reservation.UID)
-	sort.Slice(rInfo.resourceNames, func(i, j int) bool {
-		return rInfo.resourceNames[i] < rInfo.resourceNames[j]
+	sort.Slice(rInfo.ResourceNames, func(i, j int) bool {
+		return rInfo.ResourceNames[i] < rInfo.ResourceNames[j]
 	})
-	expectReservationInfo := &reservationInfo{
-		reservation: reservation,
-		resourceNames: []corev1.ResourceName{
+	expectReservationInfo := &frameworkext.ReservationInfo{
+		Reservation: reservation,
+		ResourceNames: []corev1.ResourceName{
 			corev1.ResourceCPU,
 			corev1.ResourceMemory,
 		},
-		allocatable: corev1.ResourceList{
+		Allocatable: corev1.ResourceList{
 			corev1.ResourceCPU:    resource.MustParse("4000m"),
 			corev1.ResourceMemory: resource.MustParse("4Gi"),
 		},
-		allocated: corev1.ResourceList{
+		Allocated: corev1.ResourceList{
 			corev1.ResourceCPU:    resource.MustParse("2000m"),
 			corev1.ResourceMemory: resource.MustParse("2Gi"),
 		},
-		pods: map[types.UID]*podRequirement{
+		Pods: map[types.UID]*frameworkext.PodRequirement{
 			pod.UID: {
-				namespace: pod.Namespace,
-				name:      pod.Name,
-				uid:       pod.UID,
-				requests: corev1.ResourceList{
+				Namespace: pod.Namespace,
+				Name:      pod.Name,
+				UID:       pod.UID,
+				Requests: corev1.ResourceList{
 					corev1.ResourceCPU:    resource.MustParse("2000m"),
 					corev1.ResourceMemory: resource.MustParse("2Gi"),
 				},
@@ -258,24 +259,24 @@ func TestCacheAddOrUpdateOrDeletePod(t *testing.T) {
 
 	cache.updatePod(reservation.UID, pod, pod)
 	rInfo = cache.getReservationInfoByUID(reservation.UID)
-	sort.Slice(rInfo.resourceNames, func(i, j int) bool {
-		return rInfo.resourceNames[i] < rInfo.resourceNames[j]
+	sort.Slice(rInfo.ResourceNames, func(i, j int) bool {
+		return rInfo.ResourceNames[i] < rInfo.ResourceNames[j]
 	})
-	expectReservationInfo.allocated = quotav1.SubtractWithNonNegativeResult(expectReservationInfo.allocated, corev1.ResourceList{
+	expectReservationInfo.Allocated = quotav1.SubtractWithNonNegativeResult(expectReservationInfo.Allocated, corev1.ResourceList{
 		corev1.ResourceCPU:    resource.MustParse("0"),
 		corev1.ResourceMemory: resource.MustParse("0"),
 	})
 	assert.Equal(t, expectReservationInfo, rInfo)
 
 	cache.deletePod(reservation.UID, pod)
 	rInfo = cache.getReservationInfoByUID(reservation.UID)
-	sort.Slice(rInfo.resourceNames, func(i, j int) bool {
-		return rInfo.resourceNames[i] < rInfo.resourceNames[j]
+	sort.Slice(rInfo.ResourceNames, func(i, j int) bool {
+		return rInfo.ResourceNames[i] < rInfo.ResourceNames[j]
 	})
-	expectReservationInfo.allocated = corev1.ResourceList{
+	expectReservationInfo.Allocated = corev1.ResourceList{
 		corev1.ResourceCPU:    resource.MustParse("0"),
 		corev1.ResourceMemory: resource.MustParse("0"),
 	}
-	expectReservationInfo.pods = map[types.UID]*podRequirement{}
+	expectReservationInfo.Pods = map[types.UID]*frameworkext.PodRequirement{}
 	assert.Equal(t, expectReservationInfo, rInfo)
 }

pkg/scheduler/plugins/reservation/eventhandler_test.go

@@ -106,7 +106,7 @@ func TestEventHandlerOnAdd(t *testing.T) {
 				assert.Nil(t, rInfo)
 			} else {
 				rInfo := cache.getReservationInfoByUID(tt.wantReservation.UID)
-				assert.Equal(t, tt.wantReservation, rInfo.reservation)
+				assert.Equal(t, tt.wantReservation, rInfo.Reservation)
 			}
 		})
 	}
@@ -201,7 +201,7 @@ func TestEventHandlerUpdate(t *testing.T) {
 				assert.Nil(t, rInfo)
 			} else {
 				rInfo := cache.getReservationInfoByUID(tt.wantReservation.UID)
-				assert.Equal(t, tt.wantReservation, rInfo.reservation)
+				assert.Equal(t, tt.wantReservation, rInfo.Reservation)
 			}
 		})
 	}

pkg/scheduler/plugins/reservation/nominator.go

@@ -53,11 +53,11 @@ func (pl *Plugin) NominateReservation(ctx context.Context, cycleState *framework
 
 	reservations := make([]*schedulingv1alpha1.Reservation, 0, len(rOnNode))
 	for _, rInfo := range rOnNode {
-		status := extender.RunReservationFilterPlugins(ctx, cycleState, pod, rInfo.reservation, nodeName)
+		status := extender.RunReservationFilterPlugins(ctx, cycleState, pod, rInfo.Reservation, nodeName)
 		if !status.IsSuccess() {
 			continue
 		}
-		reservations = append(reservations, rInfo.reservation)
+		reservations = append(reservations, rInfo.Reservation)
 	}
 	if len(reservations) == 0 {
 		return nil, nil