Java Code Examples for io.netty.buffer.ByteBuf#nioBuffers()

/**
 * Optimized accessor for obtaining the underlying NIO buffers for a Netty {@link ByteBuf}. Based
 * on code from Netty's {@code SslHandler}. This method returns NIO buffers that span the readable
 * region of the {@link ByteBuf}.
 */
private static ByteBuffer[] nioBuffers(ByteBuf buf, ByteBuffer[] singleBuffer) {
  // As CompositeByteBuf.nioBufferCount() can be expensive (as it needs to check all composed
  // ByteBuf to calculate the count) we will just assume a CompositeByteBuf contains more than 1
  // ByteBuf. The worst that can happen is that we allocate an extra ByteBuffer[] in
  // CompositeByteBuf.nioBuffers() which is better than walking the composed ByteBuf in most
  // cases.
  if (!(buf instanceof CompositeByteBuf) && buf.nioBufferCount() == 1) {
    // We know its only backed by 1 ByteBuffer so use internalNioBuffer to keep object
    // allocation to a minimum.
    singleBuffer[0] = buf.internalNioBuffer(buf.readerIndex(), buf.readableBytes());
    return singleBuffer;
  }

  return buf.nioBuffers();
}
 

/**
 * Write bytes form the given {@link ByteBuf} to the underlying {@link java.nio.channels.Channel}.
 * @param in the collection which contains objects to write.
 * @param buf the {@link ByteBuf} from which the bytes should be written
 * @return The value that should be decremented from the write quantum which starts at
 * {@link ChannelConfig#getWriteSpinCount()}. The typical use cases are as follows:
 * <ul>
 *     <li>0 - if no write was attempted. This is appropriate if an empty {@link ByteBuf} (or other empty content)
 *     is encountered</li>
 *     <li>1 - if a single call to write data was made to the OS</li>
 *     <li>{@link ChannelUtils#WRITE_STATUS_SNDBUF_FULL} - if an attempt to write data was made to the OS, but no
 *     data was accepted</li>
 * </ul>
 */
private int writeBytes(ChannelOutboundBuffer in, ByteBuf buf) throws Exception {
    int readableBytes = buf.readableBytes();
    if (readableBytes == 0) {
        in.remove();
        return 0;
    }

    if (buf.hasMemoryAddress() || buf.nioBufferCount() == 1) {
        return doWriteBytes(in, buf);
    } else {
        ByteBuffer[] nioBuffers = buf.nioBuffers();
        return writeBytesMultiple(in, nioBuffers, nioBuffers.length, readableBytes,
                config().getMaxBytesPerGatheringWrite());
    }
}
 

/**
 * Write bytes form the given {@link ByteBuf} to the underlying {@link java.nio.channels.Channel}.
 * @param in the collection which contains objects to write.
 * @param buf the {@link ByteBuf} from which the bytes should be written
 * @return The value that should be decremented from the write quantum which starts at
 * {@link ChannelConfig#getWriteSpinCount()}. The typical use cases are as follows:
 * <ul>
 *     <li>0 - if no write was attempted. This is appropriate if an empty {@link ByteBuf} (or other empty content)
 *     is encountered</li>
 *     <li>1 - if a single call to write data was made to the OS</li>
 *     <li>{@link ChannelUtils#WRITE_STATUS_SNDBUF_FULL} - if an attempt to write data was made to the OS, but
 *     no data was accepted</li>
 * </ul>
 */
private int writeBytes(ChannelOutboundBuffer in, ByteBuf buf) throws Exception {
    int readableBytes = buf.readableBytes();
    if (readableBytes == 0) {
        in.remove();
        return 0;
    }

    if (buf.hasMemoryAddress() || buf.nioBufferCount() == 1) {
        return doWriteBytes(in, buf);
    } else {
        ByteBuffer[] nioBuffers = buf.nioBuffers();
        return writeBytesMultiple(in, nioBuffers, nioBuffers.length, readableBytes,
                config().getMaxBytesPerGatheringWrite());
    }
}
 

/**
 * Write bytes form the given {@link ByteBuf} to the underlying {@link java.nio.channels.Channel}.
 * @param buf           the {@link ByteBuf} from which the bytes should be written
 */
private boolean writeBytes(ChannelOutboundBuffer in, ByteBuf buf, int writeSpinCount) throws Exception {
    int readableBytes = buf.readableBytes();
    if (readableBytes == 0) {
        in.remove();
        return true;
    }

    if (buf.hasMemoryAddress() || buf.nioBufferCount() == 1) {
        int writtenBytes = doWriteBytes(buf, writeSpinCount);
        in.removeBytes(writtenBytes);
        return writtenBytes == readableBytes;
    } else {
        ByteBuffer[] nioBuffers = buf.nioBuffers();
        return writeBytesMultiple(in, nioBuffers, nioBuffers.length, readableBytes, writeSpinCount);
    }
}
 

/**
 * Optimized accessor for obtaining the underlying NIO buffers for a Netty {@link ByteBuf}. Based
 * on code from Netty's {@code SslHandler}. This method returns NIO buffers that span the readable
 * region of the {@link ByteBuf}.
 */
