From 9dc8c342804f5b88d49a3d29c416a78df6579240 Mon Sep 17 00:00:00 2001
From: bohendo <bo.henderson@trailofbits.com>
Date: Wed, 13 Mar 2024 11:24:22 -0400
Subject: [PATCH] improve performance of a couple loops

---
 cmd/cloudexec/clean.go |  6 +++---
 pkg/state/state.go     | 17 ++++++++++-------
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/cmd/cloudexec/clean.go b/cmd/cloudexec/clean.go
index b000d42..b038344 100644
--- a/cmd/cloudexec/clean.go
+++ b/cmd/cloudexec/clean.go
@@ -15,13 +15,13 @@ func ConfirmDeleteDroplets(config config.Config, dropletName string, instanceToJ
 	if err != nil {
 		return confirmedToDelete, fmt.Errorf("Failed to get droplets by name: %w", err)
 	}
+	if instanceToJobs == nil {
+		return confirmedToDelete, fmt.Errorf("Given instanceToJobs argument must not be nil")
+	}
 	if len(instances) > 0 {
 		fmt.Printf("Existing %s instance(s) found:\n", dropletName)
 		for _, instance := range instances {
 			// get a pretty string describing the jobs associated with this instance
-			if instanceToJobs == nil {
-				return confirmedToDelete, fmt.Errorf("Given instanceToJobs argument must not be nil")
-			}
 			jobs := instanceToJobs[int64(instance.ID)]
 			var prettyJobs string
 			if len(jobs) == 0 {
diff --git a/pkg/state/state.go b/pkg/state/state.go
index 34d76fc..4c3f876 100644
--- a/pkg/state/state.go
+++ b/pkg/state/state.go
@@ -178,16 +178,19 @@ func (s *State) DeleteJob(jobID int64) {
 	}
 }
 
-func (s *State) CancelRunningJobs(config config.Config, bucketName string, toCancel []int64) error {
+func (s *State) CancelRunningJobs(config config.Config, bucketName string, toCancelSlice []int64) error {
+	// Gather list of ids to mark as cancelled
+	toCancel := make(map[int64]bool)
+	for _, id := range toCancelSlice {
+		toCancel[id] = true
+	}
+
 	// Mark any running jobs as cancelled
 	for i, job := range s.Jobs {
 		if job.Status == Running || job.Status == Provisioning {
-			for _, id := range toCancel {
-				if id == job.ID {
-					fmt.Printf("Setting status of job %d to 'Cancelled'\n", job.ID)
-					s.Jobs[i].Status = Cancelled
-					break
-				}
+			if _, exists := toCancel[job.ID]; exists {
+				fmt.Printf("Setting status of job %d to 'Cancelled'\n", job.ID)
+				s.Jobs[i].Status = Cancelled
 			}
 		}
 	}