bep: introduce BuildEventArtifactUploaderFactory

There can be multiple BuildEventTransports active at
the same time and we need to ensure that each transport
gets its own BuildEventArtifactUploader as these transports
might have different lifecycles.

We do that by introducing another level of indirection via
the BuildEventArtifactUploaderFactory. BlazeModules now
register a factory object instead of an uploader.

In addition, the BuildEventArtifactUploader gets a shutdown()
method that allows to free any resources associated with it.

PiperOrigin-RevId: 203752092

2 files changed, 29 insertions, 16 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceModule.java b/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceModule.java
index 115808a66a..b12f398f3f 100644
--- a/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceModule.java
+++ b/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceModule.java
@@ -27,7 +27,7 @@ import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.authandtls.AuthAndTLSOptions;
 import com.google.devtools.build.lib.buildeventservice.client.BuildEventServiceClient;
 import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploader;
-import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploaderMap;
+import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploaderFactoryMap;
 import com.google.devtools.build.lib.buildeventstream.BuildEventProtocolOptions;
 import com.google.devtools.build.lib.buildeventstream.BuildEventTransport;
 import com.google.devtools.build.lib.buildeventstream.transports.BuildEventStreamOptions;
@@ -88,7 +88,7 @@ public abstract class BuildEventServiceModule<T extends BuildEventServiceOptions
             commandEnvironment.getReporter(),
             commandEnvironment.getBlazeModuleEnvironment(),
             commandEnvironment.getRuntime().getClock(),
-            commandEnvironment.getRuntime().getBuildEventArtifactUploaders(),
+            commandEnvironment.getRuntime().getBuildEventArtifactUploaderFactoryMap(),
             commandEnvironment.getReporter(),
             commandEnvironment.getBuildRequestId().toString(),
             commandEnvironment.getCommandId().toString(),
@@ -139,12 +139,12 @@ public abstract class BuildEventServiceModule<T extends BuildEventServiceOptions
       EventHandler commandLineReporter,
       ModuleEnvironment moduleEnvironment,
       Clock clock,
-      BuildEventArtifactUploaderMap artifactUploaders,
+      BuildEventArtifactUploaderFactoryMap buildEventArtifactUploaderFactoryMap,
       Reporter reporter,
       String buildRequestId,
       String invocationId,
       String commandName) {
-    Preconditions.checkNotNull(artifactUploaders);
+    Preconditions.checkNotNull(buildEventArtifactUploaderFactoryMap);
 
     try {
       T besOptions =
@@ -152,14 +152,17 @@ public abstract class BuildEventServiceModule<T extends BuildEventServiceOptions
               optionsProvider.getOptions(optionsClass()),
               "Could not get BuildEventServiceOptions.");
       AuthAndTLSOptions authTlsOptions =
-          checkNotNull(optionsProvider.getOptions(AuthAndTLSOptions.class),
+          checkNotNull(
+              optionsProvider.getOptions(AuthAndTLSOptions.class),
               "Could not get AuthAndTLSOptions.");
       BuildEventStreamOptions bepOptions =
-          checkNotNull(optionsProvider.getOptions(BuildEventStreamOptions.class),
-          "Could not get BuildEventStreamOptions.");
+          checkNotNull(
+              optionsProvider.getOptions(BuildEventStreamOptions.class),
+              "Could not get BuildEventStreamOptions.");
       BuildEventProtocolOptions protocolOptions =
-          checkNotNull(optionsProvider.getOptions(BuildEventProtocolOptions.class),
-          "Could not get BuildEventProtocolOptions.");
+          checkNotNull(
+              optionsProvider.getOptions(BuildEventProtocolOptions.class),
+              "Could not get BuildEventProtocolOptions.");
 
       BuildEventTransport besTransport = null;
       try {
@@ -173,7 +176,7 @@ public abstract class BuildEventServiceModule<T extends BuildEventServiceOptions
                 moduleEnvironment,
                 clock,
                 protocolOptions,
-                artifactUploaders,
+                buildEventArtifactUploaderFactoryMap,
                 commandLineReporter,
                 startupOptionsProvider);
       } catch (Exception e) {
@@ -186,7 +189,10 @@ public abstract class BuildEventServiceModule<T extends BuildEventServiceOptions
 
       ImmutableSet<BuildEventTransport> bepTransports =
           BuildEventTransportFactory.createFromOptions(
-              bepOptions, protocolOptions, artifactUploaders, moduleEnvironment::exit);
+              bepOptions,
+              protocolOptions,
+              buildEventArtifactUploaderFactoryMap,
+              moduleEnvironment::exit);
 
       ImmutableSet.Builder<BuildEventTransport> transportsBuilder =
           ImmutableSet.<BuildEventTransport>builder().addAll(bepTransports);
@@ -216,7 +222,7 @@ public abstract class BuildEventServiceModule<T extends BuildEventServiceOptions
       ModuleEnvironment moduleEnvironment,
       Clock clock,
       BuildEventProtocolOptions protocolOptions,
-      BuildEventArtifactUploaderMap artifactUploaders,
+      BuildEventArtifactUploaderFactoryMap buildEventArtifactUploaderFactoryMap,
       EventHandler commandLineReporter,
       OptionsProvider startupOptionsProvider)
       throws IOException, OptionsParsingException {
@@ -224,8 +230,9 @@ public abstract class BuildEventServiceModule<T extends BuildEventServiceOptions
       logger.fine("BuildEventServiceTransport is disabled.");
       return null;
     } else {
-      logger.fine(format("Will create BuildEventServiceTransport streaming to '%s'",
-          besOptions.besBackend));
+      logger.fine(
+          format(
+              "Will create BuildEventServiceTransport streaming to '%s'", besOptions.besBackend));
 
       final String besResultsUrl;
       if (!Strings.isNullOrEmpty(besOptions.besResultsUrl)) {
@@ -247,7 +254,9 @@ public abstract class BuildEventServiceModule<T extends BuildEventServiceOptions
 
       BuildEventServiceClient client = createBesClient(besOptions, authTlsOptions);
       BuildEventArtifactUploader artifactUploader =
-          artifactUploaders.select(protocolOptions.buildEventUploadStrategy);
+          buildEventArtifactUploaderFactoryMap
+              .select(protocolOptions.buildEventUploadStrategy)
+              .create();
 
       BuildEventTransport besTransport =
           new BuildEventServiceTransport(
diff --git a/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceTransport.java b/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceTransport.java
index cb02ae3693..cd7492e0e2 100644
--- a/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceTransport.java
+++ b/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceTransport.java
@@ -404,7 +404,11 @@ public class BuildEventServiceTransport implements BuildEventTransport {
           publishBuildFinishedEvent(result);
         }
       } finally {
-        besClient.shutdown();
+        try {
+          besClient.shutdown();
+        } finally {
+          artifactUploader.shutdown();
+        }
       }
       return null;
     }