public_git
/
orchard
mirror of https://github.com/cirruslabs/orchard.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge 4945127f76 into 1c241832f5

This commit is contained in:
Jonathan Haas 2026-08-14 12:36:43 -07:00 committed by GitHub
parent 1c241832f5 4945127f76
commit 7767bf292f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
13 changed files with 816 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
go.mod
View File

@ -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

11
go.sum
View File

@ -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=

35
internal/command/controller/run.go
View File

@ -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,

19
internal/controller/controller.go
View File

@ -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()

15
internal/controller/option.go
View File

@ -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

156
internal/controller/store/etcd/etcd_events.go Normal file
View File

@ -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)
}

123
internal/controller/store/etcd/etcd_generic.go Normal file
View File

@ -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
}

82
internal/controller/store/etcd/etcd_resources.go Normal file
View File

@ -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)
}

168
internal/controller/store/etcd/etcd_store.go Normal file
View File

@ -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)
}

88
internal/controller/store/etcd/etcd_watch.go Normal file
View File

@ -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
}

54
internal/controller/store/events_test.go
View File

@ -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)
}

64
internal/controller/store/store_impl_test.go Normal file
View File

@ -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
}

19
internal/controller/store/store_test.go
View File

@ -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)
})
}