| /* |
| * 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.ByteBuffer; |
| import java.nio.channels.ClosedChannelException; |
| import java.util.concurrent.BlockingQueue; |
| |
| import com.google.common.annotations.VisibleForTesting; |
| import com.google.common.primitives.Ints; |
| |
| import io.netty.buffer.ByteBuf; |
| import io.netty.buffer.ByteBufAllocator; |
| import io.netty.buffer.Unpooled; |
| import io.netty.channel.Channel; |
| import org.apache.cassandra.io.util.RebufferingInputStream; |
| import org.apache.cassandra.streaming.StreamingDataInputPlus; |
| |
| import static org.apache.cassandra.utils.concurrent.BlockingQueues.newBlockingQueue; |
| |
| /* |
| * This class expects a single producer (Netty event loop) and single consumer thread (StreamingDeserializerTask). |
| */ |
| public class AsyncStreamingInputPlus extends RebufferingInputStream implements StreamingDataInputPlus |
| { |
| private final Channel channel; |
| |
| /** |
| * The parent, or owning, buffer of the current buffer being read from ({@link super#buffer}). |
| */ |
| private ByteBuf currentBuf; |
| |
| private final BlockingQueue<ByteBuf> queue; |
| |
| private boolean isProducerClosed = false; |
| private boolean isConsumerClosed = false; |
| |
| public AsyncStreamingInputPlus(Channel channel) |
| { |
| super(Unpooled.EMPTY_BUFFER.nioBuffer()); |
| currentBuf = Unpooled.EMPTY_BUFFER; |
| |
| queue = newBlockingQueue(); |
| |
| this.channel = channel; |
| channel.config().setAutoRead(false); |
| } |
| |
| /** |
| * Append a {@link ByteBuf} to the end of the einternal queue. |
| * |
| * Note: it's expected this method is invoked on the netty event loop. |
| */ |
| public boolean append(ByteBuf buf) throws IllegalStateException |
| { |
| if (isProducerClosed) |
| return false; // buf should be released in NettyStreamingChannel.channelRead |
| |
| queue.add(buf); |
| |
| return true; |
| } |
| |
| /** |
| * {@inheritDoc} |
| * |
| * Release open buffers and poll the {@link #queue} for more data. |
| * <p> |
| * This is invoked on a consuming thread (not the event loop) |
| * because if we block on the queue we can't fill it on the event loop (as that's where the buffers are coming from). |
| * |
| * @throws ClosedChannelException when no further reading from this instance should occur. Implies this instance is closed. |
| */ |
| @Override |
| protected void reBuffer() throws ClosedChannelException |
| { |
| if (isConsumerClosed) |
| throw new ClosedChannelException(); |
| |
| if (queue.isEmpty()) |
| channel.read(); |
| |
| currentBuf.release(); |
| currentBuf = null; |
| buffer = null; |
| |
| ByteBuf next = null; |
| do |
| { |
| try |
| { |
| next = queue.take(); // rely on sentinel being sent to terminate this loop |
| } |
| catch (InterruptedException ie) |
| { |
| // ignore interruptions, retry and rely on being shut down by requestClosure |
| } |
| } while (next == null); |
| |
| if (next == Unpooled.EMPTY_BUFFER) // the indicator that the input is closed |
| { |
| isConsumerClosed = true; |
| throw new ClosedChannelException(); |
| } |
| |
| currentBuf = next; |
| buffer = next.nioBuffer(); |
| } |
| |
| public interface Consumer |
| { |
| int accept(ByteBuffer buffer) throws IOException; |
| } |
| |
| /** |
| * Consumes bytes in the stream until the given length |
| */ |
| public void consume(Consumer consumer, long length) throws IOException |
| { |
| while (length > 0) |
| { |
| if (!buffer.hasRemaining()) |
| reBuffer(); |
| |
| final int position = buffer.position(); |
| final int limit = buffer.limit(); |
| |
| buffer.limit(position + (int) Math.min(length, limit - position)); |
| try |
| { |
| int copied = consumer.accept(buffer); |
| buffer.position(position + copied); |
| length -= copied; |
| } |
| finally |
| { |
| buffer.limit(limit); |
| } |
| } |
| } |
| |
| /** |
| * {@inheritDoc} |
| * |
| * As long as this method is invoked on the consuming thread the returned value will be accurate. |
| */ |
| @VisibleForTesting |
| public int unsafeAvailable() |
| { |
| long count = buffer != null ? buffer.remaining() : 0; |
| for (ByteBuf buf : queue) |
| count += buf.readableBytes(); |
| |
| return Ints.checkedCast(count); |
| } |
| |
| public boolean isEmpty() |
| { |
| return isConsumerClosed || (queue.isEmpty() && (buffer == null || !buffer.hasRemaining())); |
| } |
| |
| /** |
| * {@inheritDoc} |
| * |
| * Note: This should invoked on the consuming thread. |
| */ |
| @Override |
| public void close() |
| { |
| if (isConsumerClosed) |
| return; |
| |
| isConsumerClosed = true; |
| |
| if (currentBuf != null) |
| { |
| currentBuf.release(); |
| currentBuf = null; |
| buffer = null; |
| } |
| |
| while (true) |
| { |
| try |
| { |
| ByteBuf buf = queue.take(); |
| if (buf == Unpooled.EMPTY_BUFFER) |
| break; |
| buf.release(); |
| } |
| catch (InterruptedException e) |
| { |
| // ignore and rely on requestClose having been called |
| } |
| } |
| } |
| |
| /** |
| * Mark this stream as closed, but do not release any of the resources. |
| * |
| * Note: this is best to be called from the producer thread. |
| */ |
| public void requestClosure() |
| { |
| if (!isProducerClosed) |
| { |
| queue.add(Unpooled.EMPTY_BUFFER); |
| isProducerClosed = true; |
| } |
| } |
| |
| // TODO: let's remove this like we did for AsyncChannelOutputPlus |
| public ByteBufAllocator getAllocator() |
| { |
| return channel.alloc(); |
| } |
| } |