Add script execs and software installs stats to osquery-perf (#26239)

2025-02-11 12:46:53 -05:00
parent 15bbac88ed
commit 240f55b9e8
8 changed files with 258 additions and 3 deletions
@@ -0,0 +1,5 @@
+# Load testing of the unified queue story
+
+This is the Go program used to run load tests for the [unified queue story](https://github.com/fleetdm/fleet/issues/22866).
+
+It expects some software to be available for install on both macOS and Windows (including VPP apps for macOS), and some scripts too, and it enqueues installs and script execution requests on every host in the Fleet deployment for an hour.
@@ -0,0 +1,190 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"log"
+	"math/rand/v2"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/fleetdm/fleet/v4/server/fleet"
+	"github.com/fleetdm/fleet/v4/server/service"
+)
+
+func printf(format string, a ...any) {
+	fmt.Printf(time.Now().UTC().Format("2006-01-02T15:04:05Z")+": "+format, a...)
+}
+
+func main() {
+	fleetURL := flag.String("fleet_url", "", "URL (with protocol and port of Fleet server)")
+	apiToken := flag.String("api_token", "", "API authentication token to use on API calls")
+	debug := flag.Bool("debug", false, "Debug mode")
+
+	flag.Parse()
+
+	if *fleetURL == "" {
+		log.Fatal("missing fleet_url argument")
+	}
+	if *apiToken == "" {
+		log.Fatal("missing api_token argument")
+	}
+	var clientOpts []service.ClientOption
+	if *debug {
+		clientOpts = append(clientOpts, service.EnableClientDebug())
+	}
+	apiClient, err := service.NewClient(*fleetURL, true, "", "", clientOpts...)
+	if err != nil {
+		log.Fatal(err)
+	}
+	apiClient.SetToken(*apiToken)
+
+	printf("Fetching hosts...\n")
+	records, err := apiClient.GetHostsReport("id", "hostname", "platform")
+	if err != nil {
+		log.Fatal(err)
+	}
+	type smallHost struct {
+		ID       uint
+		Hostname string
+		Platform string
+	}
+	var (
+		macOSHosts   []smallHost
+		windowsHosts []smallHost
+		linuxHosts   []smallHost
+	)
+	for i, record := range records {
+		if i == 0 {
+			continue
+		}
+		hostID, _ := strconv.Atoi(record[0])
+		hostname := record[1]
+		platform := fleet.PlatformFromHost(record[2])
+		switch platform {
+		case "linux":
+			linuxHosts = append(linuxHosts, smallHost{ID: uint(hostID), Hostname: hostname, Platform: platform}) // nolint:gosec
+		case "darwin":
+			macOSHosts = append(macOSHosts, smallHost{ID: uint(hostID), Hostname: hostname, Platform: platform}) // nolint:gosec
+		case "windows":
+			windowsHosts = append(windowsHosts, smallHost{ID: uint(hostID), Hostname: hostname, Platform: platform}) // nolint:gosec
+		}
+	}
+	printf("Got linux=%d, windows=%d, macOS=%d\n", len(linuxHosts), len(windowsHosts), len(macOSHosts))
+
+	titles, err := apiClient.ListSoftwareTitles("per_page=1000&team_id=0&available_for_install=1")
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	var (
+		macOSSoftware   []fleet.SoftwareTitleListResult
+		windowsSoftware []fleet.SoftwareTitleListResult
+	)
+	for _, title := range titles {
+		if title.AppStoreApp != nil {
+			macOSSoftware = append(macOSSoftware, title)
+		} else if title.SoftwarePackage != nil {
+			if ext := filepath.Ext(title.SoftwarePackage.Name); ext == ".exe" || ext == ".msi" {
+				windowsSoftware = append(windowsSoftware, title)
+			} else {
+				macOSSoftware = append(macOSSoftware, title)
+			}
+		}
+	}
+	printf("Got software titles windows=%d, macOS=%d\n", len(windowsSoftware), len(macOSSoftware))
+
+	scripts, err := apiClient.ListScripts("per_page=1000&team_id=0")
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	var (
+		macOSScripts   []string
+		windowsScripts []string
+	)
+	for _, script := range scripts {
+		if strings.HasSuffix(script.Name, ".sh") {
+			macOSScripts = append(macOSScripts, script.Name)
+		} else if strings.HasSuffix(script.Name, ".ps1") {
+			windowsScripts = append(windowsScripts, script.Name)
+		}
+	}
+	printf("Got scripts windows=%d, macOS=%d\n", len(windowsScripts), len(macOSScripts))
+
+	var queuedScripts, queuedInstalls, hostsTargeted, errors int
+	targetedHosts := append(macOSHosts, windowsHosts...) // nolint:gocritic
+	rand.Shuffle(len(targetedHosts), func(i, j int) {
+		targetedHosts[i], targetedHosts[j] = targetedHosts[j], targetedHosts[i]
+	})
+
+	tick := time.Tick(300 * time.Millisecond)
+	for i, host := range targetedHosts {
+		<-tick
+
+		if hostsTargeted > 0 && hostsTargeted%500 == 0 {
+			printf("In progress: queued scripts=%d, queued installs=%d, hosts targeted=%d, errors=%d\n", queuedScripts, queuedInstalls, hostsTargeted, errors)
+		}
+
+		switch host.Platform {
+		case "darwin":
+			hostsTargeted++
+
+			// enqueue a software install and a couple scripts
+			_, err := apiClient.RunHostScriptAsync(host.ID, nil, macOSScripts[i%len(macOSScripts)], 0)
+			if err != nil {
+				printf("Failed to run script on host %v (%v): %v\n", host.Hostname, host.Platform, err)
+				errors++
+				continue
+			}
+			queuedScripts++
+			_, err = apiClient.RunHostScriptAsync(host.ID, nil, macOSScripts[(i+1)%len(macOSScripts)], 0)
+			if err != nil {
+				printf("Failed to run script on host %v (%v): %v\n", host.Hostname, host.Platform, err)
+				errors++
+				continue
+			}
+			queuedScripts++
+
+			err = apiClient.InstallSoftware(host.ID, macOSSoftware[i%len(macOSSoftware)].ID)
+			if err != nil {
+				printf("Failed to install software on host %v (%v): %v\n", host.Hostname, host.Platform, err)
+				errors++
+				continue
+			}
+			queuedInstalls++
+
+		case "windows":
+			hostsTargeted++
+
+			// enqueue a couple software installs and a script
+			err = apiClient.InstallSoftware(host.ID, windowsSoftware[i%len(windowsSoftware)].ID)
+			if err != nil {
+				printf("Failed to install software on host %v (%v): %v\n", host.Hostname, host.Platform, err)
+				errors++
+				continue
+			}
+			queuedInstalls++
+
+			err = apiClient.InstallSoftware(host.ID, windowsSoftware[(i+1)%len(windowsSoftware)].ID)
+			if err != nil {
+				printf("Failed to install software on host %v (%v): %v\n", host.Hostname, host.Platform, err)
+				errors++
+				continue
+			}
+			queuedInstalls++
+
+			_, err = apiClient.RunHostScriptAsync(host.ID, nil, windowsScripts[i%len(windowsScripts)], 0)
+			if err != nil {
+				printf("Failed to run script on host %v (%v): %v\n", host.Hostname, host.Platform, err)
+				errors++
+				continue
+			}
+			queuedScripts++
+		}
+	}
+
+	printf("Done: queued scripts=%d, queued installs=%d, hosts targeted=%d, errors=%d\n", queuedScripts, queuedInstalls, hostsTargeted, errors)
+}