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.dataflow.worker.windmill.client.grpc;
  
  import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList.toImmutableList;
  
  import com.google.auto.value.AutoOneOf;
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.Comparator;
  import java.util.List;
  import java.util.concurrent.CountDownLatch;
  import java.util.stream.Stream;
  import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
  import org.apache.beam.runners.dataflow.worker.windmill.Windmill.ComputationGetDataRequest;
  import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataRequest;
  import org.apache.beam.runners.dataflow.worker.windmill.Windmill.KeyedGetDataRequest;
  import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStreamShutdownException;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /** Utility data classes for {@link GrpcGetDataStream}. */
  final class GrpcGetDataStreamRequests {
    private static final Logger LOG = LoggerFactory.getLogger(GrpcGetDataStreamRequests.class);
    private static final int STREAM_CANCELLED_ERROR_LOG_LIMIT = 3;
  
    private GrpcGetDataStreamRequests() {}
  
    private static String debugFormat(long value) {
      return String.format("%016x", value);
    }
  
    static class QueuedRequest {
      private final long id;
      private final ComputationOrGlobalDataRequest dataRequest;
      private AppendableInputStream responseStream;
  
      private QueuedRequest(
          long id, ComputationOrGlobalDataRequest dataRequest, long deadlineSeconds) {
        this.id = id;
        this.dataRequest = dataRequest;
        responseStream = new AppendableInputStream(deadlineSeconds);
      }
  
      static QueuedRequest forComputation(
          long id,
          String computation,
          KeyedGetDataRequest keyedGetDataRequest,
          long deadlineSeconds) {
        ComputationGetDataRequest computationGetDataRequest =
            ComputationGetDataRequest.newBuilder()
                .setComputationId(computation)
                .addRequests(keyedGetDataRequest)
                .build();
        return new QueuedRequest(
            id,
            ComputationOrGlobalDataRequest.computation(computationGetDataRequest),
            deadlineSeconds);
      }
  
      static QueuedRequest global(
          long id, GlobalDataRequest globalDataRequest, long deadlineSeconds) {
        return new QueuedRequest(
            id, ComputationOrGlobalDataRequest.global(globalDataRequest), deadlineSeconds);
      }
  
      static Comparator<QueuedRequest> globalRequestsFirst() {
        return (QueuedRequest r1, QueuedRequest r2) -> {
          boolean r1gd = r1.dataRequest.isGlobal();
          boolean r2gd = r2.dataRequest.isGlobal();
          return r1gd == r2gd ? 0 : (r1gd ? -1 : 1);
        };
      }
  
      long id() {
        return id;
      }
  
      long byteSize() {
        return dataRequest.serializedSize();
      }
  
      AppendableInputStream getResponseStream() {
        return responseStream;
      }
  
      void resetResponseStream() {
        this.responseStream = new AppendableInputStream(responseStream.getDeadlineSeconds());
      }
  
      public ComputationOrGlobalDataRequest getDataRequest() {
        return dataRequest;
      }
  
      void addToStreamingGetDataRequest(Windmill.StreamingGetDataRequest.Builder builder) {
        builder.addRequestId(id);
        if (dataRequest.isForComputation()) {
          builder.addStateRequest(dataRequest.computation());
        } else {
          builder.addGlobalDataRequest(dataRequest.global());
        }
      }
  
      @Override
      public final String toString() {
        return "QueuedRequest{" + "dataRequest=" + dataRequest + ", id=" + id + '}';
      }
    }
  
    /**
     * Represents a batch of queued requests. Methods are not thread-safe unless commented otherwise.
     */
    static class QueuedBatch {
      private final List<QueuedRequest> requests = new ArrayList<>();
      private final CountDownLatch sent = new CountDownLatch(1);
      private long byteSize = 0;
      private volatile boolean finalized = false;
      private volatile boolean failed = false;
  
      /** Returns a read-only view of requests. */
      List<QueuedRequest> requestsReadOnly() {
        return Collections.unmodifiableList(requests);
      }
  
      /**
       * Converts the batch to a {@link
       * org.apache.beam.runners.dataflow.worker.windmill.Windmill.StreamingGetDataRequest}.
       */
      Windmill.StreamingGetDataRequest asGetDataRequest() {
        Windmill.StreamingGetDataRequest.Builder builder =
            Windmill.StreamingGetDataRequest.newBuilder();
  
        requests.stream()
            // Put all global data requests first because there is only a single repeated field for
            // request ids and the initial ids correspond to global data requests if they are present.
            .sorted(QueuedRequest.globalRequestsFirst())
            .forEach(request -> request.addToStreamingGetDataRequest(builder));
  
        return builder.build();
      }
  
      boolean isEmpty() {
        return requests.isEmpty();
      }
  
      int requestsCount() {
        return requests.size();
      }
  
      long byteSize() {
        return byteSize;
      }
  
      boolean isFinalized() {
        return finalized;
      }
  
      void markFinalized() {
        finalized = true;
      }
  
      /** Adds a request to the batch. */
      void addRequest(QueuedRequest request) {
        requests.add(request);
        byteSize += request.byteSize();
      }
  
      /**
       * Let waiting for threads know that the request has been successfully sent.
       *
       * @implNote Thread safe.
       */
      void notifySent() {
        sent.countDown();
      }
  
      /** Let waiting for threads know that a failure occurred. */
      void notifyFailed() {
        failed = true;
        for (QueuedRequest request : requests) {
          request.getResponseStream().cancel();
        }
        sent.countDown();
      }
  
      /**
       * Block until notified of a successful send via {@link #notifySent()} or a non-retryable
       * failure via {@link #notifyFailed()}. On failure, throw an exception for waiters.
       *
       * @implNote Thread safe.
       */
      void waitForSendOrFailNotification()
          throws InterruptedException, WindmillStreamShutdownException {
        sent.await();
        if (failed) {
          ImmutableList<String> cancelledRequests = createStreamCancelledErrorMessages();
          if (!cancelledRequests.isEmpty()) {
            LOG.error("Requests failed for the following batches: {}", cancelledRequests);
            throw new WindmillStreamShutdownException(
                "Requests failed for batch containing "
                    + String.join(", ", cancelledRequests)
                    + " ... requests. This is most likely due to the stream being explicitly closed"
                    + " which happens when the work is marked as invalid on the streaming"
                    + " backend when key ranges shuffle around. This is transient and corresponding"
                    + " work will eventually be retried.");
          }
  
          throw new WindmillStreamShutdownException("Stream was shutdown while waiting for send.");
        }
      }
  
      private ImmutableList<String> createStreamCancelledErrorMessages() {
        return requests.stream()
            .flatMap(
                request -> {
                  switch (request.getDataRequest().getKind()) {
                    case GLOBAL:
                      return Stream.of("GetSideInput=" + request.getDataRequest().global());
                    case COMPUTATION:
                      return request.getDataRequest().computation().getRequestsList().stream()
                          .map(
                              keyedRequest ->
                                  "KeyedGetState=["
                                      + "shardingKey="
                                      + debugFormat(keyedRequest.getShardingKey())
                                      + "cacheToken="
                                      + debugFormat(keyedRequest.getCacheToken())
                                      + "workToken"
                                      + debugFormat(keyedRequest.getWorkToken())
                                      + "]");
                    default:
                      // Will never happen switch is exhaustive.
                      throw new IllegalStateException();
                  }
                })
            .limit(STREAM_CANCELLED_ERROR_LOG_LIMIT)
            .collect(toImmutableList());
      }
    }
  
    @AutoOneOf(ComputationOrGlobalDataRequest.Kind.class)
    abstract static class ComputationOrGlobalDataRequest {
      static ComputationOrGlobalDataRequest computation(
          ComputationGetDataRequest computationGetDataRequest) {
        return AutoOneOf_GrpcGetDataStreamRequests_ComputationOrGlobalDataRequest.computation(
            computationGetDataRequest);
      }
  
      static ComputationOrGlobalDataRequest global(GlobalDataRequest globalDataRequest) {
        return AutoOneOf_GrpcGetDataStreamRequests_ComputationOrGlobalDataRequest.global(
            globalDataRequest);
      }
  
      abstract Kind getKind();
  
      abstract ComputationGetDataRequest computation();
  
      abstract GlobalDataRequest global();
  
      boolean isGlobal() {
        return getKind() == Kind.GLOBAL;
      }
  
      boolean isForComputation() {
        return getKind() == Kind.COMPUTATION;
      }
  
      long serializedSize() {
        switch (getKind()) {
          case GLOBAL:
            return global().getSerializedSize();
          case COMPUTATION:
            return computation().getSerializedSize();
            // this will never happen since the switch is exhaustive.
          default:
            throw new UnsupportedOperationException("unknown dataRequest type.");
        }
      }
  
      enum Kind {
        COMPUTATION,
        GLOBAL
      }
    }
  }
  