Support RejectWaitingPod (#98)

Gekko0114 · web-flow · commit b3d74fa17c06 · 2024-01-29T04:38:47.000-08:00
* rejectwaitingpod

* add_testcase

* move eventrecorder to handle/eventrecorder

* fix test
diff --git a/Makefile b/Makefile
@@ -9,7 +9,8 @@ examples/advanced/main.wasm: examples/advanced/main.go
 	@(cd $(@D); tinygo build -o main.wasm -scheduler=none --no-debug -target=wasi .)
 
 .PHONY: build-tinygo
-build-tinygo: examples/nodenumber/main.wasm examples/advanced/main.wasm guest/testdata/cyclestate/main.wasm guest/testdata/filter/main.wasm guest/testdata/score/main.wasm guest/testdata/bind/main.wasm guest/testdata/reserve/main.wasm
+build-tinygo: examples/nodenumber/main.wasm examples/advanced/main.wasm guest/testdata/cyclestate/main.wasm guest/testdata/filter/main.wasm guest/testdata/score/main.wasm \
+			  guest/testdata/bind/main.wasm guest/testdata/reserve/main.wasm guest/testdata/handle/main.wasm
 
 %/main-debug.wasm: %/main.go
 	@(cd $(@D); tinygo build -o main-debug.wasm -gc=custom -tags=custommalloc -scheduler=none -target=wasi .)
diff --git a/examples/advanced/main.go b/examples/advanced/main.go
@@ -28,7 +28,7 @@ import (
 	"sigs.k8s.io/kube-scheduler-wasm-extension/examples/advanced/plugin"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/config"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/enqueue"
-	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/eventrecorder"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/eventrecorder"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/klog"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/prescore"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/score"
diff --git a/examples/advanced/main.wasm b/examples/advanced/main.wasm
diff --git a/examples/advanced/plugin/plugin.go b/examples/advanced/plugin/plugin.go
@@ -33,7 +33,7 @@ import (
 
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api/proto"
-	eventrecorder "sigs.k8s.io/kube-scheduler-wasm-extension/guest/eventrecorder/api"
+	eventrecorder "sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/eventrecorder/api"
 	klog "sigs.k8s.io/kube-scheduler-wasm-extension/guest/klog/api"
 )
 
diff --git a/examples/advanced/plugin/plugin_test.go b/examples/advanced/plugin/plugin_test.go
@@ -21,7 +21,7 @@ import (
 
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api/proto"
-	eventrecorderapi "sigs.k8s.io/kube-scheduler-wasm-extension/guest/eventrecorder/api"
+	eventrecorderapi "sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/eventrecorder/api"
 	klogapi "sigs.k8s.io/kube-scheduler-wasm-extension/guest/klog/api"
 	protoapi "sigs.k8s.io/kube-scheduler-wasm-extension/kubernetes/proto/api"
 )
diff --git a/examples/nodenumber/main.go b/examples/nodenumber/main.go
@@ -25,7 +25,7 @@ import (
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api/proto"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/config"
-	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/eventrecorder"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/eventrecorder"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/klog"
 	klogapi "sigs.k8s.io/kube-scheduler-wasm-extension/guest/klog/api"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/plugin"
diff --git a/examples/nodenumber/main.wasm b/examples/nodenumber/main.wasm
diff --git a/guest/handle/eventrecorder/api/eventrecorder.go b/guest/handle/eventrecorder/api/eventrecorder.go
diff --git a/guest/handle/eventrecorder/api/eventrecorder_test.go b/guest/handle/eventrecorder/api/eventrecorder_test.go
@@ -18,7 +18,7 @@ package api_test
 
 import (
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api/proto"
-	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/eventrecorder/api"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/eventrecorder/api"
 )
 
 var (
diff --git a/guest/handle/eventrecorder/eventrecorder.go b/guest/handle/eventrecorder/eventrecorder.go
@@ -21,8 +21,8 @@ import (
 	"runtime"
 
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api/proto"
-	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/eventrecorder/api"
-	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/eventrecorder/internal"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/eventrecorder/api"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/eventrecorder/internal"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/mem"
 )
 
diff --git a/guest/handle/eventrecorder/eventrecorder_test.go b/guest/handle/eventrecorder/eventrecorder_test.go
@@ -18,7 +18,7 @@ package eventrecorder_test
 
 import (
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api/proto"
-	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/eventrecorder"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/eventrecorder"
 )
 
 var pod proto.Pod
diff --git a/guest/handle/eventrecorder/imports.go b/guest/handle/eventrecorder/imports.go
diff --git a/guest/handle/eventrecorder/imports_stub.go b/guest/handle/eventrecorder/imports_stub.go
diff --git a/guest/handle/eventrecorder/internal/eventrecorder.go b/guest/handle/eventrecorder/internal/eventrecorder.go
@@ -19,7 +19,7 @@ package internal
 
 import (
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api/proto"
-	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/eventrecorder/api"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/eventrecorder/api"
 )
 
 type EventRecorder struct {
diff --git a/guest/handle/eventrecorder/internal/eventrecorder_test.go b/guest/handle/eventrecorder/internal/eventrecorder_test.go
@@ -19,7 +19,7 @@ package internal
 import (
 	"testing"
 
-	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/eventrecorder/api"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/eventrecorder/api"
 )
 
 func TestEventf(t *testing.T) {
diff --git a/guest/handle/handle.go b/guest/handle/handle.go
@@ -0,0 +1,36 @@
+/*
+   Copyright 2023 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 prescore exports an api.PreScorePlugin to the host. Only import this
+// package when setting Plugin, as doing otherwise will cause overhead.
+package handle
+
+import (
+	"runtime"
+
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/mem"
+)
+
+func RejectWaitingPod(uid string) bool {
+	ptr, size := mem.StringToPtr(uid)
+
+	// Wrap to avoid TinyGo 0.28: cannot use an exported function as value
+	wasmBool := mem.SendAndGetUint64(ptr, size, func(input_ptr, input_size, ptr uint32, limit mem.BufLimit) {
+		rejectWaitingPod(input_ptr, input_size, ptr, limit)
+	})
+	runtime.KeepAlive(uid)
+	return wasmBool == 1
+}
diff --git a/guest/handle/imports.go b/guest/handle/imports.go
@@ -0,0 +1,24 @@
+//go:build tinygo.wasm
+
+/*
+   Copyright 2023 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 handle
+
+import "sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/mem"
+
+//go:wasmimport k8s.io/scheduler handle.reject_waiting_pod
+func rejectWaitingPod(input_ptr, input_size, ptr uint32, limit mem.BufLimit)
diff --git a/guest/handle/imports_stub.go b/guest/handle/imports_stub.go
@@ -0,0 +1,24 @@
+//go:build !tinygo.wasm
+
+/*
+   Copyright 2023 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 handle
+
+import "sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/mem"
+
+// rejectWaitingPod is stubbed for compilation outside TinyGo.
+func rejectWaitingPod(uint32, uint32, uint32, mem.BufLimit) {}
diff --git a/guest/internal/mem/mem.go b/guest/internal/mem/mem.go
@@ -16,7 +16,10 @@
 
 package mem
 
-import "unsafe"
+import (
+	"encoding/binary"
+	"unsafe"
+)
 
 // BufLimit is the possibly zero maximum length of a result value to write in
 // bytes. If the actual value is larger than this, nothing is written to
@@ -95,3 +98,8 @@ func GetString(fn func(ptr uint32, limit BufLimit) (len uint32)) string {
 	_ = fn(uint32(uintptr(ptr)), size)
 	return unsafe.String((*byte)(ptr), size /* unsafe.IntegerType */)
 }
+
+func SendAndGetUint64(input_ptr uint32, input_size uint32, fn func(input_ptr, input_size, ptr uint32, limit BufLimit)) uint64 {
+	fn(input_ptr, input_size, uint32(readBufPtr), readBufLimit)
+	return binary.LittleEndian.Uint64(readBuf)
+}
diff --git a/guest/testdata/handle/main.go b/guest/testdata/handle/main.go
@@ -0,0 +1,84 @@
+/*
+   Copyright 2023 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 main
+
+import (
+	"os"
+
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api/proto"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/filter"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/prefilter"
+)
+
+type extensionPoints interface {
+	api.PreFilterPlugin
+	api.FilterPlugin
+}
+
+// This code is for checking the method of handle.
+func main() {
+	// Multiple tests are here to reduce re-compilation time and size checked
+	// into git.
+	var plugin extensionPoints = noopPlugin{}
+	if len(os.Args) == 2 {
+		switch os.Args[1] {
+		case "rejectWaitingPod":
+			plugin = pluginForReject{}
+		}
+	}
+	prefilter.SetPlugin(plugin)
+	filter.SetPlugin(plugin)
+}
+
+// noopPlugin doesn't do anything, except evaluate each parameter.
+type noopPlugin struct{}
+
+func (noopPlugin) PreFilter(state api.CycleState, pod proto.Pod) (nodeNames []string, status *api.Status) {
+	_, _ = state.Read("ok")
+	_ = pod.Spec()
+	return
+}
+
+func (noopPlugin) Filter(state api.CycleState, pod proto.Pod, nodeInfo api.NodeInfo) (status *api.Status) {
+	_, _ = state.Read("ok")
+	_ = pod.Spec()
+	_ = nodeInfo.Node().Spec() // trigger lazy loading
+	return
+}
+
+// pluginForReject checks the function of RejectWaitingPod
+type pluginForReject struct{ noopPlugin }
+
+func (pluginForReject) Filter(_ api.CycleState, pod proto.Pod, nodeInfo api.NodeInfo) *api.Status {
+	// Call RejectWaitingPod first
+	IsRejected := handle.RejectWaitingPod(pod.GetUid())
+
+	if IsRejected {
+		// This is being skipped, note the reason.
+		return &api.Status{
+			Code:   api.StatusCodeSkip,
+			Reason: "UID is " + pod.GetUid(),
+		}
+	}
+
+	// Otherwise, this is success.
+	return &api.Status{
+		Code: api.StatusCodeSuccess,
+	}
+}
diff --git a/guest/testdata/handle/main.wasm b/guest/testdata/handle/main.wasm
diff --git a/scheduler/plugin/host.go b/scheduler/plugin/host.go
@@ -50,6 +50,7 @@ const (
 	k8sSchedulerResultStatusReason        = "result.status_reason"
 	k8sSchedulerResultNormalizedScoreList = "result.normalized_score_list"
 	k8sSchedulerHandleEventRecorderEventf = "handle.eventrecorder.eventf"
+	k8sSchedulerHandleRejectWaitingPod    = "handle.reject_waiting_pod"
 )
 
 func instantiateHostApi(ctx context.Context, runtime wazero.Runtime) (wazeroapi.Module, error) {
@@ -114,6 +115,9 @@ func instantiateHostScheduler(ctx context.Context, runtime wazero.Runtime, guest
 		NewFunctionBuilder().
 		WithGoModuleFunction(wazeroapi.GoModuleFunc(host.k8sHandleEventRecorderEventfFn), []wazeroapi.ValueType{i32, i32}, []wazeroapi.ValueType{}).
 		WithParameterNames("buf", "buf_len").Export(k8sSchedulerHandleEventRecorderEventf).
+		NewFunctionBuilder().
+		WithGoModuleFunction(wazeroapi.GoModuleFunc(host.k8sHandleRejectWaitingPodFn), []wazeroapi.ValueType{i32, i32, i32, i32}, []wazeroapi.ValueType{}).
+		WithParameterNames("buf", "buf_len").Export(k8sSchedulerHandleRejectWaitingPod).
 		Instantiate(ctx)
 }
 
@@ -493,3 +497,23 @@ func convertToObjectReference(objRef *ObjectReference) *v1.ObjectReference {
 		ResourceVersion: objRef.ResourceVersion,
 	}
 }
+
+// k8sHandleRejectWaitingPodFn is a function used by the wasm guest to call RejectWaitingPod
+func (h host) k8sHandleRejectWaitingPodFn(ctx context.Context, mod wazeroapi.Module, stack []uint64) {
+	iBuf := uint32(stack[0])
+	iBufLen := uint32(stack[1])
+	oBuf := uint32(stack[2])
+	oBufLimit := uint32(stack[3])
+
+	b, ok := mod.Memory().Read(iBuf, iBufLen)
+	if !ok {
+		panic("out of memory reading rejectWaitingPod")
+	}
+	uid := types.UID(b)
+	IsRejected := h.handle.RejectWaitingPod(uid)
+	wasmBool := uint64(0)
+	if IsRejected {
+		wasmBool = uint64(1)
+	}
+	writeUint64(mod.Memory(), wasmBool, oBuf, oBufLimit)
+}
diff --git a/scheduler/plugin/host_test.go b/scheduler/plugin/host_test.go
@@ -23,6 +23,7 @@ import (
 	"testing"
 
 	"github.com/tetratelabs/wazero/experimental/wazerotest"
+	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/klog/v2"
 	k8stest "k8s.io/klog/v2/test"
 
@@ -130,3 +131,31 @@ func Test_k8sHandleEventRecorderEventFn(t *testing.T) {
 		t.Fatalf("unexpected event: %v != %v", want, have)
 	}
 }
+
+func Test_k8sHandleRejectWaitingPodFn(t *testing.T) {
+	recorder := &test.FakeRecorder{EventMsg: ""}
+	handle := &test.FakeHandle{Recorder: recorder}
+	h := host{handle: handle}
+
+	// Create a fake wasm module, which has data the guest should write.
+	mem := wazerotest.NewMemory(wazerotest.PageSize)
+	mod := wazerotest.NewModule(mem)
+	uid := types.UID("c6feae3a-7082-42a5-a5ec-6ae2e1603727")
+	copy(mem.Bytes, uid)
+
+	// Invoke the host function in the same way the guest would have.
+	h.k8sHandleRejectWaitingPodFn(context.Background(), mod, []uint64{
+		0,
+		uint64(len(uid)),
+		0, // Ideally we should define some value, but we don't define it for now.
+		0, // Ideally we should define some value, but we don't define it for now.
+	})
+
+	// Checking the value stored on handle
+	have := handle.RejectWaitingPodValue
+	want := uid
+
+	if want != have {
+		t.Fatalf("unexpected uid: %v != %v", want, have)
+	}
+}
diff --git a/scheduler/plugin/mem.go b/scheduler/plugin/mem.go
@@ -73,3 +73,7 @@ func writeStringIfUnderLimit(mem wazeroapi.Memory, v string, buf uint32, bufLimi
 	mem.WriteString(buf, v)
 	return vLen
 }
+
+func writeUint64(mem wazeroapi.Memory, v uint64, buf uint32, bufLimit bufLimit) {
+	mem.WriteUint64Le(buf, v)
+}
diff --git a/scheduler/plugin/plugin_test.go b/scheduler/plugin/plugin_test.go
diff --git a/scheduler/test/testdata.go b/scheduler/test/testdata.go
diff --git a/scheduler/test/util.go b/scheduler/test/util.go

Original file line number	Diff line number	Diff line change
`@@ -33,7 +33,7 @@ import (`
`33`	`33`
`34`	`34`	`"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api"`
`35`	`35`	`"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api/proto"`
`36`		`- eventrecorder "sigs.k8s.io/kube-scheduler-wasm-extension/guest/eventrecorder/api"`
	`36`	`+ eventrecorder "sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/eventrecorder/api"`
`37`	`37`	`klog "sigs.k8s.io/kube-scheduler-wasm-extension/guest/klog/api"`
`38`	`38`	`)`
`39`	`39`
Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,7 @@ import (`
`21`	`21`
`22`	`22`	`"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api"`
`23`	`23`	`"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api/proto"`
`24`		`- eventrecorderapi "sigs.k8s.io/kube-scheduler-wasm-extension/guest/eventrecorder/api"`
	`24`	`+ eventrecorderapi "sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/eventrecorder/api"`
`25`	`25`	`klogapi "sigs.k8s.io/kube-scheduler-wasm-extension/guest/klog/api"`
`26`	`26`	`protoapi "sigs.k8s.io/kube-scheduler-wasm-extension/kubernetes/proto/api"`
`27`	`27`	`)`
Original file line number	Diff line number	Diff line change
`@@ -18,7 +18,7 @@ package api_test`
`18`	`18`
`19`	`19`	`import (`
`20`	`20`	`"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api/proto"`
`21`		`- "sigs.k8s.io/kube-scheduler-wasm-extension/guest/eventrecorder/api"`
	`21`	`+ "sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/eventrecorder/api"`
`22`	`22`	`)`
`23`	`23`
`24`	`24`	`var (`
Original file line number	Diff line number	Diff line change
`@@ -21,8 +21,8 @@ import (`
`21`	`21`	`"runtime"`
`22`	`22`
`23`	`23`	`"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api/proto"`
`24`		`- "sigs.k8s.io/kube-scheduler-wasm-extension/guest/eventrecorder/api"`
`25`		`- "sigs.k8s.io/kube-scheduler-wasm-extension/guest/eventrecorder/internal"`
	`24`	`+ "sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/eventrecorder/api"`
	`25`	`+ "sigs.k8s.io/kube-scheduler-wasm-extension/guest/handle/eventrecorder/internal"`
`26`	`26`	`"sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/mem"`
`27`	`27`	`)`
`28`	`28`