content: /*
   * 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.fnexecution.artifact;
  
  import com.google.auto.value.AutoValue;
  import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
  import java.io.IOException;
  import java.io.OutputStream;
  import java.nio.channels.Channels;
  import java.nio.charset.StandardCharsets;
  import java.util.ArrayDeque;
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Queue;
  import java.util.concurrent.ArrayBlockingQueue;
  import java.util.concurrent.BlockingQueue;
  import java.util.concurrent.Callable;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.ConcurrentMap;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.concurrent.Future;
  import org.apache.beam.model.jobmanagement.v1.ArtifactApi;
  import org.apache.beam.model.jobmanagement.v1.ArtifactStagingServiceGrpc;
  import org.apache.beam.model.pipeline.v1.RunnerApi;
  import org.apache.beam.sdk.fn.IdGenerator;
  import org.apache.beam.sdk.fn.IdGenerators;
  import org.apache.beam.sdk.fn.server.FnService;
  import org.apache.beam.sdk.io.FileSystems;
  import org.apache.beam.sdk.io.fs.MatchResult;
  import org.apache.beam.sdk.io.fs.MoveOptions;
  import org.apache.beam.sdk.io.fs.ResolveOptions;
  import org.apache.beam.sdk.io.fs.ResourceId;
  import org.apache.beam.sdk.util.MimeTypes;
  import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
  import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.InvalidProtocolBufferException;
  import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Status;
  import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.StatusException;
  import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.StreamObserver;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Splitter;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  @SuppressWarnings({
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  public class ArtifactStagingService
      extends ArtifactStagingServiceGrpc.ArtifactStagingServiceImplBase implements FnService {
  
    private static final Logger LOG = LoggerFactory.getLogger(ArtifactStagingService.class);
  
    private final ArtifactDestinationProvider destinationProvider;
  
    private final ConcurrentMap<String, Map<String, List<RunnerApi.ArtifactInformation>>> toStage =
        new ConcurrentHashMap<>();
  
    private final ConcurrentMap<String, Map<String, List<RunnerApi.ArtifactInformation>>> staged =
        new ConcurrentHashMap<>();
  
    public ArtifactStagingService(ArtifactDestinationProvider destinationProvider) {
      this.destinationProvider = destinationProvider;
    }
  
    /**
     * Registers a set of artifacts to be staged with this service.
     *
     * <p>A client (e.g. a Beam SDK) is expected to connect to this service with the given staging
     * token and offer resolution and retrieval of this set of artifacts.
     *
     * @param stagingToken a staging token for this job
     * @param artifacts all artifacts to stage, keyed by environment
     */
    public void registerJob(
        String stagingToken, Map<String, List<RunnerApi.ArtifactInformation>> artifacts) {
      assert !toStage.containsKey(stagingToken);
      toStage.put(stagingToken, artifacts);
    }
  
    /**
     * Returns the rewritten artifacts associated with this job, keyed by environment.
     *
     * <p>This should be called after the client has finished offering artifacts.
     *
     * @param stagingToken a staging token for this job
     */
    public Map<String, List<RunnerApi.ArtifactInformation>> getStagedArtifacts(String stagingToken) {
      toStage.remove(stagingToken);
      return staged.remove(stagingToken);
    }
  
    public void removeStagedArtifacts(String stagingToken) throws IOException {
      destinationProvider.removeStagedArtifacts(stagingToken);
    }
  
    /** Provides a concrete location to which artifacts can be staged on retrieval. */
    public interface ArtifactDestinationProvider {
      ArtifactDestination getDestination(String stagingToken, String name) throws IOException;
  
      void removeStagedArtifacts(String stagingToken) throws IOException;
    }
  
    /**
     * A pairing of a newly created artifact type and an output stream that will be readable at that
     * type.
     */
    @AutoValue
    public abstract static class ArtifactDestination {
      public static ArtifactDestination create(
          String typeUrn, ByteString typePayload, OutputStream out) {
        return new AutoValue_ArtifactStagingService_ArtifactDestination(typeUrn, typePayload, out);
      }
  
      public static ArtifactDestination fromFile(String path) throws IOException {
        return fromFile(
            path,
            Channels.newOutputStream(
                FileSystems.create(
                    FileSystems.matchNewResource(path, false /* isDirectory */), MimeTypes.BINARY)));
      }
  
      public static ArtifactDestination fromFile(String path, OutputStream out) {
        return create(
            ArtifactRetrievalService.FILE_ARTIFACT_URN,
            RunnerApi.ArtifactFilePayload.newBuilder().setPath(path).build().toByteString(),
            out);
      }
  
      public abstract String getTypeUrn();
  
      public abstract ByteString getTypePayload();
  
      public abstract OutputStream getOutputStream();
    }
  
    /**
     * An ArtifactDestinationProvider that places new artifacts as files in a Beam filesystem.
     *
     * @param root the directory in which to place all artifacts
     */
    public static ArtifactDestinationProvider beamFilesystemArtifactDestinationProvider(String root) {
      return new ArtifactDestinationProvider() {
        @Override
        public ArtifactDestination getDestination(String stagingToken, String name)
            throws IOException {
          ResourceId path =
              stagingDir(stagingToken)
                  .resolve(name, ResolveOptions.StandardResolveOptions.RESOLVE_FILE);
          return ArtifactDestination.fromFile(path.toString());
        }
  
        @Override
        public void removeStagedArtifacts(String stagingToken) throws IOException {
          // TODO(robertwb): Consider adding recursive delete.
          ResourceId stagingDir = stagingDir(stagingToken);
          List<ResourceId> toDelete = new ArrayList<>();
          for (MatchResult match :
              FileSystems.matchResources(
                  ImmutableList.of(
                      stagingDir.resolve("*", ResolveOptions.StandardResolveOptions.RESOLVE_FILE)))) {
            for (MatchResult.Metadata m : match.metadata()) {
              toDelete.add(m.resourceId());
            }
          }
          FileSystems.delete(toDelete, MoveOptions.StandardMoveOptions.IGNORE_MISSING_FILES);
          FileSystems.delete(
              ImmutableList.of(stagingDir), MoveOptions.StandardMoveOptions.IGNORE_MISSING_FILES);
        }
  
        private ResourceId stagingDir(String stagingToken) {
          return FileSystems.matchNewResource(root, true)
              .resolve(
                  Hashing.sha256().hashString(stagingToken, StandardCharsets.UTF_8).toString(),
                  ResolveOptions.StandardResolveOptions.RESOLVE_DIRECTORY);
        }
      };
    }
  
    private enum State {
      START,
      RESOLVE,
      GET,
      GETCHUNK,
      DONE,
      ERROR,
    }
  
    /**
     * Like the standard Semaphore, but allows an acquire to go over the limit if there is any room.
     *
     * <p>Also allows setting an error, to avoid issues with un-released acquires after error.
     */
    private static class OverflowingSemaphore {
      private int totalPermits;
      private int usedPermits;
      private Exception exception;
  
      public OverflowingSemaphore(int totalPermits) {
        this.totalPermits = totalPermits;
        this.usedPermits = 0;
      }
  
      synchronized void aquire(int permits) throws Exception {
        while (usedPermits >= totalPermits) {
          if (exception != null) {
            throw exception;
          }
          this.wait();
        }
        usedPermits += permits;
      }
  
      synchronized void release(int permits) {
        usedPermits -= permits;
        this.notifyAll();
      }
  
      synchronized void setException(Exception exception) {
        this.exception = exception;
        this.notifyAll();
      }
    }
  
    /** A task that pulls bytes off a queue and actually writes them to a staging location. */
    private class StoreArtifact implements Callable<RunnerApi.ArtifactInformation> {
  
      private String stagingToken;
      private String name;
      private RunnerApi.ArtifactInformation originalArtifact;
      private BlockingQueue<ByteString> bytesQueue;
      private OverflowingSemaphore totalPendingBytes;
  
      public StoreArtifact(
          String stagingToken,
          String name,
          RunnerApi.ArtifactInformation originalArtifact,
          BlockingQueue<ByteString> bytesQueue,
          OverflowingSemaphore totalPendingBytes) {
        this.stagingToken = stagingToken;
        this.name = name;
        this.originalArtifact = originalArtifact;
        this.bytesQueue = bytesQueue;
        this.totalPendingBytes = totalPendingBytes;
      }
  
      @Override
      public RunnerApi.ArtifactInformation call() throws IOException {
        try {
          ArtifactDestination dest = destinationProvider.getDestination(stagingToken, name);
          LOG.debug("Storing artifact for {}.{} at {}", stagingToken, name, dest);
          ByteString chunk = bytesQueue.take();
          while (chunk.size() > 0) {
            totalPendingBytes.release(chunk.size());
            chunk.writeTo(dest.getOutputStream());
            chunk = bytesQueue.take();
          }
          dest.getOutputStream().close();
          return originalArtifact
              .toBuilder()
              .setTypeUrn(dest.getTypeUrn())
              .setTypePayload(dest.getTypePayload())
              .build();
        } catch (IOException | InterruptedException exn) {
          // As this thread will no longer be draining the queue, we don't want to get stuck writing
          // to it.
          totalPendingBytes.setException(exn);
          LOG.error("Exception staging artifacts", exn);
          if (exn instanceof IOException) {
            throw (IOException) exn;
          } else {
            throw new RuntimeException(exn);
          }
        }
      }
    }
  
    @Override
    public StreamObserver<ArtifactApi.ArtifactResponseWrapper> reverseArtifactRetrievalService(
        StreamObserver<ArtifactApi.ArtifactRequestWrapper> responseObserver) {
  
      return new StreamObserver<ArtifactApi.ArtifactResponseWrapper>() {
  
        /** The maximum number of parallel threads to use to stage. */
        public static final int THREAD_POOL_SIZE = 10;
  
        /** The maximum number of bytes to buffer across all writes before throttling. */
        public static final int MAX_PENDING_BYTES = 100 << 20; // 100 MB
  
        IdGenerator idGenerator = IdGenerators.incrementingLongs();
  
        String stagingToken;
        Map<String, List<RunnerApi.ArtifactInformation>> toResolve;
        Map<String, List<Future<RunnerApi.ArtifactInformation>>> stagedFutures;
        ExecutorService stagingExecutor;
        OverflowingSemaphore totalPendingBytes;
  
        State state = State.START;
        Queue<String> pendingResolves;
        String currentEnvironment;
        Queue<RunnerApi.ArtifactInformation> pendingGets;
        BlockingQueue<ByteString> currentOutput;
  
        @Override
        @SuppressFBWarnings(value = "SF_SWITCH_FALLTHROUGH", justification = "fallthrough intended")
        // May be called by different threads for the same request; synchronized for memory
        // synchronization.
        public synchronized void onNext(ArtifactApi.ArtifactResponseWrapper responseWrapper) {
          switch (state) {
            case START:
              stagingToken = responseWrapper.getStagingToken();
              LOG.info("Staging artifacts for {}.", stagingToken);
              toResolve = toStage.get(stagingToken);
              if (toResolve == null) {
                responseObserver.onError(
                    new StatusException(
                        Status.INVALID_ARGUMENT.withDescription(
                            "Unknown staging token " + stagingToken)));
                return;
              }
              stagedFutures = new ConcurrentHashMap<>();
              pendingResolves = new ArrayDeque<>();
              pendingResolves.addAll(toResolve.keySet());
              stagingExecutor = Executors.newFixedThreadPool(THREAD_POOL_SIZE);
              totalPendingBytes = new OverflowingSemaphore(MAX_PENDING_BYTES);
              resolveNextEnvironment(responseObserver);
              break;
  
            case RESOLVE:
              {
                currentEnvironment = pendingResolves.remove();
                stagedFutures.put(currentEnvironment, new ArrayList<>());
                pendingGets = new ArrayDeque<>();
                for (RunnerApi.ArtifactInformation artifact :
                    responseWrapper.getResolveArtifactResponse().getReplacementsList()) {
                  Optional<RunnerApi.ArtifactInformation> fetched = getLocal();
                  if (fetched.isPresent()) {
                    stagedFutures
                        .get(currentEnvironment)
                        .add(CompletableFuture.completedFuture(fetched.get()));
                  } else {
                    pendingGets.add(artifact);
                    responseObserver.onNext(
                        ArtifactApi.ArtifactRequestWrapper.newBuilder()
                            .setGetArtifact(
                                ArtifactApi.GetArtifactRequest.newBuilder().setArtifact(artifact))
                            .build());
                  }
                }
                LOG.info(
                    "Getting {} artifacts for {}.{}.",
                    pendingGets.size(),
                    stagingToken,
                    pendingResolves.peek());
                if (pendingGets.isEmpty()) {
                  resolveNextEnvironment(responseObserver);
                } else {
                  state = State.GET;
                }
                break;
              }
  
            case GET:
              RunnerApi.ArtifactInformation currentArtifact = pendingGets.remove();
              String name = createFilename(currentEnvironment, currentArtifact);
              try {
                LOG.debug("Storing artifacts for {} as {}", stagingToken, name);
                currentOutput = new ArrayBlockingQueue<ByteString>(100);
                stagedFutures
                    .get(currentEnvironment)
                    .add(
                        stagingExecutor.submit(
                            new StoreArtifact(
                                stagingToken,
                                name,
                                currentArtifact,
                                currentOutput,
                                totalPendingBytes)));
              } catch (Exception exn) {
                LOG.error("Error submitting.", exn);
                responseObserver.onError(exn);
              }
              state = State.GETCHUNK;
              // fall through
  
            case GETCHUNK:
              try {
                ByteString chunk = responseWrapper.getGetArtifactResponse().getData();
                if (chunk.size() > 0) { // Make sure we don't accidentally send the EOF value.
                  totalPendingBytes.aquire(chunk.size());
                  currentOutput.put(chunk);
                }
                if (responseWrapper.getIsLast()) {
                  currentOutput.put(ByteString.EMPTY); // The EOF value.
                  if (pendingGets.isEmpty()) {
                    resolveNextEnvironment(responseObserver);
                  } else {
                    state = State.GET;
                    LOG.debug("Waiting for {}", pendingGets.peek());
                  }
                }
              } catch (Exception exn) {
                LOG.error("Error submitting.", exn);
                onError(exn);
              }
              break;
  
            default:
              responseObserver.onError(
                  new StatusException(
                      Status.INVALID_ARGUMENT.withDescription("Illegal state " + state)));
          }
        }
  
        private void resolveNextEnvironment(
            StreamObserver<ArtifactApi.ArtifactRequestWrapper> responseObserver) {
          if (pendingResolves.isEmpty()) {
            finishStaging(responseObserver);
          } else {
            state = State.RESOLVE;
            LOG.info("Resolving artifacts for {}.{}.", stagingToken, pendingResolves.peek());
            responseObserver.onNext(
                ArtifactApi.ArtifactRequestWrapper.newBuilder()
                    .setResolveArtifact(
                        ArtifactApi.ResolveArtifactsRequest.newBuilder()
                            .addAllArtifacts(toResolve.get(pendingResolves.peek())))
                    .build());
          }
        }
  
        private void finishStaging(
            StreamObserver<ArtifactApi.ArtifactRequestWrapper> responseObserver) {
          LOG.debug("Finishing staging for {}.", stagingToken);
          Map<String, List<RunnerApi.ArtifactInformation>> staged = new HashMap<>();
          try {
            for (Map.Entry<String, List<Future<RunnerApi.ArtifactInformation>>> entry :
                stagedFutures.entrySet()) {
              List<RunnerApi.ArtifactInformation> envStaged = new ArrayList<>();
              for (Future<RunnerApi.ArtifactInformation> future : entry.getValue()) {
                envStaged.add(future.get());
              }
              staged.put(entry.getKey(), envStaged);
            }
            ArtifactStagingService.this.staged.put(stagingToken, staged);
            stagingExecutor.shutdown();
            state = State.DONE;
            LOG.info("Artifacts fully staged for {}.", stagingToken);
            responseObserver.onCompleted();
          } catch (Exception exn) {
            LOG.error("Error staging artifacts", exn);
            responseObserver.onError(exn);
            state = State.ERROR;
            return;
          }
        }
  
        /**
         * Return an alternative artifact if we do not need to get this over the artifact API, or
         * possibly at all.
         */
        private Optional<RunnerApi.ArtifactInformation> getLocal() {
          return Optional.empty();
        }
  
        /**
         * Attempts to provide a reasonable filename for the artifact.
         *
         * @param index a monotonically increasing index, which provides uniqueness
         * @param environment the environment id
         * @param artifact the artifact itself
         */
        private String createFilename(String environment, RunnerApi.ArtifactInformation artifact) {
          String path;
          try {
            if (artifact.getRoleUrn().equals(ArtifactRetrievalService.STAGING_TO_ARTIFACT_URN)) {
              path =
                  RunnerApi.ArtifactStagingToRolePayload.parseFrom(artifact.getRolePayload())
                      .getStagedName();
            } else if (artifact.getTypeUrn().equals(ArtifactRetrievalService.FILE_ARTIFACT_URN)) {
              path = RunnerApi.ArtifactFilePayload.parseFrom(artifact.getTypePayload()).getPath();
            } else if (artifact.getTypeUrn().equals(ArtifactRetrievalService.URL_ARTIFACT_URN)) {
              path = RunnerApi.ArtifactUrlPayload.parseFrom(artifact.getTypePayload()).getUrl();
            } else {
              path = "artifact";
            }
          } catch (InvalidProtocolBufferException exn) {
            throw new RuntimeException(exn);
          }
          // Limit to the last contiguous alpha-numeric sequence. In particular, this will exclude
          // all path separators.
          List<String> components = Splitter.onPattern("[^A-Za-z-_.]]").splitToList(path);
          String base = components.get(components.size() - 1);
          return clip(
              String.format("%s-%s-%s", idGenerator.getId(), clip(environment, 25), base), 100);
        }
  
        private String clip(String s, int maxLength) {
          return s.length() < maxLength ? s : s.substring(0, maxLength);
        }
  
        @Override
        public void onError(Throwable throwable) {
          stagingExecutor.shutdownNow();
          LOG.error("Error staging artifacts", throwable);
          state = State.ERROR;
        }
  
        @Override
        public void onCompleted() {
          Preconditions.checkArgument(state == State.DONE);
        }
      };
    }
  
    @Override
    public void close() throws Exception {
      // Nothing to close.
    }
  
    /**
     * Lazily stages artifacts by letting an ArtifactStagingService resolve and request artifacts.
     *
     * @param retrievalService an ArtifactRetrievalService used to resolve and retrieve artifacts
     * @param stagingService an ArtifactStagingService stub which will request artifacts
     * @param stagingToken the staging token of the job whose artifacts will be retrieved
     * @throws InterruptedException
     * @throws IOException
     */
    public static void offer(
        ArtifactRetrievalService retrievalService,
        ArtifactStagingServiceGrpc.ArtifactStagingServiceStub stagingService,
        String stagingToken)
        throws ExecutionException, InterruptedException {
      new StagingDriver(retrievalService, stagingService, stagingToken).getCompletionFuture().get();
    }
  
    /** Actually implements the reverse retrieval protocol. */
    private static class StagingDriver implements StreamObserver<ArtifactApi.ArtifactRequestWrapper> {
  
      private final ArtifactRetrievalService retrievalService;
      private final StreamObserver<ArtifactApi.ArtifactResponseWrapper> responseObserver;
      private final CompletableFuture<Void> completionFuture;
  
      public StagingDriver(
          ArtifactRetrievalService retrievalService,
          ArtifactStagingServiceGrpc.ArtifactStagingServiceStub stagingService,
          String stagingToken) {
        this.retrievalService = retrievalService;
        completionFuture = new CompletableFuture<Void>();
        responseObserver = stagingService.reverseArtifactRetrievalService(this);
        responseObserver.onNext(
            ArtifactApi.ArtifactResponseWrapper.newBuilder().setStagingToken(stagingToken).build());
      }
  
      public CompletableFuture<?> getCompletionFuture() {
        return completionFuture;
      }
  
      @Override
      public void onNext(ArtifactApi.ArtifactRequestWrapper requestWrapper) {
  
        if (completionFuture.isCompletedExceptionally()) {
          try {
            completionFuture.get();
          } catch (Throwable th) {
            responseObserver.onError(th);
            return;
          }
        }
  
        if (requestWrapper.hasResolveArtifact()) {
          retrievalService.resolveArtifacts(
              requestWrapper.getResolveArtifact(),
              new StreamObserver<ArtifactApi.ResolveArtifactsResponse>() {
  
                @Override
                public void onNext(ArtifactApi.ResolveArtifactsResponse resolveArtifactsResponse) {
                  responseObserver.onNext(
                      ArtifactApi.ArtifactResponseWrapper.newBuilder()
                          .setResolveArtifactResponse(resolveArtifactsResponse)
                          .build());
                }
  
                @Override
                public void onError(Throwable throwable) {
                  completionFuture.completeExceptionally(throwable);
                  responseObserver.onError(throwable);
                }
  
                @Override
                public void onCompleted() {}
              });
        } else if (requestWrapper.hasGetArtifact()) {
          retrievalService.getArtifact(
              requestWrapper.getGetArtifact(),
              new StreamObserver<ArtifactApi.GetArtifactResponse>() {
  
                @Override
                public void onNext(ArtifactApi.GetArtifactResponse getArtifactResponse) {
                  responseObserver.onNext(
                      ArtifactApi.ArtifactResponseWrapper.newBuilder()
                          .setGetArtifactResponse(getArtifactResponse)
                          .build());
                }
  
                @Override
                public void onError(Throwable throwable) {
                  completionFuture.completeExceptionally(throwable);
                  responseObserver.onError(throwable);
                }
  
                @Override
                public void onCompleted() {
                  responseObserver.onNext(
                      ArtifactApi.ArtifactResponseWrapper.newBuilder()
                          .setGetArtifactResponse(
                              ArtifactApi.GetArtifactResponse.newBuilder().build())
                          .setIsLast(true)
                          .build());
                }
              });
        } else {
          Throwable exn =
              new StatusException(
                  Status.INVALID_ARGUMENT.withDescription(
                      "Expected either a resolve or get request."));
          completionFuture.completeExceptionally(exn);
          responseObserver.onError(exn);
        }
      }
  
      @Override
      public void onError(Throwable throwable) {
        completionFuture.completeExceptionally(throwable);
      }
  
      @Override
      public void onCompleted() {
        responseObserver.onCompleted();
        completionFuture.complete(null);
      }
    }
  }
  