Search blitz: updates for 3rd-party users (#64375)

The goal of this PR is to make search blitz usable for customers who
want to run it against their instance for continuous performance
analysis.

Specifically, this does two things:
1) Enables configuring the set of queries to run with the
`SEARCH_BLITZ_QUERY_FILE` env var
2) Packages the image in our standard format so we can publish it

Author: camdencheek

internal/cmd/search-blitz/BUILD.bazel

@@ -1,5 +1,7 @@
 load("//dev:go_defs.bzl", "go_test")
 load("@io_bazel_rules_go//go:def.bzl", "go_binary", "go_library")
+load("//dev:oci_defs.bzl", "image_repository", "oci_image", "oci_push", "oci_tarball", "pkg_tar")
+load("@container_structure_test//:defs.bzl", "container_structure_test")
 
 go_library(
     name = "search-blitz_lib",
@@ -50,3 +52,45 @@ go_test(
     embed = [":search-blitz_lib"],
     tags = [TAG_PLATFORM_SEARCH],
 )
+
+pkg_tar(
+    name = "tar_search_blitz",
+    srcs = [":search-blitz"],
+)
+
+oci_image(
+    name = "image",
+    base = "//wolfi-images/sourcegraph-base:base_image",
+    entrypoint = [
+        "/sbin/tini",
+        "--",
+        "/search-blitz",
+    ],
+    tars = [":tar_search_blitz"],
+    user = "sourcegraph",
+)
+
+oci_tarball(
+    name = "image_tarball",
+    image = ":image",
+    repo_tags = ["search-blitz:candidate"],
+)
+
+container_structure_test(
+    name = "image_test",
+    timeout = "short",
+    configs = ["image_test.yaml"],
+    driver = "docker",
+    image = ":image",
+    tags = [
+        "exclusive",
+        "requires-network",
+        TAG_PLATFORM_SEARCH,
+    ],
+)
+
+oci_push(
+    name = "candidate_push",
+    image = ":image",
+    repository = image_repository("search-blitz"),
+)

internal/cmd/search-blitz/README.md

@@ -1,7 +1,7 @@
 # Search-Blitz
 
 The purpose of Search-Blitz is to provide a baseline for our search performance.
-Search-Blitz calls the stream and GraphQL API of Sourcegraph.com for typical
+Search-Blitz calls the stream and GraphQL APIs for typical
 queries in regular intervals. Sourcegraph recognizes the Search-Blitz's
 `User-Agent` and sends metrics to Prometheus.
 
@@ -18,6 +18,8 @@ Add the query to [`queries.txt`](https://github.com/sourcegraph/sourcegraph/blob
 
 For attribution search add a `.txt` file to the attribution directory.
 
+A custom set of queries can be provided at runtime by pointing to a query file with the `SEARCH_BLITZ_QUERY_FILE` env var.
+
 ## How to deploy
 
 1. Merge your changes to _main_

internal/cmd/search-blitz/config.go

@@ -5,6 +5,7 @@ import (
 	"embed"
 	_ "embed"
 	"io/fs"
+	"os"
 	"strings"
 	"time"
 
@@ -26,7 +27,7 @@ type QueryConfig struct {
 	Snippet string
 	Name    string
 
-	// An unset interval defaults to 1m
+	// An unset interval defaults to 5m
 	Interval time.Duration
 
 	// An empty value for Protocols means "all"
@@ -43,7 +44,21 @@ const (
 	Stream
 )
 
-func loadQueries(env string) (_ *Config, err error) {
+// Loads queries from queryFile if set, otherwise falls back to
+// embedded set of queries.
+func loadConfig(queryFile, env string) (_ *Config, err error) {
+	if queryFile != "" {
+		queriesRaw, err := os.ReadFile(queryFile)
+		if err != nil {
+			return nil, err
+		}
+		queries, err := parseQueries(queriesRaw)
+		if err != nil {
+			return nil, err
+		}
+		return &Config{Queries: queries}, nil
+	}
+
 	if env == "" {
 		env = "cloud"
 	}
@@ -69,6 +84,16 @@ func loadQueries(env string) (_ *Config, err error) {
 	}
 	queries = append(queries, attributions...)
 
+	parsedQueries, err := parseQueries(queriesRaw)
+	if err != nil {
+		return nil, err
+	}
+	queries = append(queries, parsedQueries...)
+
+	return &Config{Queries: queries}, nil
+}
+
+func parseQueries(queriesRaw []byte) (queries []*QueryConfig, err error) {
 	var current QueryConfig
 	add := func() {
 		q := &QueryConfig{
@@ -97,10 +122,7 @@ func loadQueries(env string) (_ *Config, err error) {
 		}
 	}
 	add()
-
-	return &Config{
-		Queries: queries,
-	}, err
+	return queries, err
 }
 
 func loadAttributions() ([]*QueryConfig, error) {

internal/cmd/search-blitz/config_test.go

@@ -5,7 +5,7 @@ import "testing"
 func TestLoadQueries(t *testing.T) {
 	for _, env := range []string{"", "cloud", "dogfood"} {
 		t.Run(env, func(t *testing.T) {
-			c, err := loadQueries(env)
+			c, err := loadConfig("", env)
 			if err != nil {
 				t.Fatal(err)
 			}

internal/cmd/search-blitz/image_test.yaml

@@ -0,0 +1,9 @@
+schemaVersion: "2.0.0"
+
+commandTests:
+  - name: "not running as root"
+    command: "/usr/bin/id"
+    args:
+      - -u
+    excludedOutput: ["^0"]
+    exitCode: 0

internal/cmd/search-blitz/main.go

@@ -24,7 +24,7 @@ const (
 	envLogDir = "LOG_DIR"
 )
 
-func run(ctx context.Context, wg *sync.WaitGroup, env string) {
+func run(ctx context.Context, wg *sync.WaitGroup, config *Config) {
 	defer wg.Done()
 
 	bc, err := newClient()
@@ -37,11 +37,6 @@ func run(ctx context.Context, wg *sync.WaitGroup, env string) {
 		panic(err)
 	}
 
-	config, err := loadQueries(env)
-	if err != nil {
-		panic(err)
-	}
-
 	clientForProtocol := func(p Protocol) genericClient {
 		switch p {
 		case Batch:
@@ -196,10 +191,16 @@ func main() {
 	defer cleanup()
 
 	env := os.Getenv("SEARCH_BLITZ_ENV")
+	queryFile := os.Getenv("SEARCH_BLITZ_QUERY_FILE")
+
+	config, err := loadConfig(queryFile, env)
+	if err != nil {
+		panic(err)
+	}
 
 	wg := sync.WaitGroup{}
 	wg.Add(1)
-	go run(ctx, &wg, env)
+	go run(ctx, &wg, config)
 
 	wg.Add(1)
 	srv := startServer(&wg)