Support GetWaitingPod

guest/api/types.go

@@ -205,3 +205,10 @@ type NodeScore interface {
 	// which is keyed by the node name and valued by the score.
 	Map() map[string]int
 }
+
+// WaitingPod represents a pod currently waiting in the permit phase.
+type WaitingPod interface {
+	GetPod() proto.Pod
+	// Reject declares the waiting pod unschedulable.
+	Reject(pluginName, msg string)
+}

guest/go.mod

@@ -4,9 +4,9 @@ go 1.22.0
 
 replace sigs.k8s.io/kube-scheduler-wasm-extension/kubernetes/proto => ../kubernetes/proto
 
-require sigs.k8s.io/kube-scheduler-wasm-extension/kubernetes/proto v0.0.0-00010101000000-000000000000
-
 require (
-	github.com/google/go-cmp v0.5.9 // indirect
-	google.golang.org/protobuf v1.30.0 // indirect
+	google.golang.org/protobuf v1.30.0
+	sigs.k8s.io/kube-scheduler-wasm-extension/kubernetes/proto v0.0.0-00010101000000-000000000000
 )
+
+require github.com/google/go-cmp v0.5.9 // indirect

guest/handle/handle.go

@@ -14,14 +14,19 @@
    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 exports an api.RejectWaitingPod and GetWaitingPod to the host.
