Support PrefilterExtensions (#92)

* Support PrefilterExtensions

* Address the feedbacks

* Fix argument for plugin test

* Fix the variable name to lowercase

* Fix type errors

* Update test

* Update plugin test

* Address feedback and update testcode

* Restore unneeded changes of wasm

* Fix message

* Fix wrong URL

* Add Plugin

* Remove redundant print debug

* Add wasm file for updated go file

* Fix typos

* Remove redundant URL definition

* Add podInfoToAdd & nodeInfo to stack

* Fix `params` to pass Node and Pod

* Address the feedback

* Add `PodInfo{}` instead of cyclestate.Pod

* Remove unused type

* Update the type between host & guest

* Update the definition of guest

* Address feedback

* Fix type

* Fix the definition of PodInfo

* Fix type of testdata

* Pass PodInfo to function

* Remove unnecessary comment

* Fix e2e test for PrefilterExtensions

* Rollback unncecessary change of wasm

* Fix comment typo and oddity

* Remove redundant definition in wasm

* Merge redundant definition

* Add necessary information to testcase output

* Fix function for host api

* Fix the order of plugin definition

* Implement Pod for PodInfoToAdd

* Update the naming of the pod for add/remove

* Add comment for k8sApiPodInfoFn

* Update error message

* Tweak return string of `GetKind()`

* Upddate testdata to give concrete information

* Rename to `targetPod`

* fix: rollack wasm

Author: chansuke

guest/api/types.go

@@ -59,6 +59,14 @@ type PreFilterPlugin interface {
 	PreFilter(state CycleState, pod proto.Pod) (nodeNames []string, status *Status)
 }
 
+// PreFilterExtensions is a WebAssembly implementation of framework.PrefFilterExtensions.
+type PreFilterExtensions interface {
+	Plugin
+
+	AddPod(state CycleState, podToSchedule proto.Pod, podInfoToAdd PodInfo, nodeInfo NodeInfo) *Status
+	RemovePod(state CycleState, podToSchedule proto.Pod, podInfoToRemove PodInfo, nodeInfo NodeInfo) *Status
+}
+
 // FilterPlugin is a WebAssembly implementation of framework.FilterPlugin.
 //
 // Note: The pod and nodeInfo parameters are lazy to avoid unmarshal overhead
@@ -160,6 +168,13 @@ type NodeInfo interface {
 	Node() proto.Node
 }
 
+type PodInfo interface {
+	// Metadata is a convenience that triggers Get.
+	proto.Metadata
+
+	Pod() proto.Pod
+}
+
 // NodeToStatus contains which Node got which status during the scheduling cycle.
 type NodeToStatus interface {
 	// NodeToStatus returns a map

guest/filter/filter.go

@@ -67,43 +67,43 @@ func _filter() uint32 { //nolint
 		return 0
 	}
 
-	s := filter.Filter(cyclestate.Values, cyclestate.Pod, &nodeInfo{})
+	s := filter.Filter(cyclestate.Values, cyclestate.Pod, &NodeInfo{})
 
 	return imports.StatusToCode(s)
 }
 
-var _ api.NodeInfo = (*nodeInfo)(nil)
+var _ api.NodeInfo = (*NodeInfo)(nil)
 
 // nodeInfo is lazy so that a plugin which doesn't read fields avoids a
 // relatively expensive unmarshal penalty.
 //
 // Note: Unlike proto.Pod, this is not special cased for the scheduling cycle.
-type nodeInfo struct {
+type NodeInfo struct {
 	node proto.Node
 }
 
-func (n *nodeInfo) GetUid() string {
+func (n *NodeInfo) GetUid() string {
 	return n.lazyNode().GetUid()
 }
 
-func (n *nodeInfo) GetName() string {
+func (n *NodeInfo) GetName() string {
 	return n.lazyNode().GetName()
 }
 
-func (n *nodeInfo) GetNamespace() string {
+func (n *NodeInfo) GetNamespace() string {
 	return n.lazyNode().GetNamespace()
 }
 
-func (n *nodeInfo) GetResourceVersion() string {
+func (n *NodeInfo) GetResourceVersion() string {
 	return n.lazyNode().GetResourceVersion()
 }
 
-func (n *nodeInfo) Node() proto.Node {
+func (n *NodeInfo) Node() proto.Node {
 	return n.lazyNode()
 }
 
 // lazyNode lazy initializes node from imports.Node.
-func (n *nodeInfo) lazyNode() proto.Node {
+func (n *NodeInfo) lazyNode() proto.Node {
 	if node := n.node; node != nil {
 		return node
 	}
@@ -115,3 +115,57 @@ func (n *nodeInfo) lazyNode() proto.Node {
 	n.node = &internalproto.Node{Msg: &msg}
 	return n.node
 }
+
+type PodInfo struct {
+	pod proto.Pod
+}
+
+func (p *PodInfo) GetApiVersion() string {
+	return p.lazyPod().GetApiVersion()
+}
+
+func (p *PodInfo) GetKind() string {
+	return p.lazyPod().GetKind()
+}
+
+func (p *PodInfo) GetName() string {
+	return p.lazyPod().GetName()
+}
+
+func (p *PodInfo) GetNamespace() string {
+	return p.lazyPod().GetNamespace()
+}
+
+func (p *PodInfo) GetResourceVersion() string {
+	return p.lazyPod().GetNamespace()
+}
+
+func (p *PodInfo) GetUid() string {
+	return p.lazyPod().GetUid()
+}
+
+func (p *PodInfo) Pod() proto.Pod {
+	return p.lazyPod()
+}
+
+func (p *PodInfo) Spec() *protoapi.PodSpec {
+	return p.lazyPod().Spec()
+}
+
+func (p *PodInfo) Status() *protoapi.PodStatus {
+	return p.lazyPod().Status()
+}
+
+// lazyPod lazy initializes pod from imports.Pod.
+func (p *PodInfo) lazyPod() proto.Pod {
+	if pod := p.pod; pod != nil {
+		return pod
+	}
+
+	var msg protoapi.Pod
+	if err := imports.Pod(msg.UnmarshalVT); err != nil {
+		panic(err.Error())
+	}
+	p.pod = &internalproto.Pod{Msg: &msg}
+	return p.pod
+}

guest/internal/proto/proto.go

@@ -91,3 +91,41 @@ func (o *Node) Spec() *protoapi.NodeSpec {
 func (o *Node) Status() *protoapi.NodeStatus {
 	return o.Msg.Status
 }
+
+var _ proto.Pod = (*Pod)(nil)
+
+type Pod struct {
+	Msg *protoapi.Pod
+}
+
+func (o *Pod) GetName() string {
+	return GetName(o.Msg)
+}
+
+func (o *Pod) GetNamespace() string {
+	return GetNamespace(o.Msg)
+}
+
+func (o *Pod) GetUid() string {
+	return GetUid(o.Msg)
+}
+
+func (o *Pod) GetResourceVersion() string {
+	return GetResourceVersion(o.Msg)
+}
+
+func (o *Pod) GetKind() string {
+	return "Pod"
+}
+
+func (o *Pod) GetApiVersion() string {
+	return "v1"
+}
+
+func (o *Pod) Spec() *protoapi.PodSpec {
+	return o.Msg.Spec
+}
+
+func (o *Pod) Status() *protoapi.PodStatus {
+	return o.Msg.Status
+}

guest/plugin/plugin.go

@@ -26,6 +26,7 @@ import (
 	"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"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/prefilterextensions"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/prescore"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/reserve"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/score"
@@ -55,6 +56,9 @@ func Set(plugin api.Plugin) {
 	if plugin, ok := plugin.(api.PreFilterPlugin); ok {
 		prefilter.SetPlugin(plugin)
 	}
+	if plugin, ok := plugin.(api.PreFilterExtensions); ok {
+		prefilterextensions.SetPlugin(plugin)
+	}
 	if plugin, ok := plugin.(api.FilterPlugin); ok {
 		filter.SetPlugin(plugin)
 	}

guest/prefilterextensions/prefilterextensions.go

@@ -0,0 +1,76 @@
+/*
+   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 prefilterextensions exports an api.PreFilterExtensions to the host. Only import this
+// package when setting Plugin, as doing otherwise will cause overhead.
+
+package prefilterextensions
+
+import (
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/filter"
+	"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/plugin"
+)
+
+// prefilterextensions is the current plugin assigned with SetPlugin.
+var prefilterextensions api.PreFilterExtensions
+
+// SetPlugin is exposed to prevent package cycles.
+func SetPlugin(preFilterExtensions api.PreFilterExtensions) {
+	if preFilterExtensions == nil {
+		panic("nil preFilterExtensions")
+	}
+	prefilterextensions = preFilterExtensions
+	plugin.MustSet(prefilterextensions)
+}
+
+// prevent unused lint errors (lint is run with normal go).
+var (
+	_ func() uint32 = _addpod
+	_ func() uint32 = _removepod
+)
+
+// _addPod is only exported to the host.
+//
+//export addpod
+func _addpod() uint32 { //nolint
+	if prefilterextensions == nil { // Then, the user didn't define one.
+		// Unlike most plugins we always export reserve so that we can reset
+		// the cycle state: return success to avoid no-op overhead.
+		return 0
+	}
+
+	status := prefilterextensions.AddPod(cyclestate.Values, cyclestate.Pod, &filter.PodInfo{}, &filter.NodeInfo{})
+
+	return imports.StatusToCode(status)
+}
+
+// _removePod is only exported to the host.
+//
+//export removepod
+func _removepod() uint32 { //nolint
+	if prefilterextensions == nil { // Then, the user didn't define one.
+		// Unlike most plugins we always export unreserve so that we can reset
+		// the cycle state: return success to avoid no-op overhead.
+		return 0
+	}
+
+	status := prefilterextensions.RemovePod(cyclestate.Values, cyclestate.Pod, &filter.PodInfo{}, &filter.NodeInfo{})
+
+	return imports.StatusToCode(status)
+}

guest/testdata/bind/main.wasm

@@ -32,6 +32,7 @@ import (
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/postfilter"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/prebind"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/prefilter"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/prefilterextensions"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/prescore"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/reserve"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/score"
@@ -62,6 +63,7 @@ func main() {
 	plugin := statePlugin{}
 	enqueue.SetPlugin(plugin)
 	prefilter.SetPlugin(plugin)
+	prefilterextensions.SetPlugin(plugin)
 	filter.SetPlugin(plugin)
 	postfilter.SetPlugin(plugin)
 	prescore.SetPlugin(plugin)
@@ -229,6 +231,14 @@ func (statePlugin) PostBind(state api.CycleState, pod proto.Pod, _ string) {
 	}
 }
 
+func (statePlugin) AddPod(state api.CycleState, pod proto.Pod, podInfoToAdd api.PodInfo, nodeInfo api.NodeInfo) (status *api.Status) {
+	return
+}
+
+func (statePlugin) RemovePod(state api.CycleState, pod proto.Pod, podInfoToAdd api.PodInfo, nodeInfo api.NodeInfo) (status *api.Status) {
+	return
+}
+
 // mustNotScoreState ensures that score state, written after filter, cannot
 // be read by extension points before it.
 //

guest/testdata/cyclestate/main.go

@@ -24,12 +24,14 @@ import (
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/filter"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/postfilter"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/prefilter"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/prefilterextensions"
 )
 
 type extensionPoints interface {
 	api.PreFilterPlugin
 	api.FilterPlugin
 	api.PostFilterPlugin
+	api.PreFilterExtensions
 }
 
 func main() {
@@ -44,11 +46,14 @@ func main() {
 			plugin = preFilterPlugin{}
 		case "postFilter":
 			plugin = postFilterPlugin{}
+		case "preFilterExtensions":
+			plugin = preFilterExtensions{}
 		}
 	}
 	prefilter.SetPlugin(plugin)
 	filter.SetPlugin(plugin)
 	postfilter.SetPlugin(plugin)
+	prefilterextensions.SetPlugin(plugin)
 }
 
 // noopPlugin doesn't do anything, except evaluate each parameter.
@@ -74,6 +79,20 @@ func (noopPlugin) PostFilter(state api.CycleState, pod proto.Pod, nodeMap api.No
 	return
 }
 
+func (noopPlugin) AddPod(state api.CycleState, pod proto.Pod, podInfoToAdd api.PodInfo, nodeInfo api.NodeInfo) (status *api.Status) {
+	_, _ = state.Read("ok")
+	_ = pod.Spec()
+	_ = nodeInfo.Node().Spec() // trigger lazy loading
+	return
+}
+
+func (noopPlugin) RemovePod(state api.CycleState, pod proto.Pod, podInfoToRemove api.PodInfo, nodeInfo api.NodeInfo) (status *api.Status) {
+	_, _ = state.Read("ok")
+	_ = pod.Spec()
+	_ = nodeInfo.Node().Spec() // trigger lazy loading
+	return
+}
+
 // preFilterPlugin schedules a node if its name equals its pod spec.
 type preFilterPlugin struct{ noopPlugin }
 
@@ -133,3 +152,25 @@ func (postFilterPlugin) PostFilter(_ api.CycleState, pod proto.Pod, nodeMap api.
 		Reason: podSpecNodeName + " is unschedulable",
 	}
 }
+
+type preFilterExtensions struct{ noopPlugin }
+
+func (preFilterExtensions) AddPod(state api.CycleState, pod proto.Pod, podInfoToAdd api.PodInfo, nodeInfo api.NodeInfo) *api.Status {
+	nodeName := nodeInfo.Node().GetName()
+	state.Write(nodeName+pod.Spec().GetNodeName(), "ok")
+	status := 0
+	if nodeName == "bad" {
+		status = 1
+	}
+	return &api.Status{Code: api.StatusCode(status), Reason: "Node name is " + nodeName + " " + "and PodInfo name is " + podInfoToAdd.Pod().GetName()}
+}
+
+func (preFilterExtensions) RemovePod(state api.CycleState, pod proto.Pod, podInfoToRemove api.PodInfo, nodeInfo api.NodeInfo) *api.Status {
+	nodeName := nodeInfo.Node().GetName()
+	state.Write(nodeName+pod.Spec().GetNodeName(), "ok")
+	status := 0
+	if nodeName == "bad" {
+		status = 1
+	}
+	return &api.Status{Code: api.StatusCode(status), Reason: "Node name is " + nodeName + " " + "and PodInfo name is " + podInfoToRemove.Pod().GetName()}
+}