Repository: kubernetes/kubernetes
File: test/integration/scheduler/preemption/preemption_test.go (Go)

Code before:
}

// blockingPermitPlugin is a Permit plugin that blocks until a signal is received.
type blockingPermitPlugin struct {
	podsToBlock map[string]*blockedPod
}

const blockingPermitPluginName = "blocking-permit-plugin"

var _ fwk.PermitPlugin = &blockingPermitPlugin{}

func newBlockingPermitPlugin(_ context.Context, _ runtime.Object, h fwk.Handle) (fwk.Plugin, error) {
	return &blockingPermitPlugin{
		podsToBlock: make(map[string]*blockedPod),
	}, nil
}

func (pl *blockingPermitPlugin) Name() string {
	return blockingPermitPluginName
}

func (pl *blockingPermitPlugin) Permit(ctx context.Context, state fwk.CycleState, pod *v1.Pod, nodeName string) (*fwk.Status, time.Duration) {
	if p, ok := pl.podsToBlock[pod.Name]; ok {
		p.blocked <- struct{}{}
		delete(pl.podsToBlock, pod.Name)
		return fwk.NewStatus(fwk.Wait, "waiting"), 1 * time.Minute
	}
	return nil, 0
}

func TestPreemptionRespectsWaitingPod(t *testing.T) {
	// 1. Create a "blocking" permit plugin that signals when it's running and waits for a specific close.
	// 2. Schedule a low-priority pod (victim) that hits this plugin.
	// 3. While victim is blocked in WaitOnPermit, add a smaller node on which victim should be preempted.
	// 4. Schedule a high-priority pod (preemptor), that can only fit on big node
	// 5. High-priority pod should be scheduled and victim should be preempted.
	// 6. Victim should be either completely deleted or rescheduled or at least have the DisruptionTarget condition set.

	// Create a node with resources for only one pod.
	nodeRes := map[v1.ResourceName]string{
		v1.ResourceCPU:    "2",
		v1.ResourceMemory: "2Gi",
	}
	node := st.MakeNode().Name("node1").Capacity(nodeRes).Obj()

	// Victim requires full node resources.
	victim := st.MakePod().Name("victim").Priority(lowPriority).Req(map[v1.ResourceName]string{v1.ResourceCPU: "1", v1.ResourceMemory: "1Gi"}).Obj()
	// Preemptor also requires full node resources.
	preemptor := st.MakePod().Name("preemptor").Priority(highPriority).Req(map[v1.ResourceName]string{v1.ResourceCPU: "1.5", v1.ResourceMemory: "1.5Gi"}).Obj()

	// Register the blocking plugin

Review comment (suggestion):
nit:
```suggestion
		return fwk.NewStatus(fwk.Wait, "waiting"), time.Minute
```

Code after:
}

// blockingPermitPlugin is a Permit plugin that blocks until a signal is received.
type blockingPermitPlugin struct {
	podsToBlock map[string]*blockedPod
}

const blockingPermitPluginName = "blocking-permit-plugin"

var _ fwk.PermitPlugin = &blockingPermitPlugin{}

func newBlockingPermitPlugin(_ context.Context, _ runtime.Object, h fwk.Handle) fwk.Plugin {
	return &blockingPermitPlugin{
		podsToBlock: make(map[string]*blockedPod),
	}
}

func (pl *blockingPermitPlugin) Name() string {
	return blockingPermitPluginName
}

func (pl *blockingPermitPlugin) Permit(ctx context.Context, state fwk.CycleState, pod *v1.Pod, nodeName string) (*fwk.Status, time.Duration) {
	if p, ok := pl.podsToBlock[pod.Name]; ok {
		p.blocked <- struct{}{}
		delete(pl.podsToBlock, pod.Name)
		return fwk.NewStatus(fwk.Wait, "waiting"), time.Minute
	}
	return nil, 0
}

func TestPreemptionRespectsWaitingPod(t *testing.T) {
	// 1. Create a "blocking" permit plugin that signals when it's running and waits for a specific close.
	// 2. Create a big node on which low-priority pod will be scheduled.
	// 3. Schedule a low-priority pod (victim) that hits this plugin (after being selected to run on a big node).
	// 4. While victim is blocked in WaitOnPermit, add a smaller node on which the victim should be rescheduled.
	// 5. Schedule a high-priority pod (preemptor), that can only fit on big node.
	// 6. High-priority pod should be scheduled on a big node and victim should be preempted.
	// 7. Victim should be rescheduled on a smaller node.

	// Create a node with resources for only one pod.
	nodeRes := map[v1.ResourceName]string{
		v1.ResourceCPU:    "2",
		v1.ResourceMemory: "2Gi",
	}
	node := st.MakeNode().Name("big-node").Capacity(nodeRes).Obj()

	victim := st.MakePod().Name("victim").Priority(lowPriority).Req(map[v1.ResourceName]string{v1.ResourceCPU: "1", v1.ResourceMemory: "1Gi"}).Obj()
	// Preemptor requires more resources than the small node has.
	preemptor := st.MakePod().Name("preemptor").Priority(highPriority).Req(map[v1.ResourceName]string{v1.ResourceCPU: "1.5", v1.ResourceMemory: "1.5Gi"}).Obj()

	// Register the blocking plugin