From 9cdfd75f799bfbcef6c312704a5ecc0a2c6e7815 Mon Sep 17 00:00:00 2001 From: Nikolay Edigaryev Date: Mon, 17 Nov 2025 18:34:59 +0400 Subject: [PATCH] Badger store: avoid code duplication by using generic methods (#369) * Badger store: avoid code duplication by using generic methods * No need to return PT, can return just *T --- internal/command/controller/bootstrap.go | 7 +- internal/controller/controller.go | 4 +- .../store/badger/badger_cluster_settings.go | 39 +------- .../controller/store/badger/badger_generic.go | 97 +++++++++++++++++++ .../store/badger/badger_service_account.go | 91 ++--------------- internal/controller/store/badger/badger_vm.go | 92 ++---------------- .../controller/store/badger/badger_worker.go | 91 ++--------------- internal/controller/store/store.go | 2 +- pkg/resource/v1/cluster_settings.go | 2 + pkg/resource/v1/service_account.go | 2 + pkg/resource/v1/v1.go | 4 + pkg/resource/v1/worker.go | 2 + 12 files changed, 146 insertions(+), 287 deletions(-) create mode 100644 internal/controller/store/badger/badger_generic.go diff --git a/internal/command/controller/bootstrap.go b/internal/command/controller/bootstrap.go index c6c03e7..c4ad473 100644 --- a/internal/command/controller/bootstrap.go +++ b/internal/command/controller/bootstrap.go @@ -3,14 +3,15 @@ package controller import ( "crypto/tls" "fmt" + "os" + "strings" + "github.com/cirruslabs/orchard/internal/certificatefingerprint" "github.com/cirruslabs/orchard/internal/controller" v1 "github.com/cirruslabs/orchard/pkg/resource/v1" "github.com/pterm/pterm" "github.com/samber/lo" "github.com/sethvargo/go-password/password" - "os" - "strings" ) const BootstrapContextName = "bootstrap-context" @@ -33,7 +34,7 @@ func Bootstrap(controllerInstance *controller.Controller, controllerCert tls.Cer // However, if the BootstrapAdminName service account still exists, // return its credentials. We'll use them for updating the // BootstrapContextName context. - if serviceAccount, ok := lo.Find(serviceAccounts, func(serviceAccount *v1.ServiceAccount) bool { + if serviceAccount, ok := lo.Find(serviceAccounts, func(serviceAccount v1.ServiceAccount) bool { return serviceAccount.Name == BootstrapAdminName }); ok { return serviceAccount.Name, serviceAccount.Token, nil diff --git a/internal/controller/controller.go b/internal/controller/controller.go index cb26602..778aaf1 100644 --- a/internal/controller/controller.go +++ b/internal/controller/controller.go @@ -196,8 +196,8 @@ func New(opts ...Option) (*Controller, error) { return controller, nil } -func (controller *Controller) ServiceAccounts() ([]*v1.ServiceAccount, error) { - var serviceAccounts []*v1.ServiceAccount +func (controller *Controller) ServiceAccounts() ([]v1.ServiceAccount, error) { + var serviceAccounts []v1.ServiceAccount var err error if err := controller.store.View(func(txn storepkg.Transaction) error { diff --git a/internal/controller/store/badger/badger_cluster_settings.go b/internal/controller/store/badger/badger_cluster_settings.go index 88977bf..c1cae6a 100644 --- a/internal/controller/store/badger/badger_cluster_settings.go +++ b/internal/controller/store/badger/badger_cluster_settings.go @@ -1,46 +1,15 @@ package badger import ( - "encoding/json" "github.com/cirruslabs/orchard/pkg/resource/v1" ) var ClusterSettingsKey = []byte("/cluster-settings") -func (txn *Transaction) GetClusterSettings() (_ *v1.ClusterSettings, err error) { - defer func() { - err = mapErr(err) - }() - - item, err := txn.badgerTxn.Get(ClusterSettingsKey) - if err != nil { - return nil, err - } - - valueBytes, err := item.ValueCopy(nil) - if err != nil { - return nil, err - } - - var clusterSettings v1.ClusterSettings - - err = json.Unmarshal(valueBytes, &clusterSettings) - if err != nil { - return nil, err - } - - return &clusterSettings, nil +func (txn *Transaction) GetClusterSettings() (*v1.ClusterSettings, error) { + return genericGet[v1.ClusterSettings](txn, ClusterSettingsKey) } -func (txn *Transaction) SetClusterSettings(clusterSettings v1.ClusterSettings) (err error) { - defer func() { - err = mapErr(err) - }() - - valueBytes, err := json.Marshal(clusterSettings) - if err != nil { - return err - } - - return txn.badgerTxn.Set(ClusterSettingsKey, valueBytes) +func (txn *Transaction) SetClusterSettings(clusterSettings v1.ClusterSettings) error { + return genericSet[v1.ClusterSettings](txn, ClusterSettingsKey, clusterSettings) } diff --git a/internal/controller/store/badger/badger_generic.go b/internal/controller/store/badger/badger_generic.go new file mode 100644 index 0000000..cd63eb2 --- /dev/null +++ b/internal/controller/store/badger/badger_generic.go @@ -0,0 +1,97 @@ +package badger + +import ( + "encoding/json" + + "github.com/dgraph-io/badger/v3" +) + +func genericSet[T any](txn *Transaction, key []byte, obj T) (err error) { + defer func() { + err = mapErr(err) + }() + + valueBytes, err := json.Marshal(obj) + if err != nil { + return err + } + + return txn.badgerTxn.Set(key, valueBytes) +} + +func genericGet[T any, PT interface { + SetVersion(uint64) + *T +}](txn *Transaction, key []byte) (_ *T, err error) { + defer func() { + err = mapErr(err) + }() + + item, err := txn.badgerTxn.Get(key) + if err != nil { + return nil, err + } + + valueBytes, err := item.ValueCopy(nil) + if err != nil { + return nil, err + } + + var obj T + + err = json.Unmarshal(valueBytes, &obj) + if err != nil { + return nil, err + } + + PT(&obj).SetVersion(item.Version()) + + return &obj, nil +} + +func genericList[T any, PT interface { + SetVersion(uint64) + *T +}](txn *Transaction, prefix []byte) (_ []T, err error) { + defer func() { + err = mapErr(err) + }() + + // Declare an empty, non-nil slice to + // return [] when no objects are found + result := []T{} + + it := txn.badgerTxn.NewIterator(badger.IteratorOptions{ + Prefix: prefix, + }) + defer it.Close() + + for it.Rewind(); it.Valid(); it.Next() { + item := it.Item() + + valueBytes, err := item.ValueCopy(nil) + if err != nil { + return nil, err + } + + var obj T + + if err := json.Unmarshal(valueBytes, &obj); err != nil { + return nil, err + } + + PT(&obj).SetVersion(item.Version()) + + result = append(result, obj) + } + + return result, nil +} + +func genericDelete(txn *Transaction, key []byte) (err error) { + defer func() { + err = mapErr(err) + }() + + return txn.badgerTxn.Delete(key) +} diff --git a/internal/controller/store/badger/badger_service_account.go b/internal/controller/store/badger/badger_service_account.go index d8b8132..726565a 100644 --- a/internal/controller/store/badger/badger_service_account.go +++ b/internal/controller/store/badger/badger_service_account.go @@ -1,10 +1,9 @@ package badger import ( - "encoding/json" - "github.com/cirruslabs/orchard/pkg/resource/v1" - "github.com/dgraph-io/badger/v3" "path" + + "github.com/cirruslabs/orchard/pkg/resource/v1" ) const SpaceServiceAccounts = "/service-accounts" @@ -13,88 +12,18 @@ func ServiceAccountKey(name string) []byte { return []byte(path.Join(SpaceServiceAccounts, name)) } -func (txn *Transaction) GetServiceAccount(name string) (_ *v1.ServiceAccount, err error) { - defer func() { - err = mapErr(err) - }() - - key := ServiceAccountKey(name) - - item, err := txn.badgerTxn.Get(key) - if err != nil { - return nil, err - } - - valueBytes, err := item.ValueCopy(nil) - if err != nil { - return nil, err - } - - var serviceAccount v1.ServiceAccount - - err = json.Unmarshal(valueBytes, &serviceAccount) - if err != nil { - return nil, err - } - - return &serviceAccount, nil +func (txn *Transaction) GetServiceAccount(name string) (*v1.ServiceAccount, error) { + return genericGet[v1.ServiceAccount](txn, ServiceAccountKey(name)) } -func (txn *Transaction) SetServiceAccount(serviceAccount *v1.ServiceAccount) (err error) { - defer func() { - err = mapErr(err) - }() - - key := ServiceAccountKey(serviceAccount.Name) - - valueBytes, err := json.Marshal(serviceAccount) - if err != nil { - return err - } - - return txn.badgerTxn.Set(key, valueBytes) +func (txn *Transaction) SetServiceAccount(serviceAccount *v1.ServiceAccount) error { + return genericSet[v1.ServiceAccount](txn, ServiceAccountKey(serviceAccount.Name), *serviceAccount) } -func (txn *Transaction) DeleteServiceAccount(name string) (err error) { - defer func() { - err = mapErr(err) - }() - - key := ServiceAccountKey(name) - - return txn.badgerTxn.Delete(key) +func (txn *Transaction) DeleteServiceAccount(name string) error { + return genericDelete(txn, ServiceAccountKey(name)) } -func (txn *Transaction) ListServiceAccounts() (_ []*v1.ServiceAccount, err error) { - defer func() { - err = mapErr(err) - }() - - // Declare an empty, non-nil slice to return - // [] when no service accounts are found - result := []*v1.ServiceAccount{} - - it := txn.badgerTxn.NewIterator(badger.IteratorOptions{ - Prefix: []byte(SpaceServiceAccounts), - }) - defer it.Close() - - for it.Rewind(); it.Valid(); it.Next() { - item := it.Item() - - serviceAccountBytes, err := item.ValueCopy(nil) - if err != nil { - return nil, err - } - - var serviceAccount v1.ServiceAccount - - if err := json.Unmarshal(serviceAccountBytes, &serviceAccount); err != nil { - return nil, err - } - - result = append(result, &serviceAccount) - } - - return result, nil +func (txn *Transaction) ListServiceAccounts() ([]v1.ServiceAccount, error) { + return genericList[v1.ServiceAccount](txn, []byte(SpaceServiceAccounts)) } diff --git a/internal/controller/store/badger/badger_vm.go b/internal/controller/store/badger/badger_vm.go index 6b8b71d..4a0c020 100644 --- a/internal/controller/store/badger/badger_vm.go +++ b/internal/controller/store/badger/badger_vm.go @@ -2,11 +2,9 @@ package badger import ( - "encoding/json" "path" "github.com/cirruslabs/orchard/pkg/resource/v1" - "github.com/dgraph-io/badger/v3" ) const SpaceVMs = "/vms" @@ -15,92 +13,18 @@ func VMKey(name string) []byte { return []byte(path.Join(SpaceVMs, name)) } -func (txn *Transaction) GetVM(name string) (_ *v1.VM, err error) { - defer func() { - err = mapErr(err) - }() - - key := VMKey(name) - - item, err := txn.badgerTxn.Get(key) - if err != nil { - return nil, err - } - - valueBytes, err := item.ValueCopy(nil) - if err != nil { - return nil, err - } - - var vm v1.VM - - err = json.Unmarshal(valueBytes, &vm) - if err != nil { - return nil, err - } - - vm.Version = item.Version() - - return &vm, nil +func (txn *Transaction) GetVM(name string) (*v1.VM, error) { + return genericGet[v1.VM](txn, VMKey(name)) } -func (txn *Transaction) SetVM(vm v1.VM) (err error) { - defer func() { - err = mapErr(err) - }() - - key := VMKey(vm.Name) - - valueBytes, err := json.Marshal(vm) - if err != nil { - return err - } - - return txn.badgerTxn.Set(key, valueBytes) +func (txn *Transaction) SetVM(vm v1.VM) error { + return genericSet[v1.VM](txn, VMKey(vm.Name), vm) } -func (txn *Transaction) DeleteVM(name string) (err error) { - defer func() { - err = mapErr(err) - }() - - key := VMKey(name) - - return txn.badgerTxn.Delete(key) +func (txn *Transaction) DeleteVM(name string) error { + return genericDelete(txn, VMKey(name)) } -func (txn *Transaction) ListVMs() (_ []v1.VM, err error) { - defer func() { - err = mapErr(err) - }() - - // Declare an empty, non-nil slice to - // return [] when no VMs are found - result := []v1.VM{} - - it := txn.badgerTxn.NewIterator(badger.IteratorOptions{ - Prefix: []byte(SpaceVMs), - }) - defer it.Close() - - for it.Rewind(); it.Valid(); it.Next() { - item := it.Item() - - vmBytes, err := item.ValueCopy(nil) - if err != nil { - return nil, err - } - - var vm v1.VM - - if err := json.Unmarshal(vmBytes, &vm); err != nil { - return nil, err - } - - vm.Version = item.Version() - - result = append(result, vm) - } - - return result, nil +func (txn *Transaction) ListVMs() ([]v1.VM, error) { + return genericList[v1.VM](txn, []byte(SpaceVMs)) } diff --git a/internal/controller/store/badger/badger_worker.go b/internal/controller/store/badger/badger_worker.go index 6aa28d0..8f23913 100644 --- a/internal/controller/store/badger/badger_worker.go +++ b/internal/controller/store/badger/badger_worker.go @@ -2,10 +2,9 @@ package badger import ( - "encoding/json" - "github.com/cirruslabs/orchard/pkg/resource/v1" - "github.com/dgraph-io/badger/v3" "path" + + "github.com/cirruslabs/orchard/pkg/resource/v1" ) const SpaceWorkers = "/workers" @@ -14,88 +13,18 @@ func WorkerKey(name string) []byte { return []byte(path.Join(SpaceWorkers, name)) } -func (txn *Transaction) GetWorker(name string) (_ *v1.Worker, err error) { - defer func() { - err = mapErr(err) - }() - - key := WorkerKey(name) - - item, err := txn.badgerTxn.Get(key) - if err != nil { - return nil, err - } - - valueBytes, err := item.ValueCopy(nil) - if err != nil { - return nil, err - } - - var worker v1.Worker - - err = json.Unmarshal(valueBytes, &worker) - if err != nil { - return nil, err - } - - return &worker, nil +func (txn *Transaction) GetWorker(name string) (*v1.Worker, error) { + return genericGet[v1.Worker](txn, WorkerKey(name)) } -func (txn *Transaction) SetWorker(worker v1.Worker) (err error) { - defer func() { - err = mapErr(err) - }() - - key := WorkerKey(worker.Name) - - valueBytes, err := json.Marshal(worker) - if err != nil { - return err - } - - return txn.badgerTxn.Set(key, valueBytes) +func (txn *Transaction) SetWorker(worker v1.Worker) error { + return genericSet[v1.Worker](txn, WorkerKey(worker.Name), worker) } -func (txn *Transaction) DeleteWorker(name string) (err error) { - defer func() { - err = mapErr(err) - }() - - key := WorkerKey(name) - - return txn.badgerTxn.Delete(key) +func (txn *Transaction) DeleteWorker(name string) error { + return genericDelete(txn, WorkerKey(name)) } -func (txn *Transaction) ListWorkers() (_ []v1.Worker, err error) { - defer func() { - err = mapErr(err) - }() - - // Declare an empty, non-nil slice to - // return [] when no workers are found - result := []v1.Worker{} - - it := txn.badgerTxn.NewIterator(badger.IteratorOptions{ - Prefix: []byte(SpaceWorkers), - }) - defer it.Close() - - for it.Rewind(); it.Valid(); it.Next() { - item := it.Item() - - vmBytes, err := item.ValueCopy(nil) - if err != nil { - return nil, err - } - - var worker v1.Worker - - if err := json.Unmarshal(vmBytes, &worker); err != nil { - return nil, err - } - - result = append(result, worker) - } - - return result, nil +func (txn *Transaction) ListWorkers() ([]v1.Worker, error) { + return genericList[v1.Worker](txn, []byte(SpaceWorkers)) } diff --git a/internal/controller/store/store.go b/internal/controller/store/store.go index fcb50e6..c62f6e2 100644 --- a/internal/controller/store/store.go +++ b/internal/controller/store/store.go @@ -39,7 +39,7 @@ type Transaction interface { GetServiceAccount(name string) (result *v1.ServiceAccount, err error) SetServiceAccount(serviceAccount *v1.ServiceAccount) (err error) DeleteServiceAccount(name string) (err error) - ListServiceAccounts() (result []*v1.ServiceAccount, err error) + ListServiceAccounts() (result []v1.ServiceAccount, err error) AppendEvents(event []v1.Event, scope ...string) (err error) ListEvents(scope ...string) (result []v1.Event, err error) diff --git a/pkg/resource/v1/cluster_settings.go b/pkg/resource/v1/cluster_settings.go index ddde2c8..82df70a 100644 --- a/pkg/resource/v1/cluster_settings.go +++ b/pkg/resource/v1/cluster_settings.go @@ -14,6 +14,8 @@ type ClusterSettings struct { SchedulerProfile SchedulerProfile `json:"schedulerProfile,omitempty"` } +func (clusterSettings *ClusterSettings) SetVersion(_ uint64) {} + func NewSchedulerProfile(value string) (SchedulerProfile, error) { switch value { case string(SchedulerProfileOptimizeUtilization): diff --git a/pkg/resource/v1/service_account.go b/pkg/resource/v1/service_account.go index 1174ca3..87a961b 100644 --- a/pkg/resource/v1/service_account.go +++ b/pkg/resource/v1/service_account.go @@ -6,3 +6,5 @@ type ServiceAccount struct { Meta } + +func (serviceAccount *ServiceAccount) SetVersion(_ uint64) {} diff --git a/pkg/resource/v1/v1.go b/pkg/resource/v1/v1.go index fed4415..2434246 100644 --- a/pkg/resource/v1/v1.go +++ b/pkg/resource/v1/v1.go @@ -97,6 +97,10 @@ type VM struct { Meta } +func (vm *VM) SetVersion(version uint64) { + vm.Version = version +} + type VMSpec struct { NetSoftnetDeprecated bool `json:"net-softnet,omitempty"` NetSoftnet bool `json:"netSoftnet,omitempty"` diff --git a/pkg/resource/v1/worker.go b/pkg/resource/v1/worker.go index b0a3b67..ebe7151 100644 --- a/pkg/resource/v1/worker.go +++ b/pkg/resource/v1/worker.go @@ -30,3 +30,5 @@ type Worker struct { func (worker Worker) Offline(workerOfflineTimeout time.Duration) bool { return time.Since(worker.LastSeen) > workerOfflineTimeout } + +func (worker *Worker) SetVersion(_ uint64) {}