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

/**
 * 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();
}
 

private AsyncFileWriterImpl doWrite(Buffer buffer, long position, Handler<AsyncResult<Void>> handler) {
    Preconditions.checkNotNull(buffer, "buffer");
    Preconditions.checkArgument(position >= 0, "position must be >= 0");
    Handler<AsyncResult<Void>> wrapped = ar -> {
        if (ar.succeeded()) {
            if (handler != null) {
                handler.handle(ar);
            }
        } else {
            if (handler != null) {
                handler.handle(ar);
            } else {
                handleException(ar.cause());
            }
        }
    };
    ByteBuf buf = buffer.getByteBuf();
    if (buf.nioBufferCount() > 1) {
        doWrite(buf.nioBuffers(), position, wrapped);
    } else {
        ByteBuffer bb = buf.nioBuffer();
        doWrite(bb, position, bb.limit(), wrapped);
    }
    return this;
}
 

@Override
protected final Object filterOutboundMessage(Object msg) throws Exception {
    if (msg instanceof SctpMessage) {
        SctpMessage m = (SctpMessage) msg;
        ByteBuf buf = m.content();
        if (buf.isDirect() && buf.nioBufferCount() == 1) {
            return m;
        }

        return new SctpMessage(m.protocolIdentifier(), m.streamIdentifier(), newDirectBuffer(m, buf));
    }

    throw new UnsupportedOperationException(
            "unsupported message type: " + StringUtil.simpleClassName(msg) +
            " (expected: " + StringUtil.simpleClassName(SctpMessage.class));
}
 

/**
 * 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);
    }
}
 

/**
 * 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());
    }
}
 

@Override
protected final Object filterOutboundMessage(Object msg) throws Exception {
    if (msg instanceof SctpMessage) {
        SctpMessage m = (SctpMessage) msg;
        ByteBuf buf = m.content();
        if (buf.isDirect() && buf.nioBufferCount() == 1) {
            return m;
        }

        return new SctpMessage(m.protocolIdentifier(), m.streamIdentifier(), m.isUnordered(),
                               newDirectBuffer(m, buf));
    }

    throw new UnsupportedOperationException(
            "unsupported message type: " + StringUtil.simpleClassName(msg) +
            " (expected: " + StringUtil.simpleClassName(SctpMessage.class));
}
 

/**
 * 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();
}
 

@Override
protected boolean doWriteMessage(Object msg, ChannelOutboundBuffer in) throws Exception {
    // expects a message
    final UdtMessage message = (UdtMessage) msg;

    final ByteBuf byteBuf = message.content();

    final int messageSize = byteBuf.readableBytes();

    final long writtenBytes;
    if (byteBuf.nioBufferCount() == 1) {
        writtenBytes = javaChannel().write(byteBuf.nioBuffer());
    } else {
        writtenBytes = javaChannel().write(byteBuf.nioBuffers());
    }

    // did not write the message
    if (writtenBytes <= 0 && messageSize > 0) {
        return false;
    }

    // wrote message completely
    if (writtenBytes != messageSize) {
        throw new Error(
                "Provider error: failed to write message. Provider library should be upgraded.");
    }

    return true;
}
 

@Override
public ByteBuf decrypt(ByteBuf toDecrypt, SecretKeySpec key) throws GeneralSecurityException {
  try {
    Cipher decrypter = Cipher.getInstance(GCM_CRYPTO_INSTANCE, "BC");
    byte[] iv = deserializeIV(new ByteBufInputStream(toDecrypt));
    decrypter.init(Cipher.DECRYPT_MODE, key, new IvParameterSpec(iv));
    int outputSize = decrypter.getOutputSize(toDecrypt.readableBytes());
    ByteBuf decryptedContent = ByteBufAllocator.DEFAULT.heapBuffer(outputSize);

    boolean toRelease = false;
    if (toDecrypt.nioBufferCount() != 1) {
      toRelease = true;
      ByteBuf temp = ByteBufAllocator.DEFAULT.heapBuffer(toDecrypt.readableBytes());
      temp.writeBytes(toDecrypt);
      toDecrypt = temp;
    }

    try {
      ByteBuffer toDecryptBuffer = toDecrypt.nioBuffer();
      ByteBuffer decryptedContentBuffer = decryptedContent.nioBuffer(0, outputSize);
      int n = decrypter.doFinal(toDecryptBuffer, decryptedContentBuffer);
      decryptedContent.writerIndex(decryptedContent.writerIndex() + n);
      return decryptedContent;
    } finally {
      if (toRelease) {
        toDecrypt.release();
      } else {
        toDecrypt.skipBytes(toDecrypt.readableBytes());
      }
    }
  } catch (Exception e) {
    throw new GeneralSecurityException("Exception thrown while decrypting data", e);
  }
}
 

@Override
public void flushBuffer(final ByteBuf byteBuf, final boolean requestedSync, final List<IOCallback> callbacks) {
   final int bytes = byteBuf.readableBytes();
   if (bytes > 0) {
      final boolean releaseBuffer;
      final ByteBuffer buffer;
      if (byteBuf.nioBufferCount() == 1) {
         //any ByteBuffer is fine with the MAPPED journal
         releaseBuffer = false;
         buffer = byteBuf.internalNioBuffer(byteBuf.readerIndex(), bytes);
      } else {
         //perform the copy on buffer
         releaseBuffer = true;
         buffer = factory.newBuffer(byteBuf.capacity());
         buffer.limit(bytes);
         byteBuf.getBytes(byteBuf.readerIndex(), buffer);
         buffer.flip();
      }
      try {
         blockingWriteDirect(buffer, requestedSync, releaseBuffer);
         IOCallback.done(callbacks);
      } catch (Throwable t) {
         final int code;
         if (t instanceof IOException) {
            code = ActiveMQExceptionType.IO_ERROR.getCode();
            factory.onIOError(new ActiveMQIOErrorException(t.getMessage(), t), t.getMessage(), TimedSequentialFile.this.sequentialFile);
         } else {
            code = ActiveMQExceptionType.GENERIC_EXCEPTION.getCode();
         }
         IOCallback.onError(callbacks, code, t.getMessage());
      }
   } else {
      IOCallback.done(callbacks);
   }
}
 

public static char[] readUtf8(ByteBuf byteBuf, int length) {
  CharsetDecoder charsetDecoder = CharsetUtil.UTF_8.newDecoder();
  int en = (int) (length * (double) charsetDecoder.maxCharsPerByte());
  char[] ca = new char[en];

  CharBuffer charBuffer = CharBuffer.wrap(ca);
  ByteBuffer byteBuffer =
      byteBuf.nioBufferCount() == 1
          ? byteBuf.internalNioBuffer(byteBuf.readerIndex(), length)
          : byteBuf.nioBuffer(byteBuf.readerIndex(), length);
  byteBuffer.mark();
  try {
    CoderResult cr = charsetDecoder.decode(byteBuffer, charBuffer, true);
    if (!cr.isUnderflow()) cr.throwException();
    cr = charsetDecoder.flush(charBuffer);
    if (!cr.isUnderflow()) cr.throwException();

    byteBuffer.reset();
    byteBuf.skipBytes(length);

    return safeTrim(charBuffer.array(), charBuffer.position());
  } catch (CharacterCodingException x) {
    // Substitution is always enabled,
    // so this shouldn't happen
    throw new IllegalStateException("unable to decode char array from the given buffer", x);
  }
}
 

@Override
protected boolean doWriteMessage(Object msg, ChannelOutboundBuffer in) throws Exception {
    SctpMessage packet = (SctpMessage) msg;
    ByteBuf data = packet.content();
    int dataLen = data.readableBytes();
    if (dataLen == 0) {
        return true;
    }

    ByteBufAllocator alloc = alloc();
    boolean needsCopy = data.nioBufferCount() != 1;
    if (!needsCopy) {
        if (!data.isDirect() && alloc.isDirectBufferPooled()) {
            needsCopy = true;
        }
    }
    ByteBuffer nioData;
    if (!needsCopy) {
        nioData = data.nioBuffer();
    } else {
        data = alloc.directBuffer(dataLen).writeBytes(data);
        nioData = data.nioBuffer();
    }
    final MessageInfo mi = MessageInfo.createOutgoing(association(), null, packet.streamIdentifier());
    mi.payloadProtocolID(packet.protocolIdentifier());
    mi.streamNumber(packet.streamIdentifier());

    final int writtenBytes = javaChannel().send(nioData, mi);
    return writtenBytes > 0;
}
 

@Override
protected boolean doWriteMessage(Object msg, ChannelOutboundBuffer in) throws Exception {
    // expects a message
    final UdtMessage message = (UdtMessage) msg;

    final ByteBuf byteBuf = message.content();

    final int messageSize = byteBuf.readableBytes();
    if (messageSize == 0) {
        return true;
    }

    final long writtenBytes;
    if (byteBuf.nioBufferCount() == 1) {
        writtenBytes = javaChannel().write(byteBuf.nioBuffer());
    } else {
        writtenBytes = javaChannel().write(byteBuf.nioBuffers());
    }

    // wrote message completely
    if (writtenBytes > 0 && writtenBytes != messageSize) {
        throw new Error(
                "Provider error: failed to write message. Provider library should be upgraded.");
    }

    return writtenBytes > 0;
}
 

@Override
protected void doWrite(ChannelOutboundBuffer in) throws Exception {
    if (!writeSelector.isOpen()) {
        return;
    }
    final int size = in.size();
    final int selectedKeys = writeSelector.select(SO_TIMEOUT);
    if (selectedKeys > 0) {
        final Set<SelectionKey> writableKeys = writeSelector.selectedKeys();
        if (writableKeys.isEmpty()) {
            return;
        }
        Iterator<SelectionKey> writableKeysIt = writableKeys.iterator();
        int written = 0;
        for (;;) {
            if (written == size) {
                // all written
                return;
            }
            writableKeysIt.next();
            writableKeysIt.remove();

            SctpMessage packet = (SctpMessage) in.current();
            if (packet == null) {
                return;
            }

            ByteBuf data = packet.content();
            int dataLen = data.readableBytes();
            ByteBuffer nioData;

            if (data.nioBufferCount() != -1) {
                nioData = data.nioBuffer();
            } else {
                nioData = ByteBuffer.allocate(dataLen);
                data.getBytes(data.readerIndex(), nioData);
                nioData.flip();
            }

            final MessageInfo mi = MessageInfo.createOutgoing(association(), null, packet.streamIdentifier());
            mi.payloadProtocolID(packet.protocolIdentifier());
            mi.streamNumber(packet.streamIdentifier());

            ch.send(nioData, mi);
            written ++;
            in.remove();

            if (!writableKeysIt.hasNext()) {
                return;
            }
        }
    }
}
 

private static ByteBuffer toByteBuffer(ByteBuf out, int index, int len) {
    return out.nioBufferCount() == 1 ? out.internalNioBuffer(index, len) :
            out.nioBuffer(index, len);
}
 

/**
 * Helper function to encrypt ByteBuf with the given key. When useFixedIv is true, don't use a random iv byte
 * array, using a all zero byte array instead. Only set it to be true in test.
 * @param toEncrypt {@link ByteBuf} that needs to be encrypted
 * @param key the secret key (of type T) to use to encrypt
 * @param iv If null, will create a random byte array serve as iv bytes.
 * @return the {@link ByteBuf} containing the encrypted content. Ensure the result has all
 * the information like the IV along with the encrypted content, in order to decrypt the content with a given key
 * @throws {@link GeneralSecurityException} on any exception with encryption
 */
