diff --git a/go.mod b/go.mod index 32866f4..14bc612 100644 --- a/go.mod +++ b/go.mod @@ -35,6 +35,7 @@ require ( github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966 github.com/spf13/cobra v1.10.2 github.com/stretchr/testify v1.11.1 + go.etcd.io/etcd/client/v3 v3.6.5 go.opentelemetry.io/contrib/instrumentation/github.com/gin-gonic/gin/otelgin v0.67.0 go.opentelemetry.io/otel v1.42.0 go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetricgrpc v1.42.0 @@ -69,6 +70,8 @@ require ( github.com/clipperhouse/uax29/v2 v2.7.0 // indirect github.com/cloudwego/base64x v0.1.6 // indirect github.com/containerd/console v1.0.5 // indirect + github.com/coreos/go-semver v0.3.1 // indirect + github.com/coreos/go-systemd/v22 v22.5.0 // indirect github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect github.com/dgraph-io/ristretto v0.1.1 // indirect github.com/ebitengine/purego v0.10.0 // indirect @@ -134,6 +137,8 @@ require ( github.com/ugorji/go/codec v1.3.1 // indirect github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect github.com/yusufpapurcu/wmi v1.2.4 // indirect + go.etcd.io/etcd/api/v3 v3.6.5 // indirect + go.etcd.io/etcd/client/pkg/v3 v3.6.5 // indirect go.mongodb.org/mongo-driver/v2 v2.5.0 // indirect go.opencensus.io v0.22.5 // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect diff --git a/go.sum b/go.sum index 0edf211..77d834d 100644 --- a/go.sum +++ b/go.sum @@ -61,6 +61,10 @@ github.com/containerd/console v1.0.5/go.mod h1:YynlIjWYF8myEu6sdkwKIvGQq+cOckRm6 github.com/coreos/etcd v3.3.10+incompatible/go.mod h1:uF7uidLiAD3TWHmW31ZFd/JWoc32PjwdhPthX9715RE= github.com/coreos/go-etcd v2.0.0+incompatible/go.mod h1:Jez6KQU2B/sWsbdaef3ED8NzMklzPG4d5KIOhIy30Tk= github.com/coreos/go-semver v0.2.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk= +github.com/coreos/go-semver v0.3.1 h1:yi21YpKnrx1gt5R+la8n5WgS0kCrsPp33dmEyHReZr4= +github.com/coreos/go-semver v0.3.1/go.mod h1:irMmmIw/7yzSRPWryHsK7EYSg09caPQL03VsM8rvUec= +github.com/coreos/go-systemd/v22 v22.5.0 h1:RrqgGjYQKalulkV8NGVIfkXQf6YYmOyiJKk8iXXhfZs= +github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc= github.com/cpuguy83/go-md2man v1.0.10/go.mod h1:SmD6nW6nTyfqj6ABTjUi3V3JVMnlJmwcJI5acqYI6dE= github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= @@ -157,6 +161,7 @@ github.com/goccy/go-json v0.10.5 h1:Fq85nIqj+gXn/S5ahsiTlK3TmC85qgirsdTP/+DeaC4= github.com/goccy/go-json v0.10.5/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M= github.com/goccy/go-yaml v1.19.2 h1:PmFC1S6h8ljIz6gMRBopkjP1TVT7xuwrButHID66PoM= github.com/goccy/go-yaml v1.19.2/go.mod h1:XBurs7gK8ATbW4ZPGKgcbrY1Br56PdM69F7LkFRi1kA= +github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA= github.com/gofrs/flock v0.13.0 h1:95JolYOvGMqeH31+FC7D2+uULf6mG61mEZ/A8dRYMzw= github.com/gofrs/flock v0.13.0/go.mod h1:jxeyy9R1auM5S6JYDBhDt+E2TCo7DkratH4Pgi8P+Z0= github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q= @@ -347,6 +352,12 @@ github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9dec github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= github.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0= github.com/yusufpapurcu/wmi v1.2.4/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0= +go.etcd.io/etcd/api/v3 v3.6.5 h1:pMMc42276sgR1j1raO/Qv3QI9Af/AuyQUW6CBAWuntA= +go.etcd.io/etcd/api/v3 v3.6.5/go.mod h1:ob0/oWA/UQQlT1BmaEkWQzI0sJ1M0Et0mMpaABxguOQ= +go.etcd.io/etcd/client/pkg/v3 v3.6.5 h1:Duz9fAzIZFhYWgRjp/FgNq2gO1jId9Yae/rLn3RrBP8= +go.etcd.io/etcd/client/pkg/v3 v3.6.5/go.mod h1:8Wx3eGRPiy0qOFMZT/hfvdos+DjEaPxdIDiCDUv/FQk= +go.etcd.io/etcd/client/v3 v3.6.5 h1:yRwZNFBx/35VKHTcLDeO7XVLbCBFbPi+XV4OC3QJf2U= +go.etcd.io/etcd/client/v3 v3.6.5/go.mod h1:ZqwG/7TAFZ0BJ0jXRPoJjKQJtbFo/9NIY8uoFFKcCyo= go.mongodb.org/mongo-driver/v2 v2.5.0 h1:yXUhImUjjAInNcpTcAlPHiT7bIXhshCTL3jVBkF3xaE= go.mongodb.org/mongo-driver/v2 v2.5.0/go.mod h1:yOI9kBsufol30iFsl1slpdq1I0eHPzybRWdyYUs8K/0= go.opencensus.io v0.22.5 h1:dntmOdLpSpHlVqbW5Eay97DelsZHe+55D+xC6i0dDS0= diff --git a/internal/command/controller/run.go b/internal/command/controller/run.go index b8cfff0..32dd612 100644 --- a/internal/command/controller/run.go +++ b/internal/command/controller/run.go @@ -10,6 +10,7 @@ import ( "os" "path/filepath" "strconv" + "strings" "time" configpkg "github.com/cirruslabs/orchard/internal/config" @@ -34,6 +35,9 @@ var experimentalRPCV2 bool var noExperimentalRPCV2 bool var experimentalPingInterval time.Duration var experimentalDisableDBCompression bool +var experimentalStoreBackend string +var experimentalEtcdEndpoints string +var experimentalEtcdKeyPrefix string var workerOfflineTimeout time.Duration var execSessionRetentionTTL time.Duration var execSSHConnectionKeepaliveInterval time.Duration @@ -85,6 +89,12 @@ func newRunCommand() *cobra.Command { "smaller than the controller's default 30 second interval") cmd.Flags().BoolVar(&experimentalDisableDBCompression, "experimental-disable-db-compression", false, "disable database compression, which might reduce RAM usage in some scenarios") + cmd.Flags().StringVar(&experimentalStoreBackend, "experimental-store", string(controller.StoreBackendBadger), + "store backend to use (badger or etcd)") + cmd.Flags().StringVar(&experimentalEtcdEndpoints, "experimental-etcd-endpoints", "localhost:2379", + "comma-separated etcd endpoints used when --experimental-store=etcd") + cmd.Flags().StringVar(&experimentalEtcdKeyPrefix, "experimental-etcd-key-prefix", "/orchard", + "etcd key prefix used when --experimental-store=etcd") cmd.Flags().DurationVar(&workerOfflineTimeout, "worker-offline-timeout", 3*time.Minute, "duration (e.g. 60s or 5m30s) after which a worker is considered offline for the purposes "+ "of scheduling (no new VMs will be scheduled on such worker and already assigned VMs will be "+ @@ -219,6 +229,17 @@ func runController(cmd *cobra.Command, args []string) (err error) { controllerOpts = append(controllerOpts, controller.WithDisableDBCompression()) } + switch controller.StoreBackend(experimentalStoreBackend) { + case controller.StoreBackendBadger: + // Default store backend. + case controller.StoreBackendEtcd: + controllerOpts = append(controllerOpts, controller.WithEtcdStore( + splitEtcdEndpoints(experimentalEtcdEndpoints), experimentalEtcdKeyPrefix, + )) + default: + return fmt.Errorf("unsupported --experimental-store value %q", experimentalStoreBackend) + } + if execSSHConnectionKeepaliveInterval < 5*time.Second { return fmt.Errorf("--exec-ssh-connection-keepalive-interval's value cannot be less than 5 seconds") } @@ -246,6 +267,20 @@ func runController(cmd *cobra.Command, args []string) (err error) { return controllerInstance.Run(cmd.Context()) } +func splitEtcdEndpoints(rawEndpoints string) []string { + var endpoints []string + for _, endpoint := range strings.Split(rawEndpoints, ",") { + endpoint = strings.TrimSpace(endpoint) + if endpoint == "" { + continue + } + + endpoints = append(endpoints, endpoint) + } + + return endpoints +} + func createBootstrapContext( controllerAddress string, controllerCert tls.Certificate, diff --git a/internal/controller/controller.go b/internal/controller/controller.go index 70090d8..cb33066 100644 --- a/internal/controller/controller.go +++ b/internal/controller/controller.go @@ -17,6 +17,7 @@ import ( "github.com/cirruslabs/orchard/internal/controller/sshserver" storepkg "github.com/cirruslabs/orchard/internal/controller/store" "github.com/cirruslabs/orchard/internal/controller/store/badger" + etcdstore "github.com/cirruslabs/orchard/internal/controller/store/etcd" "github.com/cirruslabs/orchard/internal/netconstants" "github.com/cirruslabs/orchard/internal/opentelemetry" v1 "github.com/cirruslabs/orchard/pkg/resource/v1" @@ -59,6 +60,9 @@ type Controller struct { execSSHConnectionKeepaliveInterval time.Duration experimentalRPCV2 bool disableDBCompression bool + storeBackend StoreBackend + etcdEndpoints []string + etcdKeyPrefix string pingInterval time.Duration synthetic bool @@ -108,8 +112,7 @@ func New(opts ...Option) (*Controller, error) { ) // Instantiate the database - store, err := badger.NewBadgerStore(controller.dataDir.DBPath(), controller.disableDBCompression, - controller.logger) + store, err := controller.initStore() if err != nil { return nil, err } @@ -198,6 +201,18 @@ func New(opts ...Option) (*Controller, error) { return controller, nil } +func (controller *Controller) initStore() (storepkg.Store, error) { + switch controller.storeBackend { + case "", StoreBackendBadger: + return badger.NewBadgerStore(controller.dataDir.DBPath(), controller.disableDBCompression, + controller.logger) + case StoreBackendEtcd: + return etcdstore.NewEtcdStore(controller.etcdEndpoints, controller.etcdKeyPrefix, controller.logger) + default: + return nil, fmt.Errorf("%w: unsupported store backend %q", ErrInitFailed, controller.storeBackend) + } +} + func (controller *Controller) vmsEnsurePlatformDefaults() error { return controller.store.Update(func(txn storepkg.Transaction) error { vms, err := txn.ListVMs() diff --git a/internal/controller/option.go b/internal/controller/option.go index adda403..207829a 100644 --- a/internal/controller/option.go +++ b/internal/controller/option.go @@ -10,6 +10,13 @@ import ( type Option func(*Controller) +type StoreBackend string + +const ( + StoreBackendBadger StoreBackend = "badger" + StoreBackendEtcd StoreBackend = "etcd" +) + func WithDataDir(dataDir *DataDir) Option { return func(controller *Controller) { controller.dataDir = dataDir @@ -84,6 +91,14 @@ func WithDisableDBCompression() Option { } } +func WithEtcdStore(endpoints []string, keyPrefix string) Option { + return func(controller *Controller) { + controller.storeBackend = StoreBackendEtcd + controller.etcdEndpoints = endpoints + controller.etcdKeyPrefix = keyPrefix + } +} + func WithPingInterval(pingInterval time.Duration) Option { return func(controller *Controller) { controller.pingInterval = pingInterval diff --git a/internal/controller/store/etcd/etcd_events.go b/internal/controller/store/etcd/etcd_events.go new file mode 100644 index 0000000..54f8f96 --- /dev/null +++ b/internal/controller/store/etcd/etcd_events.go @@ -0,0 +1,156 @@ +package etcd + +import ( + "bytes" + "encoding/json" + "fmt" + "path" + "time" + + storepkg "github.com/cirruslabs/orchard/internal/controller/store" + "github.com/cirruslabs/orchard/pkg/resource/v1" + clientv3 "go.etcd.io/etcd/client/v3" +) + +const SpaceEvents = "/events" + +func scopePrefix(scope []string) string { + keyParts := []string{SpaceEvents} + keyParts = append(keyParts, scope...) + + return path.Join(keyParts...) +} + +func (txn *Transaction) AppendEvents(events []v1.Event, scope ...string) error { + injectionTime := time.Now().UnixNano() + + for index, event := range events { + valueBytes, err := json.Marshal(event) + if err != nil { + return err + } + + eventUID := fmt.Sprintf("/%d-%d-%06d", + event.Timestamp, + injectionTime, + index, + ) + eventKey := txn.store.key(scopePrefix(scope) + eventUID) + + txn.puts[eventKey] = string(valueBytes) + delete(txn.deletes, eventKey) + } + + return nil +} + +func (txn *Transaction) ListEvents(scope ...string) ([]v1.Event, error) { + page, err := txn.ListEventsPage(storepkg.ListOptions{}, scope...) + if err != nil { + return nil, err + } + + return page.Items, nil +} + +func (txn *Transaction) ListEventsPage(options storepkg.ListOptions, scope ...string) ( + storepkg.Page[v1.Event], + error, +) { + var result storepkg.Page[v1.Event] + result.Items = []v1.Event{} + + logicalPrefix := scopePrefix(scope) + physicalPrefix := txn.store.keyPrefix(logicalPrefix) + getKey, getOptions := txn.listEventsPageQueryOptions(physicalPrefix, logicalPrefix, options) + response, err := txn.store.client.Get(txn.ctx, getKey, getOptions...) + if err != nil { + return result, mapErr(err) + } + txn.prefixReadRevisions[physicalPrefix] = response.Header.Revision + + limit := options.Limit + for _, kv := range response.Kvs { + key := string(kv.Key) + if txn.isDeleted(key) { + continue + } + + var event v1.Event + if err := json.Unmarshal(kv.Value, &event); err != nil { + return result, err + } + if limit > 0 && len(result.Items) >= limit { + break + } + + result.Items = append(result.Items, event) + + if limit > 0 && len(result.Items) == limit && len(response.Kvs) > limit { + result.NextCursor = bytes.TrimPrefix([]byte(key), []byte(physicalPrefix)) + } + } + + return result, nil +} + +func (txn *Transaction) DeleteEvents(scope ...string) error { + physicalPrefix := txn.store.keyPrefix(scopePrefix(scope)) + response, err := txn.store.client.Get(txn.ctx, physicalPrefix, clientv3.WithPrefix(), clientv3.WithKeysOnly(), + clientv3.WithLimit(1)) + if err != nil { + return mapErr(err) + } + txn.prefixReadRevisions[physicalPrefix] = response.Header.Revision + + txn.prefixDeletes[physicalPrefix] = struct{}{} + for key := range txn.puts { + if hasPrefix(key, physicalPrefix) { + delete(txn.puts, key) + } + } + + return nil +} + +func (txn *Transaction) listEventsPageQueryOptions( + physicalPrefix string, + logicalPrefix string, + options storepkg.ListOptions, +) (string, []clientv3.OpOption) { + rangeEnd := clientv3.GetPrefixRangeEnd(physicalPrefix) + getKey := physicalPrefix + getRangeEnd := rangeEnd + if len(options.Cursor) > 0 { + cursor := eventCursor(physicalPrefix, logicalPrefix, options.Cursor) + if options.Order == storepkg.ListOrderDesc { + getRangeEnd = cursor + } else { + getKey = cursor + "\x00" + } + } + + getOptions := []clientv3.OpOption{ + clientv3.WithRange(getRangeEnd), + clientv3.WithSort(clientv3.SortByKey, clientv3.SortAscend), + } + if options.Order == storepkg.ListOrderDesc { + getOptions[1] = clientv3.WithSort(clientv3.SortByKey, clientv3.SortDescend) + } + if options.Limit > 0 { + getOptions = append(getOptions, clientv3.WithLimit(int64(options.Limit+1))) + } + + return getKey, getOptions +} + +func eventCursor(physicalPrefix, logicalPrefix string, cursor []byte) string { + if bytes.HasPrefix(cursor, []byte(physicalPrefix)) { + return string(cursor) + } + if bytes.HasPrefix(cursor, []byte(logicalPrefix)) { + return path.Join(physicalPrefix, string(bytes.TrimPrefix(cursor, []byte(logicalPrefix)))) + } + + return physicalPrefix + string(cursor) +} diff --git a/internal/controller/store/etcd/etcd_generic.go b/internal/controller/store/etcd/etcd_generic.go new file mode 100644 index 0000000..96dc1ec --- /dev/null +++ b/internal/controller/store/etcd/etcd_generic.go @@ -0,0 +1,123 @@ +package etcd + +import ( + "encoding/json" + + storepkg "github.com/cirruslabs/orchard/internal/controller/store" + clientv3 "go.etcd.io/etcd/client/v3" +) + +func genericSet[T any](txn *Transaction, key string, obj T) error { + valueBytes, err := json.Marshal(obj) + if err != nil { + return err + } + + txn.puts[key] = string(valueBytes) + delete(txn.deletes, key) + + return nil +} + +func genericGet[T any, PT interface { + SetVersion(uint64) + *T +}](txn *Transaction, key string) (*T, error) { + if value, ok := txn.puts[key]; ok { + var obj T + if err := json.Unmarshal([]byte(value), &obj); err != nil { + return nil, err + } + + return &obj, nil + } + if txn.isDeleted(key) { + return nil, storepkg.ErrNotFound + } + + response, err := txn.store.client.Get(txn.ctx, key) + if err != nil { + return nil, mapErr(err) + } + if len(response.Kvs) == 0 { + txn.readRevisions[key] = 0 + + return nil, storepkg.ErrNotFound + } + + kv := response.Kvs[0] + txn.readRevisions[key] = kv.ModRevision + + var obj T + if err := json.Unmarshal(kv.Value, &obj); err != nil { + return nil, err + } + PT(&obj).SetVersion(uint64(kv.ModRevision)) + + return &obj, nil +} + +func genericList[T any, PT interface { + SetVersion(uint64) + *T +}](txn *Transaction, logicalPrefix string) ([]T, error) { + physicalPrefix := txn.store.keyPrefix(logicalPrefix) + response, err := txn.store.client.Get(txn.ctx, physicalPrefix, clientv3.WithPrefix()) + if err != nil { + return nil, mapErr(err) + } + txn.prefixReadRevisions[physicalPrefix] = response.Header.Revision + + itemsByKey := map[string]T{} + keys := make([]string, 0, len(response.Kvs)+len(txn.puts)) + + for _, kv := range response.Kvs { + var obj T + if err := json.Unmarshal(kv.Value, &obj); err != nil { + return nil, err + } + PT(&obj).SetVersion(uint64(kv.ModRevision)) + + key := string(kv.Key) + itemsByKey[key] = obj + keys = append(keys, key) + } + + for key, value := range txn.puts { + if !hasPrefix(key, physicalPrefix) { + continue + } + + var obj T + if err := json.Unmarshal([]byte(value), &obj); err != nil { + return nil, err + } + + if _, ok := itemsByKey[key]; !ok { + keys = append(keys, key) + } + itemsByKey[key] = obj + } + + result := []T{} + for _, key := range keys { + if txn.isDeleted(key) { + continue + } + + result = append(result, itemsByKey[key]) + } + + return result, nil +} + +func genericDelete(txn *Transaction, key string) error { + delete(txn.puts, key) + txn.deletes[key] = struct{}{} + + return nil +} + +func hasPrefix(value, prefix string) bool { + return len(value) >= len(prefix) && value[:len(prefix)] == prefix +} diff --git a/internal/controller/store/etcd/etcd_resources.go b/internal/controller/store/etcd/etcd_resources.go new file mode 100644 index 0000000..509e2f1 --- /dev/null +++ b/internal/controller/store/etcd/etcd_resources.go @@ -0,0 +1,82 @@ +package etcd + +import ( + "path" + + "github.com/cirruslabs/orchard/pkg/resource/v1" +) + +const ( + SpaceVMs = "/vms" + SpaceWorkers = "/workers" + SpaceServiceAccounts = "/service-accounts" + clusterSettingsKey = "/cluster-settings" +) + +func (store *Store) VMKey(name string) string { + return store.key(path.Join(SpaceVMs, name)) +} + +func (txn *Transaction) GetVM(name string) (*v1.VM, error) { + return genericGet[v1.VM](txn, txn.store.VMKey(name)) +} + +func (txn *Transaction) SetVM(vm v1.VM) error { + return genericSet[v1.VM](txn, txn.store.VMKey(vm.Name), vm) +} + +func (txn *Transaction) DeleteVM(name string) error { + return genericDelete(txn, txn.store.VMKey(name)) +} + +func (txn *Transaction) ListVMs() ([]v1.VM, error) { + return genericList[v1.VM](txn, SpaceVMs) +} + +func (store *Store) workerKey(name string) string { + return store.key(path.Join(SpaceWorkers, name)) +} + +func (txn *Transaction) GetWorker(name string) (*v1.Worker, error) { + return genericGet[v1.Worker](txn, txn.store.workerKey(name)) +} + +func (txn *Transaction) SetWorker(worker v1.Worker) error { + return genericSet[v1.Worker](txn, txn.store.workerKey(worker.Name), worker) +} + +func (txn *Transaction) DeleteWorker(name string) error { + return genericDelete(txn, txn.store.workerKey(name)) +} + +func (txn *Transaction) ListWorkers() ([]v1.Worker, error) { + return genericList[v1.Worker](txn, SpaceWorkers) +} + +func (store *Store) serviceAccountKey(name string) string { + return store.key(path.Join(SpaceServiceAccounts, name)) +} + +func (txn *Transaction) GetServiceAccount(name string) (*v1.ServiceAccount, error) { + return genericGet[v1.ServiceAccount](txn, txn.store.serviceAccountKey(name)) +} + +func (txn *Transaction) SetServiceAccount(serviceAccount *v1.ServiceAccount) error { + return genericSet[v1.ServiceAccount](txn, txn.store.serviceAccountKey(serviceAccount.Name), *serviceAccount) +} + +func (txn *Transaction) DeleteServiceAccount(name string) error { + return genericDelete(txn, txn.store.serviceAccountKey(name)) +} + +func (txn *Transaction) ListServiceAccounts() ([]v1.ServiceAccount, error) { + return genericList[v1.ServiceAccount](txn, SpaceServiceAccounts) +} + +func (txn *Transaction) GetClusterSettings() (*v1.ClusterSettings, error) { + return genericGet[v1.ClusterSettings](txn, txn.store.key(clusterSettingsKey)) +} + +func (txn *Transaction) SetClusterSettings(clusterSettings v1.ClusterSettings) error { + return genericSet[v1.ClusterSettings](txn, txn.store.key(clusterSettingsKey), clusterSettings) +} diff --git a/internal/controller/store/etcd/etcd_store.go b/internal/controller/store/etcd/etcd_store.go new file mode 100644 index 0000000..cea718c --- /dev/null +++ b/internal/controller/store/etcd/etcd_store.go @@ -0,0 +1,168 @@ +package etcd + +import ( + "context" + "errors" + "fmt" + "path" + "strings" + "time" + + "github.com/avast/retry-go/v4" + storepkg "github.com/cirruslabs/orchard/internal/controller/store" + clientv3 "go.etcd.io/etcd/client/v3" + "go.uber.org/zap" +) + +const defaultDialTimeout = 5 * time.Second + +type Store struct { + client *clientv3.Client + prefix string +} + +type Transaction struct { + ctx context.Context + store *Store + readRevisions map[string]int64 + prefixReadRevisions map[string]int64 + puts map[string]string + deletes map[string]struct{} + prefixDeletes map[string]struct{} +} + +func NewEtcdStore(endpoints []string, keyPrefix string, logger *zap.SugaredLogger) (storepkg.Store, error) { + if len(endpoints) == 0 { + return nil, fmt.Errorf("%w: at least one etcd endpoint is required", storepkg.ErrStoreFailed) + } + + client, err := clientv3.New(clientv3.Config{ + Endpoints: endpoints, + DialTimeout: defaultDialTimeout, + Logger: logger.Desugar(), + }) + if err != nil { + return nil, mapErr(err) + } + + return &Store{ + client: client, + prefix: normalizePrefix(keyPrefix), + }, nil +} + +func (store *Store) View(cb func(txn storepkg.Transaction) error) error { + txn := store.newTransaction(context.Background()) + + return cb(txn) +} + +func (store *Store) Update(cb func(txn storepkg.Transaction) error) error { + return retry.Do(func() error { + txn := store.newTransaction(context.Background()) + if err := cb(txn); err != nil { + return err + } + + return txn.commit() + }, retry.RetryIf(func(err error) bool { + return errors.Is(err, storepkg.ErrConflict) + }), retry.Attempts(3), retry.LastErrorOnly(true)) +} + +func (store *Store) newTransaction(ctx context.Context) *Transaction { + return &Transaction{ + ctx: ctx, + store: store, + readRevisions: map[string]int64{}, + prefixReadRevisions: map[string]int64{}, + puts: map[string]string{}, + deletes: map[string]struct{}{}, + prefixDeletes: map[string]struct{}{}, + } +} + +func (store *Store) key(parts ...string) string { + keyParts := []string{store.prefix} + keyParts = append(keyParts, parts...) + + return path.Join(keyParts...) +} + +func (store *Store) keyPrefix(logicalPrefix string) string { + return store.key(logicalPrefix) +} + +func (txn *Transaction) commit() error { + if len(txn.puts) == 0 && len(txn.deletes) == 0 && len(txn.prefixDeletes) == 0 { + return nil + } + + comparisons := make([]clientv3.Cmp, 0, len(txn.readRevisions)+len(txn.prefixReadRevisions)) + for key, revision := range txn.readRevisions { + if revision == 0 { + comparisons = append(comparisons, clientv3.Compare(clientv3.CreateRevision(key), "=", 0)) + } else { + comparisons = append(comparisons, clientv3.Compare(clientv3.ModRevision(key), "=", revision)) + } + } + for prefix, revision := range txn.prefixReadRevisions { + comparisons = append(comparisons, clientv3.Compare(clientv3.ModRevision(prefix).WithPrefix(), "<", revision+1)) + } + + operations := make([]clientv3.Op, 0, len(txn.puts)+len(txn.deletes)+len(txn.prefixDeletes)) + for prefix := range txn.prefixDeletes { + operations = append(operations, clientv3.OpDelete(prefix, clientv3.WithPrefix())) + } + for key, value := range txn.puts { + if _, deleted := txn.deletes[key]; deleted { + continue + } + + operations = append(operations, clientv3.OpPut(key, value)) + } + for key := range txn.deletes { + operations = append(operations, clientv3.OpDelete(key)) + } + + response, err := txn.store.client.Txn(txn.ctx).If(comparisons...).Then(operations...).Commit() + if err != nil { + return mapErr(err) + } + if !response.Succeeded { + return storepkg.ErrConflict + } + + return nil +} + +func (txn *Transaction) isDeleted(key string) bool { + if _, deleted := txn.deletes[key]; deleted { + return true + } + + for prefix := range txn.prefixDeletes { + if hasPrefix(key, prefix) { + return true + } + } + + return false +} + +func normalizePrefix(prefix string) string { + prefix = strings.TrimSpace(prefix) + if prefix == "" || prefix == "/" { + return "/" + } + + return "/" + strings.Trim(prefix, "/") +} + +func mapErr(err error) error { + if err == nil { + return nil + } + + return fmt.Errorf("%w: %v", storepkg.ErrStoreFailed, err) +} diff --git a/internal/controller/store/etcd/etcd_watch.go b/internal/controller/store/etcd/etcd_watch.go new file mode 100644 index 0000000..2fd32b8 --- /dev/null +++ b/internal/controller/store/etcd/etcd_watch.go @@ -0,0 +1,88 @@ +package etcd + +import ( + "context" + "encoding/json" + + storepkg "github.com/cirruslabs/orchard/internal/controller/store" + "github.com/cirruslabs/orchard/pkg/resource/v1" + clientv3 "go.etcd.io/etcd/client/v3" +) + +func (store *Store) WatchVM(ctx context.Context, vmName string) (chan storepkg.WatchMessage[v1.VM], chan error, error) { + watchCh := make(chan storepkg.WatchMessage[v1.VM], 1) + errCh := make(chan error, 1) + key := store.VMKey(vmName) + + response, err := store.client.Get(ctx, key) + if err != nil { + return nil, nil, mapErr(err) + } + + exists := len(response.Kvs) != 0 + watchRevision := response.Header.Revision + 1 + if exists { + var vm v1.VM + if err := json.Unmarshal(response.Kvs[0].Value, &vm); err != nil { + return nil, nil, err + } + vm.Version = uint64(response.Kvs[0].ModRevision) + + watchCh <- storepkg.WatchMessage[v1.VM]{ + Type: storepkg.WatchMessageTypeAdded, + Object: vm, + } + } + + go func() { + defer close(watchCh) + defer close(errCh) + + watchResponses := store.client.Watch(ctx, key, clientv3.WithRev(watchRevision)) + for watchResponse := range watchResponses { + if err := watchResponse.Err(); err != nil { + errCh <- mapErr(err) + + return + } + + for _, event := range watchResponse.Events { + switch event.Type { + case clientv3.EventTypeDelete: + if !exists { + continue + } + + exists = false + select { + case watchCh <- storepkg.WatchMessage[v1.VM]{Type: storepkg.WatchMessageTypeDeleted}: + case <-ctx.Done(): + return + } + case clientv3.EventTypePut: + var vm v1.VM + if err := json.Unmarshal(event.Kv.Value, &vm); err != nil { + errCh <- err + + return + } + vm.Version = uint64(event.Kv.ModRevision) + + messageType := storepkg.WatchMessageTypeAdded + if exists { + messageType = storepkg.WatchMessageTypeModified + } + exists = true + + select { + case watchCh <- storepkg.WatchMessage[v1.VM]{Type: messageType, Object: vm}: + case <-ctx.Done(): + return + } + } + } + } + }() + + return watchCh, errCh, nil +} diff --git a/internal/controller/store/events_test.go b/internal/controller/store/events_test.go index feab1a4..6bb3b1c 100644 --- a/internal/controller/store/events_test.go +++ b/internal/controller/store/events_test.go @@ -4,7 +4,6 @@ import ( "testing" storepkg "github.com/cirruslabs/orchard/internal/controller/store" - "github.com/cirruslabs/orchard/internal/controller/store/badger" "github.com/cirruslabs/orchard/pkg/resource/v1" "github.com/stretchr/testify/require" "go.uber.org/zap" @@ -12,9 +11,16 @@ import ( func TestListEventsPage(t *testing.T) { logger := zap.NewNop().Sugar() - store, err := badger.NewBadgerStore(t.TempDir(), true, logger) - require.NoError(t, err) + for _, storeImpl := range testStores(logger) { + t.Run(storeImpl.Name, func(t *testing.T) { + store := storeImpl.Init(t) + testListEventsPage(t, store) + }) + } +} + +func testListEventsPage(t *testing.T, store storepkg.Store) { events := []v1.Event{ {Kind: v1.EventKindLogLine, Timestamp: 1, Payload: "one"}, {Kind: v1.EventKindLogLine, Timestamp: 2, Payload: "two"}, @@ -22,7 +28,7 @@ func TestListEventsPage(t *testing.T) { {Kind: v1.EventKindLogLine, Timestamp: 4, Payload: "four"}, } - err = store.Update(func(txn storepkg.Transaction) error { + err := store.Update(func(txn storepkg.Transaction) error { return txn.AppendEvents(events, "vms", "vm-uid") }) require.NoError(t, err) @@ -75,3 +81,43 @@ func TestListEventsPage(t *testing.T) { require.Equal(t, []v1.Event{events[1], events[0]}, descPage2.Items) require.Empty(t, descPage2.NextCursor) } + +func TestDeleteManyEvents(t *testing.T) { + logger := zap.NewNop().Sugar() + + for _, storeImpl := range testStores(logger) { + t.Run(storeImpl.Name, func(t *testing.T) { + store := storeImpl.Init(t) + testDeleteManyEvents(t, store) + }) + } +} + +func testDeleteManyEvents(t *testing.T, store storepkg.Store) { + events := make([]v1.Event, 200) + for i := range events { + events[i] = v1.Event{Kind: v1.EventKindLogLine, Timestamp: int64(i), Payload: "log line"} + } + + for _, eventBatch := range [][]v1.Event{events[:100], events[100:]} { + err := store.Update(func(txn storepkg.Transaction) error { + return txn.AppendEvents(eventBatch, "vms", "vm-with-many-events") + }) + require.NoError(t, err) + } + + err := store.Update(func(txn storepkg.Transaction) error { + return txn.DeleteEvents("vms", "vm-with-many-events") + }) + require.NoError(t, err) + + var remaining []v1.Event + err = store.View(func(txn storepkg.Transaction) error { + var err error + remaining, err = txn.ListEvents("vms", "vm-with-many-events") + + return err + }) + require.NoError(t, err) + require.Empty(t, remaining) +} diff --git a/internal/controller/store/store_impl_test.go b/internal/controller/store/store_impl_test.go new file mode 100644 index 0000000..3c002f7 --- /dev/null +++ b/internal/controller/store/store_impl_test.go @@ -0,0 +1,64 @@ +package store_test + +import ( + "fmt" + "os" + "strings" + "testing" + "time" + + storepkg "github.com/cirruslabs/orchard/internal/controller/store" + "github.com/cirruslabs/orchard/internal/controller/store/badger" + etcdstore "github.com/cirruslabs/orchard/internal/controller/store/etcd" + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +type storeImpl struct { + Name string + Init func(t *testing.T) storepkg.Store +} + +func testStores(logger *zap.SugaredLogger) []storeImpl { + storeImpls := []storeImpl{ + { + Name: "badger", + Init: func(t *testing.T) storepkg.Store { + store, err := badger.NewBadgerStore(t.TempDir(), true, logger) + require.NoError(t, err) + + return store + }, + }, + } + + etcdEndpointsRaw := os.Getenv("ORCHARD_TEST_ETCD_ENDPOINTS") + if etcdEndpointsRaw != "" { + storeImpls = append(storeImpls, storeImpl{ + Name: "etcd", + Init: func(t *testing.T) storepkg.Store { + prefix := fmt.Sprintf("/orchard-tests/%s-%d", strings.ReplaceAll(t.Name(), "/", "-"), time.Now().UnixNano()) + store, err := etcdstore.NewEtcdStore(splitEndpoints(etcdEndpointsRaw), prefix, logger) + require.NoError(t, err) + + return store + }, + }) + } + + return storeImpls +} + +func splitEndpoints(rawEndpoints string) []string { + var endpoints []string + for _, endpoint := range strings.Split(rawEndpoints, ",") { + endpoint = strings.TrimSpace(endpoint) + if endpoint == "" { + continue + } + + endpoints = append(endpoints, endpoint) + } + + return endpoints +} diff --git a/internal/controller/store/store_test.go b/internal/controller/store/store_test.go index 058f08d..ed370f4 100644 --- a/internal/controller/store/store_test.go +++ b/internal/controller/store/store_test.go @@ -7,7 +7,6 @@ import ( "time" storepkg "github.com/cirruslabs/orchard/internal/controller/store" - "github.com/cirruslabs/orchard/internal/controller/store/badger" "github.com/cirruslabs/orchard/pkg/resource/v1" "github.com/stretchr/testify/require" "go.uber.org/zap" @@ -152,26 +151,12 @@ func TestWatchVM(t *testing.T) { }, } - storeImpls := []struct { - Name string - Init func() (storepkg.Store, error) - }{ - { - Name: "badger", - Init: func() (storepkg.Store, error) { - return badger.NewBadgerStore(t.TempDir(), true, logger.Sugar()) - }, - }, - } - for _, testCase := range testCases { - for _, storeImpl := range storeImpls { + for _, storeImpl := range testStores(logger.Sugar()) { name := fmt.Sprintf("%s-%s", testCase.Name, storeImpl.Name) t.Run(name, func(t *testing.T) { - store, err := storeImpl.Init() - require.NoError(t, err) - + store := storeImpl.Init(t) testCase.Run(t, store) }) }