+// Only import this package when setting Plugin, as doing otherwise will cause overhead.
 package handle
 
 import (
 	"runtime"
 
+	proto "google.golang.org/protobuf/proto"
+
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api"
+	guestapi "sigs.k8s.io/kube-scheduler-wasm-extension/guest/api"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/mem"
+	protoapi "sigs.k8s.io/kube-scheduler-wasm-extension/kubernetes/proto/api"
 )
 
 func RejectWaitingPod(uid string) bool {
@@ -34,3 +39,24 @@ func RejectWaitingPod(uid string) bool {
 	runtime.KeepAlive(uid)
 	return wasmBool == 1
 }
+
+func GetWaitingPod(uid string) guestapi.WaitingPod {
+	ptr, size := mem.StringToPtr(uid)
+
+	var pod protoapi.Pod
+	err := mem.Update(
+		func(outPtr uint32, limit mem.BufLimit) (len uint32) {
+			getWaitingPod(ptr, size, outPtr, limit)
+			return limit
+		},
+		func(data []byte) error {
+			return proto.Unmarshal(data, &pod)
+		},
+	)
+	if err != nil {
+		return nil
+	}
+
+	waitingPod := make([]api.WaitingPod, size)
+	return waitingPod[0]
+}

guest/handle/imports.go

@@ -20,5 +20,8 @@ package handle
 
 import "sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/mem"
 
+//go:wasmimport k8s.io/scheduler handle.get_waiting_pod
+func getWaitingPod(input_ptr, input_size, ptr uint32, limit mem.BufLimit)
+
 //go:wasmimport k8s.io/scheduler handle.reject_waiting_pod
 func rejectWaitingPod(input_ptr, input_size, ptr uint32, limit mem.BufLimit)

guest/handle/imports_stub.go

@@ -20,5 +20,8 @@ package handle
 
 import "sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/mem"
 
+// getWaitingPod is stubbed for compilation outside TinyGo.
+func getWaitingPod(uint32, uint32, uint32, mem.BufLimit) {}
+
 // rejectWaitingPod is stubbed for compilation outside TinyGo.
 func rejectWaitingPod(uint32, uint32, uint32, mem.BufLimit) {}

guest/internal/mem/mem.go

@@ -103,3 +103,24 @@ func SendAndGetUint64(input_ptr uint32, input_size uint32, fn func(input_ptr, in
 	fn(input_ptr, input_size, uint32(readBufPtr), readBufLimit)
 	return binary.LittleEndian.Uint64(readBuf)
 }
+
+func SendAndGetPodBytes(input_ptr uint32, input_size uint32, fn func(input_ptr, input_size, ptr uint32, limit BufLimit)) []byte {
+	readBuf := make([]byte, readBufLimit)
+	readBufPtr := uint32(uintptr(unsafe.Pointer(&readBuf[0])))
+
+	fn(input_ptr, input_size, readBufPtr, readBufLimit)
+
+	var length int
+	for i := 0; i < int(readBufLimit); i++ {
+		if readBuf[i] == 0 {
+			length = i
+			break
+		}
+	}
+
+	if length == 0 {
+		return nil
+	}
+
+	return readBuf[:length]
+}

guest/testdata/handle/main.go

@@ -40,6 +40,8 @@ func main() {
 		switch os.Args[1] {
 		case "rejectWaitingPod":
 			plugin = pluginForReject{}
+		case "getWaitingPod":
+			plugin = pluginForGet{}
 		}
 	}
 	prefilter.SetPlugin(plugin)
@@ -82,3 +84,22 @@ func (pluginForReject) Filter(_ api.CycleState, pod proto.Pod, nodeInfo api.Node
 		Code: api.StatusCodeSuccess,
 	}
 }
+
+// pluginForGet checks the function of GetWaitingPod
+type pluginForGet struct{ noopPlugin }
+
+func (pluginForGet) Filter(_ api.CycleState, pod proto.Pod, nodeInfo api.NodeInfo) *api.Status {
+	// Call GetWaitingPod first
+	waitingPod := handle.GetWaitingPod(pod.GetUid())
+
+	if waitingPod == nil {
+		return &api.Status{
+			Code:   api.StatusCodeError,
+			Reason: "No waiting pod found for UID: " + pod.GetUid(),
+		}
+	}
+
+	return &api.Status{
+		Code: api.StatusCodeSuccess,
+	}
+}

scheduler/plugin/host.go

@@ -54,6 +54,7 @@ const (
 	k8sSchedulerResultNormalizedScoreList = "result.normalized_score_list"
 	k8sSchedulerHandleEventRecorderEventf = "handle.eventrecorder.eventf"
 	k8sSchedulerHandleRejectWaitingPod    = "handle.reject_waiting_pod"
+	k8sSchedulerHandleGetWaitingPod       = "handle.get_waiting_pod"
 )
 
 func instantiateHostApi(ctx context.Context, runtime wazero.Runtime, handle framework.Handle) (wazeroapi.Module, error) {
@@ -131,6 +132,9 @@ func instantiateHostScheduler(ctx context.Context, runtime wazero.Runtime, guest
 		NewFunctionBuilder().
 		WithGoModuleFunction(wazeroapi.GoModuleFunc(host.k8sHandleRejectWaitingPodFn), []wazeroapi.ValueType{i32, i32, i32, i32}, []wazeroapi.ValueType{}).
 		WithParameterNames("buf", "buf_len").Export(k8sSchedulerHandleRejectWaitingPod).
+		NewFunctionBuilder().
+		WithGoModuleFunction(wazeroapi.GoModuleFunc(host.k8sHandleGetWaitingPodFn), []wazeroapi.ValueType{i32, i32, i32, i32}, []wazeroapi.ValueType{}).
+		WithParameterNames("buf", "buf_len").Export(k8sSchedulerHandleGetWaitingPod).
 		Instantiate(ctx)
 }
 
@@ -604,3 +608,34 @@ func (h host) k8sHandleRejectWaitingPodFn(ctx context.Context, mod wazeroapi.Mod
 	}
 	writeUint64(mod.Memory(), wasmBool, oBuf, oBufLimit)
 }
+
+func (h host) k8sHandleGetWaitingPodFn(ctx context.Context, mod wazeroapi.Module, stack []uint64) {
+	iBuf := uint32(stack[0])
+	iBufLen := uint32(stack[1])
+	oBuf := uint32(stack[2])
+	oBufLimit := uint32(stack[3])
+
+	// Safely read UID from memory
+	b, ok := mod.Memory().Read(iBuf, iBufLen)
+	if !ok {
+		panic("out of memory reading getWaitingPod")
+	}
+	uid := types.UID(b)
+
+	// Fetch the waiting pod using the handle
+	waitingPod := h.handle.GetWaitingPod(uid)
+	if waitingPod == nil {
+		// Write an error message or handle this scenario in a non-crashing way
+		writeUint64(mod.Memory(), uint64(0), oBuf, oBufLimit)
+		stack[0] = uint64(0) // Indicate failure (return 0 to WebAssembly)
+		return
+	}
+
+	// Marshal the found pod and ensure no out-of-bounds memory writes
+	vLen := marshalIfUnderLimit(mod.Memory(), waitingPod.GetPod(), oBuf, oBufLimit)
+	if vLen == 0 {
+		panic("out of memory writing getWaitingPod result")
+	}
+	// Indicate success (return 1=success to WebAssembly)
+	stack[0] = uint64(1)
+}

scheduler/plugin/host_test.go

@@ -132,6 +132,34 @@ func Test_k8sHandleEventRecorderEventFn(t *testing.T) {
 	}
 }
 
+func Test_k8sHandleGetWaitingPodFn(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.k8sHandleGetWaitingPodFn(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.GetWaitingPodValue
+	want := uid
+
+	if want != have.GetPod().UID {
+		t.Fatalf("unexpected uid: %v != %v", want, have)
+	}
+}
+
 func Test_k8sHandleRejectWaitingPodFn(t *testing.T) {
 	recorder := &test.FakeRecorder{EventMsg: ""}
 	handle := &test.FakeHandle{Recorder: recorder}

scheduler/plugin/plugin.go

@@ -95,7 +95,6 @@ func newWasmPlugin(ctx context.Context, pluginName string, runtime wazero.Runtim
 		guestModuleConfig: wazero.NewModuleConfig(),
 		instanceCounter:   atomic.Uint64{},
 	}
-
 	if pl.pool, err = newGuestPool(ctx, pl.newGuest); err != nil {
 		return nil, fmt.Errorf("failed to create a guest pool: %w", err)
 	}