From b99366e6fb78e98e02942ea7585853010f881aeb Mon Sep 17 00:00:00 2001
From: bufdev <4228796+bufdev@users.noreply.github.com>
Date: Wed, 31 Jan 2024 17:54:26 -0500
Subject: [PATCH] Add hidden --parallelism flag (#2745)

---
 private/pkg/app/appflag/builder.go | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/private/pkg/app/appflag/builder.go b/private/pkg/app/appflag/builder.go
index 5baabae2c4..749651aba2 100644
--- a/private/pkg/app/appflag/builder.go
+++ b/private/pkg/app/appflag/builder.go
@@ -24,6 +24,7 @@ import (
 	"github.com/bufbuild/buf/private/pkg/app/applog"
 	"github.com/bufbuild/buf/private/pkg/app/appverbose"
 	"github.com/bufbuild/buf/private/pkg/observabilityzap"
+	"github.com/bufbuild/buf/private/pkg/thread"
 	"github.com/pkg/profile"
 	"github.com/spf13/pflag"
 	"go.uber.org/multierr"
@@ -44,6 +45,8 @@ type builder struct {
 	profileType       string
 	profileAllowError bool
 
+	parallelism int
+
 	timeout time.Duration
 
 	defaultTimeout time.Duration
@@ -83,6 +86,8 @@ func (b *builder) BindRoot(flagSet *pflag.FlagSet) {
 	// We do not officially support this flag, this is for testing, where we need warnings turned off.
 	flagSet.BoolVar(&b.noWarn, "no-warn", false, "Turn off warn logging")
 	_ = flagSet.MarkHidden("no-warn")
+	flagSet.IntVar(&b.parallelism, "parallelism", 0, "Manually control the parallelism")
+	_ = flagSet.MarkHidden("parallelism")
 }
 
 func (b *builder) NewRunFunc(
@@ -120,6 +125,10 @@ func (b *builder) run(
 		return err
 	}
 
+	if b.parallelism > 0 {
+		thread.SetParallelism(b.parallelism)
+	}
+
 	var cancel context.CancelFunc
 	if !b.profile && b.timeout != 0 {
 		ctx, cancel = context.WithTimeout(ctx, b.timeout)