Resolve issues with orphaned VMs when process restarts with in-flight jobs

examples/.env

@@ -22,7 +22,7 @@ GITHUB_APP_PRIVATE_KEY = ""
 # GITHUB_URL="https://github.com/<account-name>/<repo-name>" where <account-name> is your account name and <repo-name> is your repo name
 # For repositories under an organization, use the format:
 # GITHUB_URL="https://github.com/<org-name>" where <org-name> 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://<my-enterprise-github-url>/<org-name>" where <my-enterprise-github-url> is the url endpoint for your GHES instance,
 # and <org-name> 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

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)
 

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"
 

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

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

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())
+		})
+	})
 })

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 != "" {