Implement server-side filtering for VMs by worker (#392)

* Implement server-side filtering for VMs by worker

* Parse more than one filter but error out when more than one is provided

* Fix off-by-one

* No need to use "\n" in Debugf()

api/openapi.yaml

@@ -243,6 +243,13 @@ paths:
       summary: "List VMs"
       tags:
         - vms
+      parameters:
+        - in: query
+          name: filter
+          description: "Filter VMs using `path=value` syntax; currently only `worker=<name>` is supported to return VMs assigned to the given worker"
+          schema:
+            type: string
+          required: false
       responses:
         '200':
           description: OK

internal/controller/api_service_accounts.go

@@ -2,14 +2,15 @@ package controller
 
 import (
 	"errors"
+	"net/http"
+	"time"
+
 	storepkg "github.com/cirruslabs/orchard/internal/controller/store"
 	"github.com/cirruslabs/orchard/internal/responder"
 	"github.com/cirruslabs/orchard/internal/simplename"
 	v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
 	"github.com/gin-gonic/gin"
 	"github.com/google/uuid"
-	"net/http"
-	"time"
 )
 
 func (controller *Controller) createServiceAccount(ctx *gin.Context) responder.Responder {

internal/controller/api_vms.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"net/http"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/cirruslabs/orchard/internal/controller/lifecycle"
@@ -298,8 +299,30 @@ func (controller *Controller) listVMs(ctx *gin.Context) responder.Responder {
 		return responder
 	}
 
+	var opts []storepkg.ListOption
+
+	if filterRaw := ctx.Query("filter"); filterRaw != "" {
+		var filters []v1.Filter
+
+		for _, filterRaw := range strings.Split(filterRaw, ",") {
+			filter, err := v1.NewFilter(filterRaw)
+			if err != nil {
+				return responder.JSON(http.StatusPreconditionFailed, NewErrorResponse("%v", err))
+			}
+
+			filters = append(filters, filter)
+		}
+
+		if len(filters) > 1 {
+			return responder.JSON(http.StatusPreconditionFailed, NewErrorResponse("only "+
+				"a single filter is currently supported"))
+		}
+
+		opts = append(opts, storepkg.WithListFilters(filters...))
+	}
+
 	return controller.storeView(func(txn storepkg.Transaction) responder.Responder {
-		vms, err := txn.ListVMs()
+		vms, err := txn.ListVMs(opts...)
 		if err != nil {
 			return responder.Error(err)
 		}

internal/controller/store/badger/badger_generic.go

@@ -3,6 +3,8 @@ package badger
 import (
 	"encoding/json"
 
+	storepkg "github.com/cirruslabs/orchard/internal/controller/store"
+	v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
 	"github.com/dgraph-io/badger/v3"
 )
 
@@ -51,12 +53,20 @@ func genericGet[T any, PT interface {
 
 func genericList[T any, PT interface {
 	SetVersion(uint64)
+	Match(v1.Filter) bool
 	*T
-}](txn *Transaction, prefix []byte) (_ []T, err error) {
+}](txn *Transaction, prefix []byte, opts ...storepkg.ListOption) (_ []T, err error) {
 	defer func() {
 		err = mapErr(err)
 	}()
 
+	// Apply options
+	listInput := &storepkg.ListInput{}
+
+	for _, opt := range opts {
+		opt(listInput)
+	}
+
 	// Declare an empty, non-nil slice to
 	// return [] when no objects are found
 	result := []T{}
@@ -66,6 +76,7 @@ func genericList[T any, PT interface {
 	})
 	defer it.Close()
 
+Outer:
 	for it.Rewind(); it.Valid(); it.Next() {
 		item := it.Item()
 
@@ -80,6 +91,12 @@ func genericList[T any, PT interface {
 			return nil, err
 		}
 
+		for _, filter := range listInput.Filters {
+			if !PT(&obj).Match(filter) {
+				continue Outer
+			}
+		}
+
 		PT(&obj).SetVersion(item.Version())
 
 		result = append(result, obj)

internal/controller/store/badger/badger_vm.go

@@ -4,6 +4,7 @@ package badger
 import (
 	"path"
 
+	storepkg "github.com/cirruslabs/orchard/internal/controller/store"
 	"github.com/cirruslabs/orchard/pkg/resource/v1"
 )
 
@@ -25,6 +26,6 @@ func (txn *Transaction) DeleteVM(name string) error {
 	return genericDelete(txn, VMKey(name))
 }
 
-func (txn *Transaction) ListVMs() ([]v1.VM, error) {
+func (txn *Transaction) ListVMs(opts ...storepkg.ListOption) ([]v1.VM, error) {
-	return genericList[v1.VM](txn, []byte(SpaceVMs))
+	return genericList[v1.VM](txn, []byte(SpaceVMs), opts...)
 }

internal/controller/store/option.go

@@ -0,0 +1,15 @@
+package store
+
+import v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
+
+type ListInput struct {
+	Filters []v1.Filter
+}
+
+type ListOption func(listInput *ListInput)
+
+func WithListFilters(filters ...v1.Filter) ListOption {
+	return func(listInput *ListInput) {
+		listInput.Filters = filters
+	}
+}

internal/controller/store/store.go

@@ -29,7 +29,7 @@ type Transaction interface {
 	GetVM(name string) (result *v1.VM, err error)
 	SetVM(vm v1.VM) (err error)
 	DeleteVM(name string) (err error)
-	ListVMs() (result []v1.VM, err error)
+	ListVMs(opts ...ListOption) (result []v1.VM, err error)
 
 	GetWorker(name string) (result *v1.Worker, err error)
 	SetWorker(worker v1.Worker) (err error)

internal/worker/worker.go

@@ -339,7 +339,7 @@ func (worker *Worker) syncVMs(ctx context.Context, updateVM func(context.Context
 		action := transitions[remoteState][localState]
 
 		worker.logger.Debugf("processing VM: %s, remote state: %s, local state: %s, "+
-			"local conditions: [%s], action: %v\n", onDiskName, optionToString(remoteState),
+			"local conditions: [%s], action: %v", onDiskName, optionToString(remoteState),
 			optionToString(localState), v1.ConditionsHumanize(localConditions), action)
 
 		switch action {

pkg/client/option.go

@@ -2,8 +2,11 @@ package client
 
 import (
 	"crypto/x509"
+	"fmt"
+	"strings"
 
 	"github.com/cirruslabs/orchard/internal/dialer"
+	v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
 )
 
 type Option func(*Client)
@@ -32,3 +35,25 @@ func WithDialer(dialer dialer.Dialer) Option {
 		client.dialer = dialer
 	}
 }
+
+type ListInput struct {
+	Filters []v1.Filter
+}
+
+type ListOption func(params map[string]string)
+
+func WithListFilters(filters ...v1.Filter) ListOption {
+	return func(params map[string]string) {
+		if len(filters) == 0 {
+			return
+		}
+
+		var pairs []string
+
+		for _, filter := range filters {
+			pairs = append(pairs, fmt.Sprintf("%s=%s", filter.Path, filter.Value))
+		}
+
+		params["filter"] = strings.Join(pairs, ",")
+	}
+}

pkg/client/vms.go

@@ -57,13 +57,18 @@ func (service *VMsService) Create(ctx context.Context, vm *v1.VM) error {
 }
 
 func (service *VMsService) FindForWorker(ctx context.Context, worker string) ([]v1.VM, error) {
-	allVms, err := service.List(ctx)
+	allVms, err := service.List(ctx, WithListFilters(v1.Filter{
+		Path:  "worker",
+		Value: worker,
+	}))
 	if err != nil {
 		return nil, err
 	}
 
 	var result []v1.VM
 
+	// Backwards compatibility with older Orchard Controllers
+	// that do not support the "filter" query parameter
 	for _, vmResource := range allVms {
 		if vmResource.Worker != worker {
 			continue
@@ -75,11 +80,18 @@ func (service *VMsService) FindForWorker(ctx context.Context, worker string) ([]
 	return result, nil
 }
 
-func (service *VMsService) List(ctx context.Context) ([]v1.VM, error) {
+func (service *VMsService) List(ctx context.Context, opts ...ListOption) ([]v1.VM, error) {
+	params := map[string]string{}
+
+	// Apply options
+	for _, opt := range opts {
+		opt(params)
+	}
+
 	var vms []v1.VM
 
 	err := service.client.request(ctx, http.MethodGet, "vms",
-		nil, &vms, nil)
+		nil, &vms, params)
 	if err != nil {
 		return nil, err
 	}

pkg/resource/v1/filter.go

@@ -0,0 +1,31 @@
+package v1
+
+import (
+	"errors"
+	"fmt"
+	"strings"
+)
+
+var ErrInvalidFilter = errors.New("invalid filter")
+
+type Filter struct {
+	Path  string
+	Value string
+}
+
+func NewFilter(s string) (Filter, error) {
+	parts := strings.SplitN(s, "=", 2)
+
+	if len(parts) != 2 {
+		return Filter{}, fmt.Errorf("%w: expected path=value", ErrInvalidFilter)
+	}
+
+	if parts[0] == "" {
+		return Filter{}, fmt.Errorf("%w: path cannot be empty", ErrInvalidFilter)
+	}
+
+	return Filter{
+		Path:  parts[0],
+		Value: parts[1],
+	}, nil
+}

pkg/resource/v1/filter_test.go

@@ -0,0 +1,55 @@
+package v1_test
+
+import (
+	"testing"
+
+	v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewFilter(t *testing.T) {
+	testCases := []struct {
+		Name  string
+		Input string
+		Err   error
+		Path  string
+		Value string
+	}{
+		{
+			Name:  "simple",
+			Input: "a.b.c=value",
+			Path:  "a.b.c",
+			Value: "value",
+		},
+		{
+			Name:  "value with equals",
+			Input: "a.b.c=d=e",
+			Path:  "a.b.c",
+			Value: "d=e",
+		},
+		{
+			Name:  "empty value",
+			Input: "a.b.c=",
+			Path:  "a.b.c",
+		},
+		{
+			Name:  "missing value",
+			Input: "abc",
+			Err:   v1.ErrInvalidFilter,
+		},
+		{
+			Name:  "missing path",
+			Input: "=value",
+			Err:   v1.ErrInvalidFilter,
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.Name, func(t *testing.T) {
+			filter, err := v1.NewFilter(testCase.Input)
+			require.ErrorIs(t, err, testCase.Err)
+			require.Equal(t, testCase.Path, filter.Path)
+			require.Equal(t, testCase.Value, filter.Value)
+		})
+	}
+}

pkg/resource/v1/service_account.go

@@ -8,3 +8,7 @@ type ServiceAccount struct {
 }
 
 func (serviceAccount *ServiceAccount) SetVersion(_ uint64) {}
+
+func (serviceAccount *ServiceAccount) Match(filter Filter) bool {
+	return false
+}

pkg/resource/v1/v1.go

@@ -102,6 +102,15 @@ func (vm *VM) SetVersion(version uint64) {
 	vm.Version = version
 }
 
+func (vm *VM) Match(filter Filter) bool {
+	switch filter.Path {
+	case "worker":
+		return vm.Worker == filter.Value
+	default:
+		return false
+	}
+}
+
 func (vm *VM) IsScheduled() bool {
 	if ConditionExists(vm.Conditions, ConditionTypeScheduled) {
 		return ConditionIsTrue(vm.Conditions, ConditionTypeScheduled)

pkg/resource/v1/worker.go

@@ -32,3 +32,7 @@ func (worker Worker) Offline(workerOfflineTimeout time.Duration) bool {
 }
 
 func (worker *Worker) SetVersion(_ uint64) {}
+
+func (worker *Worker) Match(filter Filter) bool {
+	return false
+}