volcano.sh/volcano@v1.9.0/pkg/scheduler/cache/cache_mock.go (about)

     1  package cache
     2  
     3  import (
     4  	"math"
     5  	"os"
     6  	"strconv"
     7  
     8  	schedulingv1 "k8s.io/api/scheduling/v1"
     9  	"k8s.io/client-go/kubernetes/fake"
    10  	"k8s.io/client-go/tools/record"
    11  	"k8s.io/client-go/util/workqueue"
    12  
    13  	fakevcClient "volcano.sh/apis/pkg/client/clientset/versioned/fake"
    14  	"volcano.sh/volcano/cmd/scheduler/app/options"
    15  	schedulingapi "volcano.sh/volcano/pkg/scheduler/api"
    16  )
    17  
    18  // NewCustomMockSchedulerCache returns a mock scheduler cache with custom interface
    19  func NewCustomMockSchedulerCache(schedulerName string,
    20  	binder Binder,
    21  	evictor Evictor,
    22  	statusUpdater StatusUpdater,
    23  	PodGroupBinder BatchBinder,
    24  	volumeBinder VolumeBinder,
    25  	recorder record.EventRecorder,
    26  ) *SchedulerCache {
    27  	msc := newMockSchedulerCache(schedulerName)
    28  	// add all events handlers
    29  	msc.addEventHandler()
    30  	msc.Recorder = recorder
    31  	msc.Binder = binder
    32  	msc.Evictor = evictor
    33  	msc.StatusUpdater = statusUpdater
    34  	msc.PodGroupBinder = PodGroupBinder
    35  	// use custom volume binder
    36  	msc.VolumeBinder = volumeBinder
    37  	checkAndSetDefaultInterface(msc)
    38  	return msc
    39  }
    40  
    41  // NewDefaultMockSchedulerCache returns a mock scheduler cache with interface mocked with default fake clients
    42  // Notes that default events recorder's buffer only has a length 100;
    43  // when use it do performance test, should use a &FakeRecorder{} without length limit to avoid block
    44  func NewDefaultMockSchedulerCache(schedulerName string) *SchedulerCache {
    45  	msc := newMockSchedulerCache(schedulerName)
    46  	// add all events handlers
    47  	msc.addEventHandler()
    48  	checkAndSetDefaultInterface(msc)
    49  	return msc
    50  }
    51  
    52  func checkAndSetDefaultInterface(sc *SchedulerCache) {
    53  	if sc.Recorder == nil {
    54  		sc.Recorder = record.NewFakeRecorder(100) // to avoid blocking, we can pass in &FakeRecorder{} to NewCustomMockSchedulerCache
    55  	}
    56  	if sc.Binder == nil {
    57  		sc.Binder = &DefaultBinder{
    58  			kubeclient: sc.kubeClient,
    59  			recorder:   sc.Recorder,
    60  		}
    61  	}
    62  	if sc.Evictor == nil {
    63  		sc.Evictor = &defaultEvictor{
    64  			kubeclient: sc.kubeClient,
    65  			recorder:   sc.Recorder,
    66  		}
    67  	}
    68  	if sc.StatusUpdater == nil {
    69  		sc.StatusUpdater = &defaultStatusUpdater{
    70  			kubeclient: sc.kubeClient,
    71  			vcclient:   sc.vcClient,
    72  		}
    73  	}
    74  	if sc.PodGroupBinder == nil {
    75  		sc.PodGroupBinder = &podgroupBinder{
    76  			kubeclient: sc.kubeClient,
    77  			vcclient:   sc.vcClient,
    78  		}
    79  	}
    80  	// finally, init default fake volume binder which has dependencies on other informers
    81  	if sc.VolumeBinder == nil {
    82  		sc.setDefaultVolumeBinder()
    83  	}
    84  }
    85  
    86  func getNodeWorkers() uint32 {
    87  	if options.ServerOpts != nil && options.ServerOpts.NodeWorkerThreads > 0 {
    88  		return options.ServerOpts.NodeWorkerThreads
    89  	}
    90  	threads, err := strconv.Atoi(os.Getenv("NODE_WORKER_THREADS"))
    91  	if err == nil && threads > 0 && threads <= math.MaxUint32 {
    92  		return uint32(threads)
    93  	}
    94  	return 2 //default 2
    95  }
    96  
    97  // newMockSchedulerCache init the mock scheduler cache structure
    98  func newMockSchedulerCache(schedulerName string) *SchedulerCache {
    99  	msc := &SchedulerCache{
   100  		Jobs:                make(map[schedulingapi.JobID]*schedulingapi.JobInfo),
   101  		Nodes:               make(map[string]*schedulingapi.NodeInfo),
   102  		Queues:              make(map[schedulingapi.QueueID]*schedulingapi.QueueInfo),
   103  		PriorityClasses:     make(map[string]*schedulingv1.PriorityClass),
   104  		errTasks:            workqueue.NewRateLimitingQueue(workqueue.DefaultControllerRateLimiter()),
   105  		nodeQueue:           workqueue.NewRateLimitingQueue(workqueue.DefaultControllerRateLimiter()),
   106  		DeletedJobs:         workqueue.NewRateLimitingQueue(workqueue.DefaultControllerRateLimiter()),
   107  		kubeClient:          fake.NewSimpleClientset(),
   108  		vcClient:            fakevcClient.NewSimpleClientset(),
   109  		restConfig:          nil,
   110  		defaultQueue:        "default",
   111  		schedulerNames:      []string{schedulerName},
   112  		nodeSelectorLabels:  make(map[string]string),
   113  		NamespaceCollection: make(map[string]*schedulingapi.NamespaceCollection),
   114  		CSINodesStatus:      make(map[string]*schedulingapi.CSINodeStatusInfo),
   115  		imageStates:         make(map[string]*imageState),
   116  
   117  		NodeList: []string{},
   118  	}
   119  	if options.ServerOpts != nil && len(options.ServerOpts.NodeSelector) > 0 {
   120  		msc.updateNodeSelectors(options.ServerOpts.NodeSelector)
   121  	}
   122  	msc.setBatchBindParallel()
   123  	msc.nodeWorkers = getNodeWorkers()
   124  
   125  	return msc
   126  }