diff --git a/examples/.env b/examples/.env index e0c2e81..053802f 100644 --- a/examples/.env +++ b/examples/.env @@ -22,7 +22,7 @@ GITHUB_APP_PRIVATE_KEY = "" # GITHUB_URL="https://github.com//" where is your account name and is your repo name # For repositories under an organization, use the format: # GITHUB_URL="https://github.com/" where is the name of your organisation -# For Github Enterprise Self-Hosted instances (GHES) repositoroes under an organiztion, use the format: +# For Github Enterprise Self-Hosted instances (GHES) repositoroes under an organiztion, use the format: # GITHUB_URL="https:///" where is the url endpoint for your GHES instance, # and is the name of your organization in the GHES instance GITHUB_URL="https://github.com/macstadium" @@ -78,6 +78,11 @@ ORKA_VM_METADATA="key1=value1,key2=value2" # If not provided, it defaults to 300 seconds. VM_TRACKER_INTERVAL="300s" +# [Optional] VM_TRACKER_SEED_ON_START specifies whether the VM tracker should perform an initial seeding of existing VMs on startup. +# If set to true, the VM tracker will check for any existing VMs known by Orka and add them to its tracking list. +# This can help ensure that any VMs that were created before the process started are properly tracked. +VM_TRACKER_SEED_ON_START=true + # Prometheus metrics (optional) ENABLE_METRICS=true METRICS_ADDR=:8080 diff --git a/main.go b/main.go index e629edb..55c7bc4 100644 --- a/main.go +++ b/main.go @@ -148,7 +148,7 @@ func run(ctx context.Context, actionsClient *actions.ActionsClient, orkaClient * runnerProvisioner := provisioner.NewRunnerProvisioner(runnerScaleSet, actionsClient, orkaClient, envData) vmTracker := runners.NewVMTracker(orkaClient, actionsClient, logger) - go vmTracker.Start(ctx, envData.VMTrackerInterval) + go vmTracker.Start(ctx, envData.VMTrackerInterval, runnerScaleSet.Name, envData.VMTrackerSeedOnStart) runnerMessageProcessor := runners.NewRunnerMessageProcessor(ctx, runnerManager, runnerProvisioner, vmTracker, runnerScaleSet) diff --git a/pkg/env/constants.go b/pkg/env/constants.go index 7b3b8b5..d19881a 100644 --- a/pkg/env/constants.go +++ b/pkg/env/constants.go @@ -27,7 +27,8 @@ const ( RunnerDeregistrationTimeoutEnvName = "RUNNER_DEREGISTRATION_TIMEOUT" RunnerDeregistrationPollIntervalEnvName = "RUNNER_DEREGISTRATION_POLL_INTERVAL" - VMTrackerIntervalEnvName = "VM_TRACKER_INTERVAL" + VMTrackerIntervalEnvName = "VM_TRACKER_INTERVAL" + VMTrackerSeedOnStartEnvName = "VM_TRACKER_SEED_ON_START" LogLevelEnvName = "LOG_LEVEL" diff --git a/pkg/env/env.go b/pkg/env/env.go index e77aefb..8af19fe 100644 --- a/pkg/env/env.go +++ b/pkg/env/env.go @@ -48,7 +48,8 @@ type Data struct { RunnerDeregistrationTimeout time.Duration RunnerDeregistrationPollInterval time.Duration - VMTrackerInterval time.Duration + VMTrackerInterval time.Duration + VMTrackerSeedOnStart bool LogLevel string @@ -85,7 +86,8 @@ func ParseEnv() *Data { RunnerDeregistrationTimeout: getDurationEnv(RunnerDeregistrationTimeoutEnvName, 30*time.Second), RunnerDeregistrationPollInterval: getDurationEnv(RunnerDeregistrationPollIntervalEnvName, 2*time.Second), - VMTrackerInterval: getDurationEnv(VMTrackerIntervalEnvName, 300*time.Second), + VMTrackerInterval: getDurationEnv(VMTrackerIntervalEnvName, 300*time.Second), + VMTrackerSeedOnStart: getBoolEnv(VMTrackerSeedOnStartEnvName, true), LogLevel: getEnvWithDefault(LogLevelEnvName, logging.LogLevelInfo), diff --git a/pkg/github/runners/vm_tracker.go b/pkg/github/runners/vm_tracker.go index 7e95717..7aae37b 100644 --- a/pkg/github/runners/vm_tracker.go +++ b/pkg/github/runners/vm_tracker.go @@ -43,7 +43,32 @@ func (tracker *VMTracker) Untrack(vmName string) { delete(tracker.trackedVMs, vmName) } -func (tracker *VMTracker) Start(ctx context.Context, interval time.Duration) { +func (tracker *VMTracker) seedFromOrka(ctx context.Context, runnerScaleSetName string) { + tracker.logger.Infof("Seeding VM tracker from Orka for scale set prefix %q", runnerScaleSetName) + + vms, err := tracker.orkaClient.ListVMs(ctx) + if err != nil { + tracker.logger.Errorf("Failed to list VMs during seeding: %v", err) + return + } + + var matched int + for _, vm := range vms { + if !strings.HasPrefix(vm.Name, runnerScaleSetName) { + continue + } + matched++ + tracker.Track(vm.Name) + } + + tracker.logger.Infof("Seeded %d VMs matching prefix %q out of %d total", matched, runnerScaleSetName, len(vms)) +} + +func (tracker *VMTracker) Start(ctx context.Context, interval time.Duration, runnerScaleSetName string, seedOnStart bool) { + if seedOnStart { + tracker.seedFromOrka(ctx, runnerScaleSetName) + } + ticker := time.NewTicker(interval) defer ticker.Stop() diff --git a/pkg/github/runners/vm_tracker_test.go b/pkg/github/runners/vm_tracker_test.go index 8f837e1..af70a86 100644 --- a/pkg/github/runners/vm_tracker_test.go +++ b/pkg/github/runners/vm_tracker_test.go @@ -19,10 +19,18 @@ func TestVMTracker(t *testing.T) { } type MockOrkaClient struct { + ListVMsFunc func(ctx context.Context) ([]orka.OrkaVMDeployResponseModel, error) DeleteVMFunc func(ctx context.Context, name string) error DeployVMFunc func(ctx context.Context, namePrefix, vmConfig string) (*orka.OrkaVMDeployResponseModel, error) } +func (m *MockOrkaClient) ListVMs(ctx context.Context) ([]orka.OrkaVMDeployResponseModel, error) { + if m.ListVMsFunc != nil { + return m.ListVMsFunc(ctx) + } + return nil, nil +} + func (m *MockOrkaClient) DeleteVM(ctx context.Context, name string) error { if m.DeleteVMFunc != nil { return m.DeleteVMFunc(ctx, name) @@ -209,4 +217,62 @@ var _ = Describe("VMTracker", func() { }) }) }) + + Describe("seedFromOrka", func() { + const scaleSetName = "my-runner" + + It("should track VMs with matching prefix", func() { + mockOrka.ListVMsFunc = func(ctx context.Context) ([]orka.OrkaVMDeployResponseModel, error) { + return []orka.OrkaVMDeployResponseModel{ + {Name: "my-runner-abc12", Status: orka.VMRunning}, + {Name: "my-runner-xyz99", Status: orka.VMRunning}, + }, nil + } + + tracker.seedFromOrka(ctx, scaleSetName) + + Expect(tracker.trackedVMs).To(HaveLen(2)) + Expect(tracker.trackedVMs).To(HaveKey("my-runner-abc12")) + Expect(tracker.trackedVMs).To(HaveKey("my-runner-xyz99")) + }) + + It("should ignore VMs that do not match the scale set name prefix", func() { + mockOrka.ListVMsFunc = func(ctx context.Context) ([]orka.OrkaVMDeployResponseModel, error) { + return []orka.OrkaVMDeployResponseModel{ + {Name: "other-service-vm1", Status: orka.VMRunning}, + }, nil + } + + tracker.seedFromOrka(ctx, scaleSetName) + + Expect(tracker.trackedVMs).To(BeEmpty()) + }) + + It("should only track VMs with matching prefix in a mixed list", func() { + mockOrka.ListVMsFunc = func(ctx context.Context) ([]orka.OrkaVMDeployResponseModel, error) { + return []orka.OrkaVMDeployResponseModel{ + {Name: "my-runner-abc12", Status: orka.VMRunning}, + {Name: "unrelated-vm", Status: orka.VMRunning}, + {Name: "my-runner-xyz99", Status: orka.VMPending}, + }, nil + } + + tracker.seedFromOrka(ctx, scaleSetName) + + Expect(tracker.trackedVMs).To(HaveLen(2)) + Expect(tracker.trackedVMs).To(HaveKey("my-runner-abc12")) + Expect(tracker.trackedVMs).To(HaveKey("my-runner-xyz99")) + Expect(tracker.trackedVMs).NotTo(HaveKey("unrelated-vm")) + }) + + It("should handle ListVMs errors gracefully", func() { + mockOrka.ListVMsFunc = func(ctx context.Context) ([]orka.OrkaVMDeployResponseModel, error) { + return nil, errors.New("connection refused") + } + + tracker.seedFromOrka(ctx, scaleSetName) + + Expect(tracker.trackedVMs).To(BeEmpty()) + }) + }) }) diff --git a/pkg/orka/client.go b/pkg/orka/client.go index 77d8de9..d6e0d9e 100644 --- a/pkg/orka/client.go +++ b/pkg/orka/client.go @@ -13,6 +13,7 @@ import ( ) type OrkaService interface { + ListVMs(ctx context.Context) ([]OrkaVMDeployResponseModel, error) DeployVM(ctx context.Context, namePrefix, vmConfig string) (*OrkaVMDeployResponseModel, error) DeleteVM(ctx context.Context, name string) error } @@ -21,6 +22,15 @@ type OrkaClient struct { envData *env.Data } +func (client *OrkaClient) ListVMs(ctx context.Context) ([]OrkaVMDeployResponseModel, error) { + res, err := exec.ExecJSONCommand[[]OrkaVMDeployResponseModel]("orka3", []string{"vm", "list", "--namespace", client.envData.OrkaNamespace, "-o", "json"}) + if err != nil { + return nil, err + } + + return *res, nil +} + func (client *OrkaClient) DeployVM(ctx context.Context, namePrefix, vmConfig string) (*OrkaVMDeployResponseModel, error) { args := []string{"vm", "deploy", namePrefix, "--config", vmConfig, "--generate-name", "-o", "json", "--namespace", client.envData.OrkaNamespace} if client.envData.OrkaVMMetadata != "" {