Enable artifact staging during Prism Runner lifecycle (#32084)

runners/prism/java/build.gradle

@@ -30,6 +30,7 @@ dependencies {
     implementation project(path: ":model:pipeline", configuration: "shadow")
     implementation project(path: ":sdks:java:core", configuration: "shadow")
     implementation project(path: ":sdks:java:harness", configuration: "shadow")
+    implementation project(":runners:java-fn-execution")
     implementation project(":runners:portability:java")
 
     implementation library.java.joda_time

runners/prism/java/src/main/java/org/apache/beam/runners/prism/PrismArtifactStager.java

@@ -0,0 +1,173 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.prism;
+
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.Optional;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import org.apache.beam.model.jobmanagement.v1.ArtifactStagingServiceGrpc;
+import org.apache.beam.model.jobmanagement.v1.JobApi;
+import org.apache.beam.model.jobmanagement.v1.JobServiceGrpc;
+import org.apache.beam.model.pipeline.v1.Endpoints;
+import org.apache.beam.runners.fnexecution.artifact.ArtifactRetrievalService;
+import org.apache.beam.runners.fnexecution.artifact.ArtifactStagingService;
+import org.apache.beam.sdk.fn.channel.ManagedChannelFactory;
+import org.apache.beam.vendor.grpc.v1p60p1.io.grpc.ManagedChannel;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Stages {@link org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline} artifacts of prepared jobs.
+ */
+@AutoValue
+abstract class PrismArtifactStager implements AutoCloseable {
+
+  private static final Logger LOG = LoggerFactory.getLogger(PrismArtifactStager.class);
+
+  /**
+   * Instantiate a {@link PrismArtifactStager} via call to {@link #of(String, String)}, assigning
+   * {@link Builder#setStagingEndpoint} using {@param prepareJobResponse} {@link
+   * JobApi.PrepareJobResponse#getArtifactStagingEndpoint} and {@link
+   * JobApi.PrepareJobResponse#getStagingSessionToken}.
+   */
+  static PrismArtifactStager of(JobApi.PrepareJobResponse prepareJobResponse) {
+    return of(
+        prepareJobResponse.getArtifactStagingEndpoint().getUrl(),
+        prepareJobResponse.getStagingSessionToken());
+  }
+
+  /**
+   * Instantiates a {@link PrismArtifactStager} from the {@param stagingEndpoint} URL and {@param
+   * stagingSessionToken} to instantiate the {@link #getRetrievalService}, {@link
+   * #getManagedChannel}, and {@link #getStagingServiceStub} defaults. See the referenced getters
+   * for more details.
+   */
+  static PrismArtifactStager of(String stagingEndpoint, String stagingSessionToken) {
+    return PrismArtifactStager.builder()
+        .setStagingEndpoint(stagingEndpoint)
+        .setStagingSessionToken(stagingSessionToken)
+        .build();
+  }
+
+  static Builder builder() {
+    return new AutoValue_PrismArtifactStager.Builder();
+  }
+
+  /**
+   * Stage the {@link org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline} artifacts via {@link
+   * ArtifactStagingService#offer} supplying {@link #getRetrievalService}, {@link
+   * #getStagingServiceStub}, and {@link #getStagingSessionToken}.
+   */
+  void stage() throws ExecutionException, InterruptedException {
+    LOG.info("staging artifacts at {}", getStagingEndpoint());
+    ArtifactStagingService.offer(
+        getRetrievalService(), getStagingServiceStub(), getStagingSessionToken());
+  }
+
+  /** The URL of the {@link ArtifactStagingService}. */
+  abstract String getStagingEndpoint();
+
+  /**
+   * Token associated with a staging session and acquired from a {@link
+   * JobServiceGrpc.JobServiceStub#prepare}'s {@link JobApi.PrepareJobResponse}.
+   */
+  abstract String getStagingSessionToken();
+
+  /**
+   * The service that retrieves artifacts; defaults to instantiating from the default {@link
+   * ArtifactRetrievalService#ArtifactRetrievalService()} constructor.
+   */
+  abstract ArtifactRetrievalService getRetrievalService();
+
+  /**
+   * Used to instantiate the {@link #getStagingServiceStub}. By default, instantiates using {@link
+   * ManagedChannelFactory#forDescriptor(Endpoints.ApiServiceDescriptor)}, where {@link
+   * Endpoints.ApiServiceDescriptor} is instantiated via {@link
+   * Endpoints.ApiServiceDescriptor.Builder#setUrl(String)} and the URL provided by {@link
+   * #getStagingEndpoint}.
+   */
+  abstract ManagedChannel getManagedChannel();
+
+  /**
+   * Required by {@link ArtifactStagingService#offer}. By default, instantiates using {@link
+   * ArtifactStagingServiceGrpc#newStub} and {@link #getManagedChannel}.
+   */
+  abstract ArtifactStagingServiceGrpc.ArtifactStagingServiceStub getStagingServiceStub();
+
+  @Override
+  public void close() {
+    LOG.info("shutting down {}", PrismArtifactStager.class);
+    getRetrievalService().close();
+    getManagedChannel().shutdown();
+    try {
+      getManagedChannel().awaitTermination(3000L, TimeUnit.MILLISECONDS);
+    } catch (InterruptedException ignored) {
+    }
+  }
+
+  @AutoValue.Builder
+  abstract static class Builder {
+
+    abstract Builder setStagingEndpoint(String stagingEndpoint);
+
+    abstract Optional<String> getStagingEndpoint();
+
+    abstract Builder setStagingSessionToken(String stagingSessionToken);
+
+    abstract Builder setRetrievalService(ArtifactRetrievalService retrievalService);
+
+    abstract Optional<ArtifactRetrievalService> getRetrievalService();
+
+    abstract Builder setManagedChannel(ManagedChannel managedChannel);
+
+    abstract Optional<ManagedChannel> getManagedChannel();
+
+    abstract Builder setStagingServiceStub(
+        ArtifactStagingServiceGrpc.ArtifactStagingServiceStub stub);
+
+    abstract Optional<ArtifactStagingServiceGrpc.ArtifactStagingServiceStub>
+        getStagingServiceStub();
+
+    abstract PrismArtifactStager autoBuild();
+
+    final PrismArtifactStager build() {
+
+      checkState(getStagingEndpoint().isPresent(), "missing staging endpoint");
+      ManagedChannelFactory channelFactory = ManagedChannelFactory.createDefault();
+
+      if (!getManagedChannel().isPresent()) {
+        Endpoints.ApiServiceDescriptor descriptor =
+            Endpoints.ApiServiceDescriptor.newBuilder().setUrl(getStagingEndpoint().get()).build();
+        setManagedChannel(channelFactory.forDescriptor(descriptor));
+      }
+
+      if (!getStagingServiceStub().isPresent()) {
+        setStagingServiceStub(ArtifactStagingServiceGrpc.newStub(getManagedChannel().get()));
+      }
+
+      if (!getRetrievalService().isPresent()) {
+        setRetrievalService(new ArtifactRetrievalService());
+      }
+
+      return autoBuild();
+    }
+  }
+}

runners/prism/java/src/test/java/org/apache/beam/runners/prism/PrismArtifactStagerTest.java

@@ -0,0 +1,143 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.prism;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.apache.beam.runners.fnexecution.artifact.ArtifactRetrievalService.EMBEDDED_ARTIFACT_URN;
+import static org.junit.Assert.assertThrows;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ExecutionException;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.runners.fnexecution.artifact.ArtifactRetrievalService;
+import org.apache.beam.runners.fnexecution.artifact.ArtifactStagingService;
+import org.apache.beam.vendor.grpc.v1p60p1.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p60p1.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p60p1.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p60p1.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p60p1.io.grpc.testing.GrpcCleanupRule;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.apache.commons.io.output.ByteArrayOutputStream;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link PrismArtifactStager}. */
+@RunWith(JUnit4.class)
+public class PrismArtifactStagerTest {
+
+  @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
+
+  final ArtifactStagingService stagingService =
+      new ArtifactStagingService(new TestDestinationProvider());
+
+  @Test
+  public void givenValidArtifacts_stages()
+      throws IOException, ExecutionException, InterruptedException {
+    PrismArtifactStager underTest = prismArtifactStager(validArtifacts());
+    assertThat(underTest.getManagedChannel().isShutdown()).isFalse();
+    underTest.stage();
+    assertThat(stagingService.getStagedArtifacts(underTest.getStagingSessionToken())).isNotEmpty();
+    underTest.close();
+    assertThat(underTest.getManagedChannel().isShutdown()).isTrue();
+  }
+
+  @Test
+  public void givenErrors_performsGracefulCleanup() throws IOException {
+    PrismArtifactStager underTest = prismArtifactStager(invalidArtifacts());
+    assertThat(underTest.getManagedChannel().isShutdown()).isFalse();
+    ExecutionException error = assertThrows(ExecutionException.class, underTest::stage);
+    assertThat(error.getMessage()).contains("Unexpected artifact type: invalid-type-urn");
+    assertThat(underTest.getManagedChannel().isShutdown()).isFalse();
+    underTest.close();
+    assertThat(underTest.getManagedChannel().isShutdown()).isTrue();
+  }
+
+  private PrismArtifactStager prismArtifactStager(
+      Map<String, List<RunnerApi.ArtifactInformation>> artifacts) throws IOException {
+    String serverName = InProcessServerBuilder.generateName();
+    ArtifactRetrievalService retrievalService = new ArtifactRetrievalService();
+    String stagingToken = "staging-token";
+    stagingService.registerJob(stagingToken, artifacts);
+
+    grpcCleanup.register(
+        InProcessServerBuilder.forName(serverName)
+            .directExecutor()
+            .addService(stagingService)
+            .addService(retrievalService)
+            .build()
+            .start());
+
+    ManagedChannel channel =
+        grpcCleanup.register(InProcessChannelBuilder.forName(serverName).build());
+
+    return PrismArtifactStager.builder()
+        .setStagingEndpoint("ignore")
+        .setStagingSessionToken(stagingToken)
+        .setManagedChannel(channel)
+        .build();
+  }
+
+  private Map<String, List<RunnerApi.ArtifactInformation>> validArtifacts() {
+    return ImmutableMap.of(
+        "env1",
+        Collections.singletonList(
+            RunnerApi.ArtifactInformation.newBuilder()
+                .setTypeUrn(EMBEDDED_ARTIFACT_URN)
+                .setTypePayload(
+                    RunnerApi.EmbeddedFilePayload.newBuilder()
+                        .setData(ByteString.copyFromUtf8("type-payload"))
+                        .build()
+                        .toByteString())
+                .setRoleUrn("role-urn")
+                .build()));
+  }
+
+  private Map<String, List<RunnerApi.ArtifactInformation>> invalidArtifacts() {
+    return ImmutableMap.of(
+        "env1",
+        Collections.singletonList(
+            RunnerApi.ArtifactInformation.newBuilder()
+                .setTypeUrn("invalid-type-urn")
+                .setTypePayload(
+                    RunnerApi.EmbeddedFilePayload.newBuilder()
+                        .setData(ByteString.copyFromUtf8("type-payload"))
+                        .build()
+                        .toByteString())
+                .setRoleUrn("role-urn")
+                .build()));
+  }
+
+  private static class TestDestinationProvider
+      implements ArtifactStagingService.ArtifactDestinationProvider {
+
+    @Override
+    public ArtifactStagingService.ArtifactDestination getDestination(
+        String stagingToken, String name) throws IOException {
+      return ArtifactStagingService.ArtifactDestination.create(
+          EMBEDDED_ARTIFACT_URN, ByteString.EMPTY, new ByteArrayOutputStream());
+    }
+
+    @Override
+    public void removeStagedArtifacts(String stagingToken) throws IOException {}
+  }
+}