github.com/kubewharf/katalyst-core@v0.5.3/pkg/agent/evictionmanager/rule/queue.go (about)

     1  /*
     2  Copyright 2022 The Katalyst Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package rule
    18  
    19  import (
    20  	"sync"
    21  
    22  	"k8s.io/apimachinery/pkg/types"
    23  )
    24  
    25  // EvictionQueue aims to build a queue for eviction EvictPods, based on this
    26  // queue, eviction manager can use it to perform more efficient logics, e.g.
    27  // rate limiter, priority-based sorting strategy, eviction withdraw and so on
    28  type EvictionQueue interface {
    29  	// Add EvictPods into queue, whether to perform sorting logic should be added here;
    30  	// and the parameter "override" will decide whether we should withdraw all
    31  	// candidate EvictPods that don't appear in the given pod list
    32  	Add(rpList RuledEvictPodList, override bool)
    33  	// Withdraw withdraws EvictPods from candidate list if exists
    34  	Withdraw(rpList RuledEvictPodList)
    35  	// Pop EvictPods for eviction
    36  	Pop() RuledEvictPodList
    37  
    38  	// List returns all EvictPods in this queue (without deleting)
    39  	List() RuledEvictPodList
    40  }
    41  
    42  // FIFOEvictionQueue is the default implementation for EvictionQueue;
    43  // it will use the FIFO strategy for one batch of EvictPods (sort for
    44  // those EvictPods in the same way though), and support to pop a limited amount
    45  // of EvictPods when pop is triggered.
    46  type FIFOEvictionQueue struct {
    47  	// limited set slo
    48  	limited int
    49  
    50  	sync.Mutex
    51  	rpList RuledEvictPodList
    52  	podIDs map[types.UID]interface{}
    53  }
    54  
    55  func NewFIFOEvictionQueue(limited int) EvictionQueue {
    56  	return &FIFOEvictionQueue{
    57  		limited: limited,
    58  		podIDs:  make(map[types.UID]interface{}),
    59  	}
    60  }
    61  
    62  func (f *FIFOEvictionQueue) Add(rpList RuledEvictPodList, override bool) {
    63  	f.Lock()
    64  	defer f.Unlock()
    65  
    66  	if override {
    67  		f.rpList = RuledEvictPodList{}
    68  		f.podIDs = make(map[types.UID]interface{})
    69  	}
    70  
    71  	for _, ep := range rpList {
    72  		if _, ok := f.podIDs[ep.Pod.UID]; !ok {
    73  			f.podIDs[ep.Pod.UID] = struct{}{}
    74  			f.rpList = append(f.rpList, ep)
    75  		}
    76  	}
    77  }
    78  
    79  func (f *FIFOEvictionQueue) Withdraw(rpList RuledEvictPodList) {
    80  	if len(rpList) == 0 {
    81  		return
    82  	}
    83  
    84  	deleteSets := make(map[types.UID]interface{})
    85  	for _, ep := range rpList {
    86  		deleteSets[ep.Pod.UID] = struct{}{}
    87  	}
    88  
    89  	f.Lock()
    90  	defer f.Unlock()
    91  
    92  	newRpList := RuledEvictPodList{}
    93  	evictPodIDs := make(map[types.UID]interface{})
    94  	for _, ep := range f.rpList {
    95  		if _, ok := deleteSets[ep.Pod.UID]; !ok {
    96  			newRpList = append(newRpList, ep)
    97  			evictPodIDs[ep.Pod.UID] = struct{}{}
    98  		}
    99  	}
   100  
   101  	f.rpList = newRpList
   102  	f.podIDs = evictPodIDs
   103  }
   104  
   105  func (f *FIFOEvictionQueue) Pop() RuledEvictPodList {
   106  	f.Lock()
   107  	defer f.Unlock()
   108  
   109  	amount := f.limited
   110  	if amount < 0 || amount > len(f.rpList) {
   111  		amount = f.rpList.Len()
   112  	}
   113  
   114  	rpList := f.rpList[:amount]
   115  	f.rpList = f.rpList[amount:]
   116  	for _, ep := range rpList {
   117  		delete(f.podIDs, ep.Pod.UID)
   118  	}
   119  
   120  	return rpList
   121  }
   122  
   123  func (f *FIFOEvictionQueue) List() RuledEvictPodList {
   124  	f.Lock()
   125  	defer f.Unlock()
   126  
   127  	return f.rpList
   128  }