Support permit plugin

Author: kazukousen

Makefile

@@ -10,7 +10,7 @@ examples/advanced/main.wasm: examples/advanced/main.go
 
 .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 guest/testdata/handle/main.wasm
+			  guest/testdata/bind/main.wasm guest/testdata/reserve/main.wasm guest/testdata/handle/main.wasm guest/testdata/permit/main.wasm
 
 %/main-debug.wasm: %/main.go
 	@(cd $(@D); tinygo build -o main-debug.wasm -gc=custom -tags=custommalloc -scheduler=none -target=wasi .)

guest/api/types.go

@@ -16,7 +16,11 @@
 
 package api
 
-import "sigs.k8s.io/kube-scheduler-wasm-extension/guest/api/proto"
+import (
+	"time"
+
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api/proto"
+)
 
 // CycleState is a WebAssembly implementation of framework.CycleState.
 //
@@ -122,6 +126,13 @@ type ReservePlugin interface {
 	Unreserve(state CycleState, p proto.Pod, nodeName string)
 }
 
+// PermitPlugin is a WebAssembly implementation of framework.PermitPlugin.
+type PermitPlugin interface {
+	Plugin
+
+	Permit(state CycleState, p proto.Pod, nodeName string) (*Status, time.Duration)
+}
+
 // PreBindPlugin is a WebAssembly implementation of framework.PreBindPlugin.
 type PreBindPlugin interface {
 	Plugin

guest/internal/mem/mem.go

@@ -53,6 +53,15 @@ func BytesToPtr(b []byte) (uint32, uint32) {
 	return uint32(uintptr(ptr)), uint32(len(b))
 }
 
+// Int64ToPtr returns a pointer for the given int64 number in a way compatible
+// with WebAssembly numeric types.
+// The returned pointer aliases the number hence it must be kept alive until ptr
+// is no longer needed.
+func Int64ToPtr(n int64) uint32 {
+	ptr := unsafe.Pointer(&n)
+	return uint32(uintptr(ptr))
+}
+
 // Update is for decoding values from memory. The updater doesn't keep a
 // reference to the underlying bytes, so we don't need to copy them.
 func Update(

guest/permit/imports.go

@@ -0,0 +1,6 @@
+//go:build tinygo.wasm
+
+package permit
+
+//go:wasmimport k8s.io/scheduler result.timeout
+func setTimeoutResult(uint32)

guest/permit/imports_stub.go

@@ -0,0 +1,6 @@
+//go:build !tinygo.wasm
+
+package permit
+
+// setTimeoutResult is stubbed for compilation outside TinyGo.
+func setTimeoutResult(uint32) {}

guest/permit/permit.go

@@ -0,0 +1,78 @@
+/*
+   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 permit exports an api.PermitPlugin to the host.
+package permit
+
+import (
+	"runtime"
+
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/cyclestate"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/imports"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/mem"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/plugin"
+)
+
+// permit is the current plugin assigned with SetPlugin.
+var permit api.PermitPlugin
+
+// SetPlugin should be called in `main` to assign an api.PermitPlugin
+// instance.
+//
+// For example:
+//
+//	func main() {
+//		plugin := permitPlugin{}
+//		permit.SetPlugin(plugin)
+//	}
+//
+//	type permitPlugin struct{}
+//
+//	func (permitPlugin) Permit(state api.CycleState, p proto.Pod, nodeName string) (status *api.Status, timeout time.Duration)
+//		// Write state you need on Permit
+//	}
+func SetPlugin(permitPlugin api.PermitPlugin) {
+	if permitPlugin == nil {
+		panic("nil permitPlugin")
+	}
+	permit = permitPlugin
+	plugin.MustSet(permit)
+}
+
+// prevent unused lint errors (lint is run with normal go).
+var _ func() uint32 = _permit
+
+// _permit is only exported to the host.
+//
+//export permit
+func _permit() uint32 {
+	if permit == nil { // Then, the user didn't define one.
+		// Unlike most plugins we always export permit so that we can reset
+		// the cycle state: return success to avoid no-op overhead.
+		return 0
+	}
+
+	pod := cyclestate.Pod
+	nodeName := imports.NodeName()
+	status, timeout := permit.Permit(cyclestate.Values, pod, nodeName)
+
+	ptr := mem.Int64ToPtr(int64(timeout))
+	setTimeoutResult(ptr)
+	runtime.KeepAlive(timeout) // untir ptr is no longer needed.
+
+	return imports.StatusToCode(status)
+}

guest/plugin/plugin.go

@@ -22,6 +22,7 @@ import (
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/enqueue"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/filter"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/prefilter"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/permit"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/postbind"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/postfilter"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/prebind"
@@ -72,6 +73,9 @@ func Set(plugin api.Plugin) {
 	if plugin, ok := plugin.(api.ReservePlugin); ok {
 		reserve.SetPlugin(plugin)
 	}
+	if plugin, ok := plugin.(api.PermitPlugin); ok {
+		permit.SetPlugin(plugin)
+	}
 	if plugin, ok := plugin.(api.PreBindPlugin); ok {
 		prebind.SetPlugin(plugin)
 	}

guest/testdata/cyclestate/main.go

@@ -20,13 +20,15 @@ package main
 
 import (
 	"os"
+	"time"
 	"unsafe"
 
 	"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/bind"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/enqueue"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/filter"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/permit"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/postbind"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/postfilter"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/prebind"
@@ -67,6 +69,7 @@ func main() {
 	score.SetPlugin(plugin)
 	scoreextensions.SetPlugin(plugin)
 	reserve.SetPlugin(plugin)
+	permit.SetPlugin(plugin)
 	prebind.SetPlugin(plugin)
 	bind.SetPlugin(plugin)
 	postbind.SetPlugin(plugin)
@@ -187,6 +190,11 @@ func (statePlugin) Unreserve(state api.CycleState, pod proto.Pod, nodeName strin
 	mustFilterState(state)
 }
 
+func (statePlugin) Permit(state api.CycleState, pod proto.Pod, nodeName string) (status *api.Status, timeout time.Duration) {
+	mustFilterState(state)
+	return
+}
+
 func (statePlugin) PreBind(state api.CycleState, pod proto.Pod, _ string) (status *api.Status) {
 	if unsafe.Pointer(pod.Spec()) != unsafe.Pointer(podSpec) {
 		panic("didn't cache pod from pre-filter")

guest/testdata/cyclestate/main.wasm

@@ -0,0 +1,47 @@
+/*
+   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 (
+	"time"
+
+	"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/permit"
+)
+
+type extensionPoints interface {
+	api.PermitPlugin
+}
+
+func main() {
+	var plugin extensionPoints = permitPlugin{}
+	permit.SetPlugin(plugin)
+}
+
+type permitPlugin struct{}
+
+func (permitPlugin) Permit(state api.CycleState, pod proto.Pod, nodeName string) (*api.Status, time.Duration) {
+	status, timeout := api.StatusCodeSuccess, time.Duration(0)
+	if nodeName == "bad" {
+		status = api.StatusCodeError
+	} else if nodeName == "wait" {
+		status = api.StatusCodeWait
+		timeout = 10 * time.Second
+	}
+	return &api.Status{Code: status, Reason: "name is " + nodeName}, timeout
+}