add configurable retry and batch size options for benchmarking (#20)

- add query-retries flag to control query retry behavior (defaults to 0 for client defaults)
- add namespace-setup-batch-size flag for configuring batch size during setup (default 250k)
- add namespace-setup-concurrency-max flag to cap total concurrent goroutines across namespaces
- make query retries conditional: only apply WithMaxRetries when query-retries is non-zero
- use new batch size flag in upsert operations instead of hardcoded value

Author: utaal

flags.go

@@ -91,7 +91,19 @@ var logNormalSigma = flag.Float64(
 var namespaceSetupConcurrency = flag.Int(
 	"namespace-setup-concurrency",
 	4,
-	"the number of concurrent goroutines to use for namespace setup",
+	"the number of concurrent goroutines to use for namespace setup (concurrency per namespace)",
+)
+
+var namespaceSetupConcurrencyMax = flag.Int(
+	"namespace-setup-concurrency-max",
+	64,
+	"maximum number of concurrent goroutines to use for namespace setup (total across all namespaces)",
+)
+
+var namespaceSetupBatchSize = flag.Int(
+	"namespace-setup-batch-size",
+	250_000,
+	"the number of documents to process in each batch during namespace setup",
 )
 
 // Benchmark settings
@@ -174,6 +186,12 @@ var benchmarkDuration = flag.Duration(
 	"duration of the benchmark. if 0, will run indefinitely",
 )
 
+var benchmarkQueryRetries = flag.Int(
+	"query-retries",
+	0,
+	"number of times to retry failed queries (default 0 = default tpuf client retries)",
+)
+
 var outputDir = flag.String(
 	"output-dir",
 	defaultOutputDir(),

namespace.go

@@ -227,7 +227,13 @@ func (n *Namespace) Query(ctx context.Context) (*turbopuffer.QueryPerformance, t
 
 	url := fmt.Sprintf("/v2/namespaces/%s/query", n.ID())
 	var response turbopuffer.NamespaceQueryResponse
-	if err := n.client.Post(ctx, url, buf.Bytes(), &response); err != nil {
+
+	var opts []option.RequestOption
+	if *benchmarkQueryRetries != 0 {
+		opts = append(opts, option.WithMaxRetries(*benchmarkQueryRetries))
+	}
+
+	if err := n.client.Post(ctx, url, buf.Bytes(), &response, opts...); err != nil {
 		var apiErr *turbopuffer.Error
 		if errors.As(err, &apiErr) && apiErr.StatusCode == http.StatusNotFound {
 			return nil, 0, nil

upsert.go

@@ -59,7 +59,7 @@ func UpsertDocumentsToNamespaces(
 		totalUpserts += int64(size)
 	}
 
-	concurrentRequests := min(max(1, (*namespaceSetupConcurrency)*len(namespaces)), 64)
+	concurrentRequests := min(max(1, (*namespaceSetupConcurrency)*len(namespaces)), *namespaceSetupConcurrencyMax)
 	log.Printf("upserting documents with %d concurrent batches\n", concurrentRequests)
 	pb := progressbar.Default(totalUpserts, "upserting documents")
 
@@ -122,7 +122,7 @@ func makeProgressOn(
 
 	var (
 		largest = upserts[len(upserts)-1].Pending
-		batch   = min(largest, 250_000)
+		batch   = min(largest, *namespaceSetupBatchSize)
 	)
 	if batch == 0 {
 		return nil, errors.New("batch size is zero")