content: /*
   * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
   * (the "License"). You may not use this work except in compliance with the License, which is
   * available at www.apache.org/licenses/LICENSE-2.0
   *
   * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
   * either express or implied, as more fully set forth in the License.
   *
   * See the NOTICE file distributed with this work for information regarding copyright ownership.
   */
  
  package alluxio.client.block.stream;
  
  import alluxio.Constants;
  import alluxio.exception.status.AlluxioStatusException;
  import alluxio.exception.status.CancelledException;
  import alluxio.exception.status.DeadlineExceededException;
  import alluxio.exception.status.UnavailableException;
  import alluxio.resource.LockResource;
  import alluxio.util.LogUtils;
  
  import io.grpc.Status;
  import io.grpc.StatusRuntimeException;
  import io.grpc.stub.ClientCallStreamObserver;
  import io.grpc.stub.ClientResponseObserver;
  import io.grpc.stub.StreamObserver;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.io.IOException;
  import java.util.Optional;
  import java.util.concurrent.ArrayBlockingQueue;
  import java.util.concurrent.BlockingQueue;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.locks.Condition;
  import java.util.concurrent.locks.ReentrantLock;
  import java.util.function.Function;
  import javax.annotation.concurrent.GuardedBy;
  import javax.annotation.concurrent.NotThreadSafe;
  
  /**
   * A helper class for accessing gRPC bi-directional stream synchronously.
   *
   * @param <ReqT> type of the request
   * @param <ResT> type of the response
   */
  @NotThreadSafe
  public class GrpcBlockingStream<ReqT, ResT> {
    private static final Logger LOG = LoggerFactory.getLogger(GrpcBlockingStream.class);
    private final StreamObserver<ResT> mResponseObserver;
    private final ClientCallStreamObserver<ReqT> mRequestObserver;
    /** Buffer that stores responses to be consumed by {@link GrpcBlockingStream#receive(long)}. */
    private final BlockingQueue<Object> mResponses;
    private final String mDescription;
    private volatile boolean mCompleted = false;
    private volatile boolean mClosed = false;
    private volatile boolean mCanceled = false;
  
    /**
     * Uses to guarantee the operation ordering.
     *
     * NOTE: {@link StreamObserver} events are async.
     * gRPC worker threads executes the response {@link StreamObserver} events.
     */
    private final ReentrantLock mLock = new ReentrantLock();
  
    @GuardedBy("mLock")
    private Throwable mError;
    /** This condition is met if mError != null or client is ready to send data. */
    private final Condition mReadyOrFailed = mLock.newCondition();
    /** This is set by the grpc threads, and checked/read by the client. */
    private volatile boolean mClosedFromRemote = false;
  
    /**
     * @param rpcFunc the gRPC bi-directional stream stub function
     * @param bufferSize maximum number of incoming messages the buffer can hold
     * @param description description of this stream
     */
    public GrpcBlockingStream(Function<StreamObserver<ResT>, StreamObserver<ReqT>> rpcFunc,
        int bufferSize, String description) {
      LOG.debug("Opening stream ({})", description);
      mResponses = new ArrayBlockingQueue<>(bufferSize);
      mResponseObserver = new ResponseStreamObserver();
      mRequestObserver = (ClientCallStreamObserver) rpcFunc.apply(mResponseObserver);
      mDescription = description;
    }
  
    /**
     * Sends a request. Will wait until the stream is ready before sending or timeout if the
     * given timeout is reached.
     *
     * @param request the request
     * @param timeoutMs maximum wait time before throwing a {@link DeadlineExceededException}
     * @throws IOException if any error occurs
     */
    public void send(ReqT request, long timeoutMs) throws IOException {
      if (mClosed || mCanceled || mClosedFromRemote) {
        throw new CancelledException(formatErrorMessage(
            "Failed to send request %s: stream is already closed or cancelled. clientClosed: %s "
                + "clientCancelled: %s serverClosed: %s",
            LogUtils.truncateMessageLineLength(request), mClosed, mCanceled, mClosedFromRemote));
      }
      try (LockResource ignored = new LockResource(mLock)) {
        long startMs = System.currentTimeMillis();
        while (true) {
          checkError();
          if (mRequestObserver.isReady()) {
            break;
          }
  
          long waitedForMs = System.currentTimeMillis() - startMs;
          if (waitedForMs >= timeoutMs) {
            throw new DeadlineExceededException(formatErrorMessage(
                "Timeout sending request %s after %dms. clientClosed: %s clientCancelled: %s "
                    + "serverClosed: %s",
                LogUtils.truncateMessageLineLength(request), timeoutMs, mClosed, mCanceled,
                mClosedFromRemote));
          }
  
          try {
            // Wait for a minute max
            long awaitMs = Math.min(timeoutMs - waitedForMs, Constants.MINUTE_MS);
            if (!mReadyOrFailed.await(awaitMs, TimeUnit.MILLISECONDS)) {
              // Log a warning before looping again
              LOG.warn(
                  "Stream is not ready for client to send request, will wait again. totalWaitMs: {} "
                      + "clientClosed: {} clientCancelled: {} serverClosed: {} description: {}",
                  System.currentTimeMillis() - startMs, mClosed, mCanceled, mClosedFromRemote,
                  mDescription);
            }
          } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new CancelledException(
                formatErrorMessage("Failed to send request %s: interrupted while waiting for server.",
                    LogUtils.truncateMessageLineLength(request)), e);
          }
        }
      }
      mRequestObserver.onNext(request);
    }
  
    /**
     * Sends a request. Will not wait for the stream to be ready. If the stream is closed or cancelled
     * this method will return without an error.
     *
     * @param request the request
     * @throws IOException if any error occurs
     */
    public void send(ReqT request) throws IOException {
      if (mClosed || mCanceled || mClosedFromRemote) {
        if (LOG.isDebugEnabled()) {
          LOG.debug(
              "Failed to send request {}: stream is already closed or cancelled. clientClosed: {} "
                  + "clientCancelled: {} serverClosed: {} ({})",
              LogUtils.truncateMessageLineLength(request), mClosed, mCanceled, mClosedFromRemote,
              mDescription);
        }
        return;
      }
      try (LockResource ignored = new LockResource(mLock)) {
        checkError();
      }
      mRequestObserver.onNext(request);
    }
  
    /**
     * Receives a response from the server. Will wait until a response is received, or
     * throw an exception if times out.
     *
     * @param timeoutMs maximum time to wait before giving up and throwing
     *                  a {@link DeadlineExceededException}
     * @return the response message, or null if the inbound stream is completed
     * @throws IOException if any error occurs
     */
    public ResT receive(long timeoutMs) throws IOException {
      if (mCompleted) {
        return null;
      }
      if (mCanceled) {
        throw new CancelledException(formatErrorMessage("Stream is already canceled."));
      }
  
      long startMs = System.currentTimeMillis();
      while (true) {
        long waitedForMs = System.currentTimeMillis() - startMs;
        if (waitedForMs >= timeoutMs) {
          throw new DeadlineExceededException(formatErrorMessage(
              "Timeout waiting for response after %dms. clientClosed: %s clientCancelled: %s "
                  + "serverClosed: %s", timeoutMs, mClosed, mCanceled, mClosedFromRemote));
        }
  
        // Wait for a minute max
        long waitMs = Math.min(timeoutMs - waitedForMs, Constants.MINUTE_MS);
        try {
          Object response = mResponses.poll(waitMs, TimeUnit.MILLISECONDS);
          if (response == null) {
            checkError(); // The stream could have errored while we were waiting
            // Log a warning before looping again
            LOG.warn("Client did not receive message from stream, will wait again. totalWaitMs: {} "
                    + "clientClosed: {} clientCancelled: {} serverClosed: {} description: {}",
                System.currentTimeMillis() - startMs, mClosed, mCanceled, mClosedFromRemote,
                mDescription);
            continue;
          }
          if (response == mResponseObserver) {
            mCompleted = true;
            return null;
          }
          checkError();
          return (ResT) response;
        } catch (InterruptedException e) {
          Thread.currentThread().interrupt();
          throw new CancelledException(
              formatErrorMessage("Interrupted while waiting for response."), e);
        }
      }
    }
  
    /**
     * Closes the outbound stream. If the stream is already closed then invoking this method has no
     * effect.
     */
    public void close() {
      if (isOpen()) {
        LOG.debug("Closing stream ({})", mDescription);
        mClosed = true;
        mRequestObserver.onCompleted();
      }
    }
  
    /**
     * Cancels the stream. If the stream is already cancelled then invoking this method has no
     * effect.
     */
    public void cancel() {
      if (isOpen()) {
        LOG.debug("Cancelling stream ({})", mDescription);
        mCanceled = true;
        mRequestObserver.cancel("Request is cancelled by user.", null);
      }
    }
  
    /**
     * Wait for server to complete the inbound stream.
     *
     * @param timeoutMs maximum time to wait for server response
     * @return the last response of the stream
     */
    public Optional<ResT> waitForComplete(long timeoutMs) throws IOException {
      if (mCompleted || mCanceled) {
        return Optional.empty();
      }
      ResT prevResponse;
      ResT response = null;
      do {
        // wait until inbound stream is closed from server.
        prevResponse = response;
        response = receive(timeoutMs);
      } while (response != null);
      return Optional.ofNullable(prevResponse);
    }
  
    /**
     * @return whether the stream is open
     */
    public boolean isOpen() {
      try (LockResource ignored = new LockResource(mLock)) {
        return !mClosed && !mCanceled && mError == null;
      }
    }
  
    /**
     * @return whether the stream is closed
     */
    public boolean isClosed() {
      return mClosed;
    }
  
    /**
     * @return whether the stream is canceled
     */
    public boolean isCanceled() {
      return mCanceled;
    }
  
    private void checkError() throws IOException {
      try (LockResource ignored = new LockResource(mLock)) {
        if (mError != null) {
          // prevents rethrowing the same error
          mCanceled = true;
          throw toAlluxioStatusException(mError);
        }
      }
    }
  
    private AlluxioStatusException toAlluxioStatusException(Throwable t) {
      AlluxioStatusException ex;
      if (t instanceof StatusRuntimeException) {
        ex = AlluxioStatusException.fromStatusRuntimeException((StatusRuntimeException) t);
        if (ex.getStatusCode() == Status.Code.CANCELLED) {
          // Streams are canceled when server is shutdown. Convert it to UnavailableException for
          // client to retry.
          ex = new UnavailableException(formatErrorMessage("Stream is canceled by server."), ex);
        }
      } else {
        ex = AlluxioStatusException.fromThrowable(mError);
      }
      // attaches description to the exception while maintaining the cause
      return AlluxioStatusException
          .from(ex.getStatus().withDescription(formatErrorMessage(ex.getMessage())));
    }
  
    private String formatErrorMessage(String format, Object... args) {
      return (format == null ? "Unknown error" : String.format(format, args))
          + String.format(" (%s)", mDescription);
    }
  
    private final class ResponseStreamObserver
        implements ClientResponseObserver<ReqT, ResT> {
  
      @Override
      public void onNext(ResT response) {
        try {
          mResponses.put(response);
        } catch (InterruptedException e) {
          handleInterruptedException(e);
        }
      }
  
      @Override
      public void onError(Throwable t) {
        try (LockResource ignored = new LockResource(mLock)) {
          LOG.warn("Received error on stream ({})", mDescription, t);
          updateException(t);
          mReadyOrFailed.signal();
        }
      }
  
      @Override
      public void onCompleted() {
        try {
          LOG.debug("Received completed event for stream ({})", mDescription);
          mResponses.put(this);
          mClosedFromRemote = true;
        } catch (InterruptedException e) {
          handleInterruptedException(e);
        }
      }
  
      @Override
      public void beforeStart(ClientCallStreamObserver<ReqT> requestStream) {
        requestStream.setOnReadyHandler(() -> {
          try (LockResource ignored = new LockResource(mLock)) {
            mReadyOrFailed.signal();
          }
        });
      }
  
      private void handleInterruptedException(InterruptedException e) {
        Thread.currentThread().interrupt();
        try (LockResource ignored = new LockResource(mLock)) {
          updateException(e);
        }
        throw new RuntimeException(e);
      }
  
      /**
       * Updates the channel exception to be the given exception e, or adds e to
       * suppressed exceptions.
       *
       * @param e Exception received
       */
      @GuardedBy("mLock")
      private void updateException(Throwable e) {
        if (mError == null || mError == e) {
          mError = e;
          mResponses.offer(e);
        } else {
          mError.addSuppressed(e);
        }
      }
    }
  }
  