Merge pull request #106 from depot/profiling

depot · Apr 21, 2023 · bb18a57 · bb18a57
2 parents 0dd8d4b + 17d5bd6
commit bb18a57
Show file tree

Hide file tree

Showing 4 changed files with 350 additions and 259 deletions.
diff --git a/pkg/helpers/build.go b/pkg/helpers/build.go
@@ -26,6 +26,7 @@ func BeginBuild(ctx context.Context, project string, token string) (build Build,
 
 	build.Token = token
 	build.ID = os.Getenv("DEPOT_BUILD_ID")
+	profilerToken := ""
 	if build.ID == "" {
 		req := cliv1.CreateBuildRequest{ProjectId: project}
 		var b *connect.Response[cliv1.CreateBuildResponse]
@@ -36,13 +37,18 @@ func BeginBuild(ctx context.Context, project string, token string) (build Build,
 		build.ID = b.Msg.BuildId
 		build.Token = b.Msg.BuildToken
 
+		if b.Msg.Profiler != nil {
+			profilerToken = b.Msg.Profiler.Token
+		}
+
 		build.UseLocalRegistry = b.Msg.GetRegistry() != nil && b.Msg.GetRegistry().CanUseLocalRegistry
 		if os.Getenv("DEPOT_USE_LOCAL_REGISTRY") != "" {
 			build.UseLocalRegistry = true
 		}
 	}
 
-	profiler.StartProfiler(build.ID)
+	profiler.StartProfiler(build.ID, profilerToken)
+
 	build.Finish = func(buildErr error) {
 		req := cliv1.FinishBuildRequest{BuildId: build.ID}
 		req.Result = &cliv1.FinishBuildRequest_Success{Success: &cliv1.FinishBuildRequest_BuildSuccess{}}

diff --git a/pkg/profiler/profiler.go b/pkg/profiler/profiler.go
@@ -9,8 +9,10 @@ import (
 )
 
 // StartProfiler starts a profiler if DEPOT_ENABLE_DEBUG_PROFILING is set.
-func StartProfiler(buildID string) {
-	profileToken := os.Getenv("DEPOT_ENABLE_DEBUG_PROFILING")
+func StartProfiler(buildID, profileToken string) {
+	if profileToken == "" {
+		profileToken = os.Getenv("DEPOT_ENABLE_DEBUG_PROFILING")
+	}
 	if profileToken != "" {
 		runtime.SetMutexProfileFraction(5)
 		runtime.SetBlockProfileRate(10000)