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.cassandra.net;
  
  import java.io.IOException;
  import java.nio.channels.WritableByteChannel;
  import java.util.concurrent.locks.LockSupport;
  
  import io.netty.channel.Channel;
  import io.netty.channel.ChannelPromise;
  import org.apache.cassandra.io.util.BufferedDataOutputStreamPlus;
  import org.apache.cassandra.io.util.DataOutputStreamPlus;
  
  import static java.lang.Math.max;
  
  /**
   * A {@link DataOutputStreamPlus} that writes ASYNCHRONOUSLY to a Netty Channel.
   *
   * The close() and flush() methods synchronously wait for pending writes, and will propagate any exceptions
   * encountered in writing them to the wire.
   *
   * The correctness of this class depends on the ChannelPromise we create against a Channel always being completed,
   * which appears to be a guarantee provided by Netty so long as the event loop is running.
   *
   * There are two logical threads accessing the state in this class: the eventLoop of the channel, and the writer
   * (the writer thread may change, so long as only one utilises the class at any time).
   * Each thread has exclusive write access to certain state in the class, with the other thread only viewing the state,
   * simplifying concurrency considerations.
   */
  public abstract class AsyncChannelOutputPlus extends BufferedDataOutputStreamPlus
  {
      public static class FlushException extends IOException
      {
          public FlushException(String message)
          {
              super(message);
          }
  
          public FlushException(String message, Throwable cause)
          {
              super(message, cause);
          }
      }
  
      final Channel channel;
  
      /** the number of bytes we have begun flushing; updated only by writer */
      private volatile long flushing;
      /** the number of bytes we have finished flushing, successfully or otherwise; updated only by eventLoop */
      private volatile long flushed;
      /** the number of bytes we have finished flushing to the network; updated only by eventLoop */
      private          long flushedToNetwork;
      /** any error that has been thrown during a flush; updated only by eventLoop */
      private volatile Throwable flushFailed;
  
      /**
       * state for pausing until flushing has caught up - store the number of bytes we need to be flushed before
       * we should be signalled, and store ourselves in {@link #waiting}; once the flushing thread exceeds this many
       * total bytes flushed, any Thread stored in waiting will be signalled.
       *
       * This works exactly like using a WaitQueue, except that we only need to manage a single waiting thread.
       */
      private volatile long signalWhenFlushed; // updated only by writer
      private volatile Thread waiting; // updated only by writer
  
      public AsyncChannelOutputPlus(Channel channel)
      {
          super(null, null);
          this.channel = channel;
      }
  
      /**
       * Create a ChannelPromise for a flush of the given size.
       * <p>
       * This method will not return until the write is permitted by the provided watermarks and in flight bytes,
       * and on its completion will mark the requested bytes flushed.
       * <p>
       * If this method returns normally, the ChannelPromise MUST be writtenAndFlushed, or else completed exceptionally.
       */
      protected ChannelPromise beginFlush(long byteCount, long lowWaterMark, long highWaterMark) throws IOException
      {
          waitForSpace(byteCount, lowWaterMark, highWaterMark);
  
          return AsyncChannelPromise.withListener(channel, future -> {
              if (future.isSuccess() && null == flushFailed)
              {
                  flushedToNetwork += byteCount;
                  releaseSpace(byteCount);
              }
              else if (null == flushFailed)
              {
                  Throwable cause = future.cause();
                  if (cause == null)
                  {
                      cause = new FlushException("Flush failed for unknown reason");
                      cause.fillInStackTrace();
                  }
                  flushFailed = cause;
                  releaseSpace(flushing - flushed);
              }
              else
              {
                  assert flushing == flushed;
              }
          });
      }
  
      /**
       * Imposes our lowWaterMark/highWaterMark constraints, and propagates any exceptions thrown by prior flushes.
       *
       * If we currently have lowWaterMark or fewer bytes flushing, we are good to go.
       * If our new write will not take us over our highWaterMark, we are good to go.
       * Otherwise, we wait until either of these conditions are met.
       *
       * This may only be invoked by the writer thread, never by the eventLoop.
       *
       * @throws IOException if a prior asynchronous flush failed
       */
      private void waitForSpace(long bytesToWrite, long lowWaterMark, long highWaterMark) throws IOException
      {
          // decide when we would be willing to carry on writing
          // we are always writable if we have lowWaterMark or fewer bytes, no matter how many bytes we are flushing
          // our callers should not be supplying more than (highWaterMark - lowWaterMark) bytes, but we must work correctly if they do
          long wakeUpWhenFlushing = highWaterMark - bytesToWrite;
          waitUntilFlushed(max(lowWaterMark, wakeUpWhenFlushing), lowWaterMark);
          flushing += bytesToWrite;
      }
  
      /**
       * Implementation of waitForSpace, which calculates what flushed points we need to wait for,
       * parks if necessary and propagates flush failures.
       *
       * This may only be invoked by the writer thread, never by the eventLoop.
       */
      void waitUntilFlushed(long wakeUpWhenExcessBytesWritten, long signalWhenExcessBytesWritten) throws IOException
      {
          // we assume that we are happy to wake up at least as early as we will be signalled; otherwise we will never exit
          assert signalWhenExcessBytesWritten <= wakeUpWhenExcessBytesWritten;
          // flushing shouldn't change during this method invocation, so our calculations for signal and flushed are consistent
          long wakeUpWhenFlushed = flushing - wakeUpWhenExcessBytesWritten;
          if (flushed < wakeUpWhenFlushed)
              parkUntilFlushed(wakeUpWhenFlushed, flushing - signalWhenExcessBytesWritten);
          propagateFailedFlush();
      }
  
      /**
       * Utility method for waitUntilFlushed, which actually parks the current thread until the necessary
       * number of bytes have been flushed
       *
       * This may only be invoked by the writer thread, never by the eventLoop.
       */
      protected void parkUntilFlushed(long wakeUpWhenFlushed, long signalWhenFlushed)
      {
          assert wakeUpWhenFlushed <= signalWhenFlushed;
          assert waiting == null;
          this.waiting = Thread.currentThread();
          this.signalWhenFlushed = signalWhenFlushed;
  
          while (flushed < wakeUpWhenFlushed)
              LockSupport.park();
          waiting = null;
      }
  
      /**
       * Update our flushed count, and signal any waiters.
       *
       * This may only be invoked by the eventLoop, never by the writer thread.
       */
      protected void releaseSpace(long bytesFlushed)
      {
          long newFlushed = flushed + bytesFlushed;
          flushed = newFlushed;
  
          Thread thread = waiting;
          if (thread != null && signalWhenFlushed <= newFlushed)
              LockSupport.unpark(thread);
      }
  
      private void propagateFailedFlush() throws IOException
      {
          Throwable t = flushFailed;
          if (t != null)
          {
              if (SocketFactory.isCausedByConnectionReset(t))
                  throw new FlushException("The channel this output stream was writing to has been closed", t);
              throw new FlushException("This output stream is in an unsafe state after an asynchronous flush failed", t);
          }
      }
  
      @Override
      abstract protected void doFlush(int count) throws IOException;
  
      abstract public long position();
  
      public long flushed()
      {
          // external flushed (that which has had flush() invoked implicitly or otherwise) == internal flushing
          return flushing;
      }
  
      public long flushedToNetwork()
      {
          return flushedToNetwork;
      }
  
      /**
       * Perform an asynchronous flush, then waits until all outstanding flushes have completed
       *
       * @throws IOException if any flush fails
       */
      @Override
      public void flush() throws IOException
      {
          doFlush(0);
          waitUntilFlushed(0, 0);
      }
  
      /**
       * Flush any remaining writes, and release any buffers.
       *
       * The channel is not closed, as it is assumed to be managed externally.
       *
       * WARNING: This method requires mutual exclusivity with all other producer methods to run safely.
       * It should only be invoked by the owning thread, never the eventLoop; the eventLoop should propagate
       * errors to {@link #flushFailed}, which will propagate them to the producer thread no later than its
       * final invocation to {@link #close()} or {@link #flush()} (that must not be followed by any further writes).
       */
      @Override
      public void close() throws IOException
      {
          try
          {
              flush();
          }
          finally
          {
              discard();
          }
      }
  
      /**
       * Discard any buffered data, and the buffers that contain it.
       * May be invoked instead of {@link #close()} if we terminate exceptionally.
       */
      public abstract void discard();
  
      @Override
      protected WritableByteChannel newDefaultChannel()
      {
          throw new UnsupportedOperationException();
      }
  
  }
  