private static ByteBuffer[] nioBuffers(ByteBuf buf, ByteBuffer[] singleBuffer) {
  // As CompositeByteBuf.nioBufferCount() can be expensive (as it needs to check all composed
  // ByteBuf to calculate the count) we will just assume a CompositeByteBuf contains more than 1
  // ByteBuf. The worst that can happen is that we allocate an extra ByteBuffer[] in
  // CompositeByteBuf.nioBuffers() which is better than walking the composed ByteBuf in most
  // cases.
  if (!(buf instanceof CompositeByteBuf) && buf.nioBufferCount() == 1) {
    // We know its only backed by 1 ByteBuffer so use internalNioBuffer to keep object
    // allocation to a minimum.
    singleBuffer[0] = buf.internalNioBuffer(buf.readerIndex(), buf.readableBytes());
    return singleBuffer;
  }

  return buf.nioBuffers();
}
 

/**
 * Ensure that the data filled in is exactly identical to the original content.
 */
private void assertDataIdentity(ClusterMap clusterMap) throws IOException {
  if (!testEncryption) {
    ByteBuffer dest = ByteBuffer.allocate(totalSizeWritten);
    for (ByteBuf buf : compositeBuffers) {
      Assert.assertNotNull("All chunks should have come in", buf);
      for (ByteBuffer buffer: buf.nioBuffers()) {
        dest.put(buffer);
      }
    }
    Assert.assertTrue("Filled chunk contents must exactly match the input buffer ",
        Arrays.equals(putContent, dest.array()));
  } else {
    byte[] content = new byte[blobSize];
    AtomicInteger offset = new AtomicInteger(0);
    for (int i = 0; i < numChunks; i++) {
      DecryptJob decryptJob =
          new DecryptJob(compositeBlobIds[i], compositeEncryptionKeys[i], compositeBuffers[i], null, cryptoService,
              kms, new CryptoJobMetricsTracker(routerMetrics.decryptJobMetrics), (result, exception) -> {
            Assert.assertNull("Exception shouldn't have been thrown", exception);
            int chunkSize = result.getDecryptedBlobContent().remaining();
            result.getDecryptedBlobContent().get(content, offset.get(), chunkSize);
            offset.addAndGet(chunkSize);
          });
      decryptJob.run();
    }
    Assert.assertTrue("Filled chunk contents must exactly match the input buffer ",
        Arrays.equals(putContent, content));
  }
}
 

private SSLEngineResult wrap(ByteBufAllocator alloc, SSLEngine engine, ByteBuf in, ByteBuf out)
        throws SSLException {
    ByteBuf newDirectIn = null;
    try {
        int readerIndex = in.readerIndex();
        int readableBytes = in.readableBytes();

        // We will call SslEngine.wrap(ByteBuffer[], ByteBuffer) to allow efficient handling of
        // CompositeByteBuf without force an extra memory copy when CompositeByteBuffer.nioBuffer() is called.
        final ByteBuffer[] in0;
        if (in.isDirect() || !engineType.wantsDirectBuffer) {
            // As CompositeByteBuf.nioBufferCount() can be expensive (as it needs to check all composed ByteBuf
            // to calculate the count) we will just assume a CompositeByteBuf contains more then 1 ByteBuf.
            // The worst that can happen is that we allocate an extra ByteBuffer[] in CompositeByteBuf.nioBuffers()
            // which is better then walking the composed ByteBuf in most cases.
            if (!(in instanceof CompositeByteBuf) && in.nioBufferCount() == 1) {
                in0 = singleBuffer;
                // We know its only backed by 1 ByteBuffer so use internalNioBuffer to keep object allocation
                // to a minimum.
                in0[0] = in.internalNioBuffer(readerIndex, readableBytes);
            } else {
                in0 = in.nioBuffers();
            }
        } else {
            // We could even go further here and check if its a CompositeByteBuf and if so try to decompose it and
            // only replace the ByteBuffer that are not direct. At the moment we just will replace the whole
            // CompositeByteBuf to keep the complexity to a minimum
            newDirectIn = alloc.directBuffer(readableBytes);
            newDirectIn.writeBytes(in, readerIndex, readableBytes);
            in0 = singleBuffer;
            in0[0] = newDirectIn.internalNioBuffer(newDirectIn.readerIndex(), readableBytes);
        }

        for (;;) {
            ByteBuffer out0 = out.nioBuffer(out.writerIndex(), out.writableBytes());
            SSLEngineResult result = engine.wrap(in0, out0);
            in.skipBytes(result.bytesConsumed());
            out.writerIndex(out.writerIndex() + result.bytesProduced());

            switch (result.getStatus()) {
            case BUFFER_OVERFLOW:
                out.ensureWritable(engine.getSession().getPacketBufferSize());
                break;
            default:
                return result;
            }
        }
    } finally {
        // Null out to allow GC of ByteBuffer
        singleBuffer[0] = null;

        if (newDirectIn != null) {
            newDirectIn.release();
        }
    }
}
 

