Merge pull request kubernetes-sigs#66 from codefromthecrypt/simplify-…

…host

guest: inlines host.GetConfig to config.Get

examples/nodenumber/main.go

@@ -26,16 +26,16 @@ import (
 	_ "github.com/wasilibs/nottinygc"
 
 	"sigs.k8s.io/kube-scheduler-wasm-extension/examples/nodenumber/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/host"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/prescore"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/score"
 )
 
 // main is compiled to a WebAssembly function named "_start", called by the
 // wasm scheduler plugin during initialization.
 func main() {
-	plugin, err := plugin.New(host.Get())
+	plugin, err := plugin.New(config.Get())
 	if err != nil {
 		panic(err)
 	}

examples/nodenumber/plugin/plugin.go

@@ -52,11 +52,15 @@ type NodeNumber struct {
 	reverse bool
 }
 
-// New creates a new NodeNumber plugin for the given host or returns an error.
-func New(host api.Host) (*NodeNumber, error) {
+// New creates a new NodeNumber plugin with the given jsonConfig or returns an
+// error.
+//
+// Note: This accepts config instead of implicitly calling config.Get for
+// testing.
+func New(jsonConfig []byte) (*NodeNumber, error) {
 	var args nodeNumberArgs
-	if config := host.GetConfig(); config != nil {
-		if err := json.Unmarshal(config, &args); err != nil {
+	if jsonConfig != nil {
+		if err := json.Unmarshal(jsonConfig, &args); err != nil {
 			return nil, fmt.Errorf("decode arg into NodeNumberArgs: %w", err)
 		}
 	}

guest/host/host.go → guest/config/config.go

@@ -14,44 +14,29 @@
    limitations under the License.
 */
 
-// Package host imports an api.Host from the WebAssembly host.
-package host
+// Package config allows reading configuration from the WebAssembly host.
+package config
 
 import (
-	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/api"
+	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/config/internal"
 	"sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/mem"
 )
 
-// Get can be called at any time, to use features exported by the host.
+// Get can be called at any time, to read configuration set by the host.
 //
 // For example:
 //
 //	func main() {
-//		config := host.Get().GetConfig()
+//		config := config.Get()
 //		// decode yaml
 //	}
-func Get() api.Host {
-	return currentHost
+func Get() []byte {
+	return internal.Get(readConfig)
 }
 
-var currentHost api.Host = &host{}
-
-type host struct {
-	config []byte
-}
-
-func (n *host) GetConfig() []byte {
-	return n.lazyConfig()
-}
-
-func (n *host) lazyConfig() []byte {
-	if config := n.config; config != nil {
-		return config
-	}
-
+func readConfig() string {
 	// Wrap to avoid TinyGo 0.28: cannot use an exported function as value
-	n.config = mem.GetBytes(func(ptr uint32, limit mem.BufLimit) (len uint32) {
+	return mem.GetString(func(ptr uint32, limit mem.BufLimit) (len uint32) {
 		return getConfig(ptr, limit)
 	})
-	return n.config
 }

guest/host/imports.go → guest/config/imports.go

@@ -16,7 +16,7 @@
    limitations under the License.
 */
 
-package host
+package config
 
 import "sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/mem"
 

guest/host/imports_stub.go → guest/config/imports_stub.go

@@ -16,7 +16,7 @@
    limitations under the License.
 */
 
-package host
+package config
 
 import "sigs.k8s.io/kube-scheduler-wasm-extension/guest/internal/mem"
 

guest/config/internal/config.go

@@ -0,0 +1,43 @@
+/*
+   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 internal allows unit testing without requiring wasm imports.
+package internal
+
+import "unsafe"
+
+var (
+	// config is lazy read on Get.
+	config string
+	// configRead is required to differentiate an empty read from never read.
+	configRead bool
+)
+
+// Get lazy reads configuration from the given function and returns a byte
+// slice of the result.
+func Get(readConfig func() string) []byte {
+	if !configRead {
+		config = readConfig()
+		configRead = true
+	}
+	if config == "" {
+		return nil // don't call unsafe.StringData("")
+	}
+	// Return the bytes under `config`. This is safe because `config` is
+	// package-scoped, so is always kept alive. This is an alternative to
+	// maintaining `mem.GetBytes` or `[]byte(stringConfig)`, which allocates.
+	return unsafe.Slice(unsafe.StringData(config), len(config))
+}

guest/config/internal/config_test.go

@@ -0,0 +1,64 @@
+/*
+   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 internal
+
+import (
+	"bytes"
+	"testing"
+)
+
+func TestGet(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    string
+		expected []byte
+	}{
+		{name: "empty"},
+		{
+			name:     "NUL terminated",
+			input:    string([]byte{'a', 0, 't', 'w', 'o', 0}),
+			expected: []byte{'a', 0, 't', 'w', 'o', 0},
+		},
+		{
+			name:     "unicode",
+			input:    "fóo",
+			expected: []byte{'f', 0xc3, 0xb3, 'o'},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			configRead = false
+			defer func() {
+				configRead = false
+			}()
+
+			for _, fn := range []func() string{
+				func() string {
+					return tc.input
+				},
+				func() string {
+					panic("should cache the first read")
+				},
+			} {
+				if want, have := tc.expected, Get(fn); (want == nil && have != nil) || !bytes.Equal(want, have) {
+					t.Fatalf("unexpected config: %v != %v", want, have)
+				}
+			}
+		})
+	}
+}

guest/internal/mem/mem.go

@@ -68,27 +68,6 @@ func Update(
 	return updater(readBuf)
 }
 
-func GetBytes(fn func(ptr uint32, limit BufLimit) (len uint32)) []byte {
-	size := fn(uint32(readBufPtr), readBufLimit)
-	if size == 0 {
-		return nil
-	}
-
-	buf := make([]byte, size)
-
-	// If the function result fit in our read buffer, copy it out.
-	if size <= readBufLimit {
-		copy(buf, readBuf[:size])
-		return buf
-	}
-
-	// If the size returned from the function was larger than our read buffer,
-	// we need to execute it again.
-	ptr := unsafe.Pointer(&buf[0])
-	_ = fn(uint32(uintptr(ptr)), size)
-	return buf
-}
-
 func GetString(fn func(ptr uint32, limit BufLimit) (len uint32)) string {
 	size := fn(uint32(readBufPtr), readBufLimit)
 	if size == 0 {