From 37d08e88068073eb19de92882adc466c7d6dfcad Mon Sep 17 00:00:00 2001
From: Chris Goller <goller@gmail.com>
Date: Tue, 19 Sep 2023 16:43:52 -0500
Subject: [PATCH] feat: enable run-time override of fetch/push concurrency

Signed-off-by: Chris Goller <goller@gmail.com>
---
 proto/depot/cloud/v3/machine.proto      | 1 +
 src/gen/ts/depot/cloud/v3/machine_pb.ts | 6 ++++++
 src/tasks/buildkit.ts                   | 4 ++++
 3 files changed, 11 insertions(+)

diff --git a/proto/depot/cloud/v3/machine.proto b/proto/depot/cloud/v3/machine.proto
index 8a0cd64..47fee38 100644
--- a/proto/depot/cloud/v3/machine.proto
+++ b/proto/depot/cloud/v3/machine.proto
@@ -77,6 +77,7 @@ message RegisterMachineResponse {
 
     // Disable fstrim on shutdown
     optional bool disable_fstrim = 14;
+    optional int32 resolver_concurrency = 15;
   }
 
   // Specifies sending buildkit profiling data to a remote endpoint.
diff --git a/src/gen/ts/depot/cloud/v3/machine_pb.ts b/src/gen/ts/depot/cloud/v3/machine_pb.ts
index 96151e8..4e93c4c 100644
--- a/src/gen/ts/depot/cloud/v3/machine_pb.ts
+++ b/src/gen/ts/depot/cloud/v3/machine_pb.ts
@@ -465,6 +465,11 @@ export class RegisterMachineResponse_BuildKitTask extends Message<RegisterMachin
    */
   disableFstrim?: boolean
 
+  /**
+   * @generated from field: optional int32 resolver_concurrency = 15;
+   */
+  resolverConcurrency?: number
+
   constructor(data?: PartialMessage<RegisterMachineResponse_BuildKitTask>) {
     super()
     proto3.util.initPartial(data, this)
@@ -486,6 +491,7 @@ export class RegisterMachineResponse_BuildKitTask extends Message<RegisterMachin
     {no: 12, name: 'disable_merge_to', kind: 'scalar', T: 8 /* ScalarType.BOOL */, opt: true},
     {no: 13, name: 'max_parallelism', kind: 'scalar', T: 5 /* ScalarType.INT32 */},
     {no: 14, name: 'disable_fstrim', kind: 'scalar', T: 8 /* ScalarType.BOOL */, opt: true},
+    {no: 15, name: 'resolver_concurrency', kind: 'scalar', T: 5 /* ScalarType.INT32 */, opt: true},
   ])
 
   static fromBinary(bytes: Uint8Array, options?: Partial<BinaryReadOptions>): RegisterMachineResponse_BuildKitTask {
diff --git a/src/tasks/buildkit.ts b/src/tasks/buildkit.ts
index e2bee7e..d5f0178 100644
--- a/src/tasks/buildkit.ts
+++ b/src/tasks/buildkit.ts
@@ -120,6 +120,10 @@ keepBytes = ${cacheSizeBytes}
     env.BUILDKIT_SCHEDULER_DEBUG = '1'
   }
 
+  if (task.resolverConcurrency)  {
+    env.DEPOT_RESOLVER_CONCURRENCY = task.resolverConcurrency.toString()
+  }
+
   async function runBuildKit() {
     try {
       await execa('/usr/bin/buildkitd', [], {stdio: 'inherit', signal, env})