private SSLEngineResult wrap(ByteBufAllocator alloc, SSLEngine engine, ByteBuf in, ByteBuf out)
        throws SSLException {
    ByteBuf newDirectIn = null;
    try {
        int readerIndex = in.readerIndex();
        int readableBytes = in.readableBytes();

        // We will call SslEngine.wrap(ByteBuffer[], ByteBuffer) to allow efficient handling of
        // CompositeByteBuf without force an extra memory copy when CompositeByteBuffer.nioBuffer() is called.
        final ByteBuffer[] in0;
        if (in.isDirect() || !wantsDirectBuffer) {
            // As CompositeByteBuf.nioBufferCount() can be expensive (as it needs to check all composed ByteBuf
            // to calculate the count) we will just assume a CompositeByteBuf contains more then 1 ByteBuf.
            // The worst that can happen is that we allocate an extra ByteBuffer[] in CompositeByteBuf.nioBuffers()
            // which is better then walking the composed ByteBuf in most cases.
            if (!(in instanceof CompositeByteBuf) && in.nioBufferCount() == 1) {
                in0 = singleBuffer;
                // We know its only backed by 1 ByteBuffer so use internalNioBuffer to keep object allocation
                // to a minimum.
                in0[0] = in.internalNioBuffer(readerIndex, readableBytes);
            } else {
                in0 = in.nioBuffers();
            }
        } else {
            // We could even go further here and check if its a CompositeByteBuf and if so try to decompose it and
            // only replace the ByteBuffer that are not direct. At the moment we just will replace the whole
            // CompositeByteBuf to keep the complexity to a minimum
            newDirectIn = alloc.directBuffer(readableBytes);
            newDirectIn.writeBytes(in, readerIndex, readableBytes);
            in0 = singleBuffer;
            in0[0] = newDirectIn.internalNioBuffer(0, readableBytes);
        }

        for (;;) {
            ByteBuffer out0 = out.nioBuffer(out.writerIndex(), out.writableBytes());
            SSLEngineResult result = engine.wrap(in0, out0);
            in.skipBytes(result.bytesConsumed());
            out.writerIndex(out.writerIndex() + result.bytesProduced());

            switch (result.getStatus()) {
            case BUFFER_OVERFLOW:
                out.ensureWritable(maxPacketBufferSize);
                break;
            default:
                return result;
            }
        }
    } finally {
        // Null out to allow GC of ByteBuffer
        singleBuffer[0] = null;

        if (newDirectIn != null) {
            newDirectIn.release();
        }
    }
}
 

/**
 * Write data in {@code src} to the channel and the {@code callback} will be invoked once the write succeeds or fails.
 * This method is the counterpart for {@link ByteBuf}. It shares the same guarantee as the {@link #write(ByteBuffer, Callback)}.
 * Whoever implements this interface, shouldn't release the {@code src}. If releasing is expected after finishing writing
 * to channel, please release this {@link ByteBuf} in the callback method.
 * @param src The data taht needs to be written to the channel.
 * @param callback The {@link Callback} that will be invoked once the write succeeds/fails. This can be null.
 * @return a {@link Future} that will eventually contain the result of the write operation (the number of bytes
 *         written).
 */
default Future<Long> write(ByteBuf src, Callback<Long> callback) {
  if (src == null) {
    throw new IllegalArgumentException("Source ByteBuf cannot be null");
  }
  int numBuffers = src.nioBufferCount();
  FutureResult<Long> futureResult = new FutureResult<>();
  Callback<Long> singleBufferCallback = (result, exception) -> {
    if (result != 0) {
      src.readerIndex(src.readerIndex() + (int) result.longValue());
    }
    futureResult.done(result, exception);
    if (callback != null) {
      callback.onCompletion(result, exception);
    }
  };
  if (numBuffers < 1) {
    ByteBuffer byteBuffer = ByteBuffer.allocate(src.readableBytes());
    src.getBytes(src.readerIndex(), byteBuffer);
    write(byteBuffer, singleBufferCallback);
  } else if (numBuffers == 1) {
    write(src.nioBuffer(), singleBufferCallback);
  } else {
    ByteBuffer[] buffers = src.nioBuffers();
    AtomicLong size = new AtomicLong(0);
    AtomicInteger index = new AtomicInteger(0);
    AtomicBoolean callbackInvoked = new AtomicBoolean(false);
    Callback<Long> cb = (result, exception) -> {
      index.addAndGet(1);
      size.addAndGet(result);
      if (result != 0) {
        src.readerIndex(src.readerIndex() + (int) result.longValue());
      }
      if ((exception != null || index.get() == buffers.length) && callbackInvoked.compareAndSet(false, true)) {
        futureResult.done(size.get(), exception);
        if (callback != null) {
          callback.onCompletion(size.get(), exception);
        }
      }
    };
    for (int i = 0; i < buffers.length; i++) {
      write(buffers[i], cb);
    }
  }
  return futureResult;
}