public ByteBuf encrypt(ByteBuf toEncrypt, SecretKeySpec key, byte[] iv) throws GeneralSecurityException {
  try {
    Cipher encrypter = Cipher.getInstance(GCM_CRYPTO_INSTANCE, "BC");
    if (iv == null) {
      iv = new byte[ivValSize];
      random.nextBytes(iv);
    }
    encrypter.init(Cipher.ENCRYPT_MODE, key, new IvParameterSpec(iv));
    int outputSize = encrypter.getOutputSize(toEncrypt.readableBytes());

    // stick with heap memory for now so to compare with the java.nio.ByteBuffer.
    ByteBuf encryptedContent =
        ByteBufAllocator.DEFAULT.heapBuffer(IVRecord_Format_V1.getIVRecordSize(iv) + outputSize);
    IVRecord_Format_V1.serializeIVRecord(encryptedContent, iv);

    boolean toRelease = false;
    if (toEncrypt.nioBufferCount() != 1) {
      toRelease = true;
      ByteBuf temp = ByteBufAllocator.DEFAULT.heapBuffer(toEncrypt.readableBytes());
      temp.writeBytes(toEncrypt);
      toEncrypt = temp;
    }
    try {
      ByteBuffer toEncryptBuffer = toEncrypt.nioBuffer();
      ByteBuffer encryptedContentBuffer = encryptedContent.nioBuffer(encryptedContent.writerIndex(),
          encryptedContent.capacity() - encryptedContent.writerIndex());
      int n = encrypter.doFinal(toEncryptBuffer, encryptedContentBuffer);
      encryptedContent.writerIndex(encryptedContent.writerIndex() + n);
      return encryptedContent;
    } finally {
      if (toRelease) {
        toEncrypt.release();
      } else {
        toEncrypt.skipBytes(toEncrypt.readableBytes());
      }
    }
  } catch (Exception e) {
    throw new GeneralSecurityException("Exception thrown while encrypting data", e);
  }
}
 

private ByteBuffer safeNioBuffer(ByteBuf buffer) {
    return buffer.nioBufferCount() == 1 ? buffer.internalNioBuffer(buffer.readerIndex(), this.compressedLength)
            : buffer.nioBuffer(buffer.readerIndex(), this.compressedLength);
}
 

/**
 * Checks if the specified buffer is a direct buffer and is composed of a single NIO buffer.
 * (We check this because otherwise we need to make it a non-composite buffer.)
 */
private static boolean isSingleDirectBuffer(ByteBuf buf) {
    return buf.isDirect() && buf.nioBufferCount() == 1;
}
 

/**
 * Checks if the specified buffer is a direct buffer and is composed of a single NIO buffer.
 * (We check this because otherwise we need to make it a non-composite buffer.)
 */
private static boolean isSingleDirectBuffer(ByteBuf buf) {
    return buf.isDirect() && buf.nioBufferCount() == 1;
}
 

/**
 * Checks if the specified buffer is a direct buffer and is composed of a single NIO buffer.
 * (We check this because otherwise we need to make it a non-composite buffer.)
 */
private static boolean isSingleDirectBuffer(ByteBuf buf) {
    return buf.isDirect() && buf.nioBufferCount() == 1;
}