Add profile spans to BazelOutputService RPC calls

coeuvre · Apr 8, 2024 · bf16de3 · bf16de3
1 parent a0c17b0
commit bf16de3
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/src/main/java/com/google/devtools/build/lib/remote/BazelOutputService.java b/src/main/java/com/google/devtools/build/lib/remote/BazelOutputService.java
@@ -27,6 +27,7 @@
 import com.google.devtools.build.lib.actions.ExecException;
 import com.google.devtools.build.lib.actions.cache.OutputMetadataStore;
 import com.google.devtools.build.lib.events.EventHandler;
+import com.google.devtools.build.lib.profiler.Profiler;
 import com.google.devtools.build.lib.remote.BazelOutputServiceProto.BatchStatRequest;
 import com.google.devtools.build.lib.remote.BazelOutputServiceProto.BatchStatResponse;
 import com.google.devtools.build.lib.remote.BazelOutputServiceProto.CleanRequest;
@@ -250,7 +251,7 @@ private StartBuildResponse startBuild(StartBuildRequest request)
         () ->
             channel.withChannelBlocking(
                 channel -> {
-                  try {
+                  try (var sc = Profiler.instance().profile("BazelOutputService.StartBuild")) {
                     return BazelOutputServiceGrpc.newBlockingStub(channel).startBuild(request);
                   } catch (StatusRuntimeException e) {
                     throw new IOException(e);
@@ -295,7 +296,7 @@ private StageArtifactsResponse stageArtifacts(StageArtifactsRequest request)
         () ->
             channel.withChannelBlocking(
                 channel -> {
-                  try {
+                  try (var sc = Profiler.instance().profile("BazelOutputService.StageArtifacts")) {
                     return BazelOutputServiceGrpc.newBlockingStub(channel).stageArtifacts(request);
                   } catch (StatusRuntimeException e) {
                     throw new IOException(e);
@@ -335,7 +336,7 @@ private FinalizeBuildResponse finalizeBuild(FinalizeBuildRequest request)
         () ->
             channel.withChannelBlocking(
                 channel -> {
-                  try {
+                  try (var sc = Profiler.instance().profile("BazelOutputService.FinalizeBuild")) {
                     return BazelOutputServiceGrpc.newBlockingStub(channel).finalizeBuild(request);
                   } catch (StatusRuntimeException e) {
                     throw new IOException(e);
@@ -376,7 +377,8 @@ private FinalizeArtifactsResponse finalizeArtifacts(FinalizeArtifactsRequest req
         () ->
             channel.withChannelBlocking(
                 channel -> {
-                  try {
+                  try (var sc =
+                      Profiler.instance().profile("BazelOutputService.FinalizeArtifacts")) {
                     return BazelOutputServiceGrpc.newBlockingStub(channel)
                         .finalizeArtifacts(request);
                   } catch (StatusRuntimeException e) {
@@ -643,7 +645,7 @@ private CleanResponse clean(CleanRequest request) throws IOException, Interrupte
         () ->
             channel.withChannelBlocking(
                 channel -> {
-                  try {
+                  try (var sc = Profiler.instance().profile("BazelOutputService.Clean")) {
                     return BazelOutputServiceGrpc.newBlockingStub(channel).clean(request);
                   } catch (StatusRuntimeException e) {
                     throw new IOException(e);
@@ -657,7 +659,7 @@ private BatchStatResponse batchStat(BatchStatRequest request)
         () ->
             channel.withChannelBlocking(
                 channel -> {
-                  try {
+                  try (var sc = Profiler.instance().profile("BazelOutputService.BatchStat")) {
                     return BazelOutputServiceGrpc.newBlockingStub(channel).batchStat(request);
                   } catch (StatusRuntimeException e) {
                     throw new IOException(e);