Java Code Examples for io.netty.buffer.ByteBufAllocator#buffer()

@Nonnull
@Override
public ByteBuf writeObject(ByteBufAllocator bufAllocator, @Nullable Object object) throws Exception {
    // 这里的测试结果是：拷贝字节数组，比先计算一次大小，再写入ByteBuf快，而且快很多。
    // 此外，即使使用输入输出流构造，其内部还是做了缓存(创建了字节数组)，因此一定要有自己的缓冲字节数组
    final byte[] localBuffer = BufferPool.allocateBuffer();
    try {
        // 写入字节数组缓存
        final CodedDataOutputStream outputStream = new CodedDataOutputStream(localBuffer);
        encodeObject(outputStream, object);

        // 写入byteBuf
        final ByteBuf buffer = bufAllocator.buffer(outputStream.writeIndex());
        buffer.writeBytes(localBuffer, 0, outputStream.writeIndex());
        return buffer;
    } finally {
        BufferPool.releaseBuffer(localBuffer);
    }
}
 

@Override
void decode(ByteBufAllocator alloc, ByteBuf headerBlock, SpdyHeadersFrame frame) throws Exception {
    if (headerBlock == null) {
        throw new NullPointerException("headerBlock");
    }
    if (frame == null) {
        throw new NullPointerException("frame");
    }

    if (cumulation == null) {
        decodeHeaderBlock(headerBlock, frame);
        if (headerBlock.isReadable()) {
            cumulation = alloc.buffer(headerBlock.readableBytes());
            cumulation.writeBytes(headerBlock);
        }
    } else {
        cumulation.writeBytes(headerBlock);
        decodeHeaderBlock(cumulation, frame);
        if (cumulation.isReadable()) {
            cumulation.discardReadBytes();
        } else {
            releaseBuffer();
        }
    }
}
 

static ByteBuf encodeBytes(ByteBufAllocator alloc, byte[] value) {
    int size = value.length;

    if (size == 0) {
        // It is zero of var int, not terminal.
        return alloc.buffer(Byte.BYTES).writeByte(0);
    }

    ByteBuf buf = alloc.buffer(VarIntUtils.varIntBytes(size) + size);

    try {
        VarIntUtils.writeVarInt(buf, size);
        return buf.writeBytes(value);
    } catch (Throwable e) {
        buf.release();
        throw e;
    }
}
 

@Override
public ByteBuf encodeHeader(ByteBufAllocator allocator) {
    ByteBuf buffer = allocator.buffer(4 + Trace.TRACING_HEADER_LENGTH);
    boolean release = true;
    try {
        // ttl: 4
        buffer.writeInt((int) getTTL());

        // tracing: 25
        CodecUtils.encodeTrace(getTracing(), buffer);
        release = false;
    } finally {
        if (release) {
            buffer.release();
        }
    }
    return buffer;
}
 

@Override
public ByteBuf encodeHeader(ByteBufAllocator allocator) {
    // Allocate new ByteBuf
    ByteBuf buffer = allocator.buffer(256);

    boolean release = true;
    try {
        // version:2
        buffer.writeShort(getVersion());

        // headers -> nh:2 (key~2 value~2){nh}
        CodecUtils.encodeHeaders(getHeaders(), buffer);
        release = false;
    } finally {
        if (release) {
            buffer.release();
        }
    }
    return buffer;
}
 

private static ByteBuf encodeMessageWithOnlySingleByteFixedHeaderAndPacketId(
        ByteBufAllocator byteBufAllocator,
        MqttMessage message) {
    MqttFixedHeader mqttFixedHeader = message.fixedHeader();
    MqttPacketIdVariableHeader variableHeader = (MqttPacketIdVariableHeader) message.variableHeader();
    int msgId = variableHeader.packetId();

    int variableHeaderBufferSize = 2; // variable part only has a message id
    int fixedHeaderBufferSize = 1 + getVariableLengthInt(variableHeaderBufferSize);
    ByteBuf buf = byteBufAllocator.buffer(fixedHeaderBufferSize + variableHeaderBufferSize);
    buf.writeByte(getFixedHeaderByte1(mqttFixedHeader));
    writeVariableLengthInt(buf, variableHeaderBufferSize);
    buf.writeShort(msgId);

    return buf;
}
 

public static ByteBuf mapToByteBuf(ByteBufAllocator allocator, Map<String, String> map) {
  if (map == null || map.isEmpty()) {
    return Unpooled.EMPTY_BUFFER;
  }

  ByteBuf byteBuf = allocator.buffer();

  for (Map.Entry<String, String> entry : map.entrySet()) {
    String key = entry.getKey();
    int keyLength = NumberUtils.requireUnsignedShort(ByteBufUtil.utf8Bytes(key));
    byteBuf.writeShort(keyLength);
    ByteBufUtil.reserveAndWriteUtf8(byteBuf, key, keyLength);

    String value = entry.getValue();
    int valueLength = NumberUtils.requireUnsignedShort(ByteBufUtil.utf8Bytes(value));
    byteBuf.writeShort(valueLength);
    ByteBufUtil.reserveAndWriteUtf8(byteBuf, value, keyLength);
  }

  return byteBuf;
}
 

@Override
public ByteBuf serializeRequestAsBinary(final RequestMessage requestMessage, final ByteBufAllocator allocator) throws SerializationException {
    ByteBuf encodedMessage = null;
    try {
        final byte[] header = obtainHeader();
        final byte[] payload = mapper.writeValueAsBytes(requestMessage);

        encodedMessage = allocator.buffer(header.length + payload.length);
        encodedMessage.writeBytes(header);
        encodedMessage.writeBytes(payload);

        return encodedMessage;
    } catch (Exception ex) {
        if (encodedMessage != null) ReferenceCountUtil.release(encodedMessage);

        logger.warn(String.format("Request [%s] could not be serialized by %s.", requestMessage, AbstractGraphSONMessageSerializerV2d0.class.getName()), ex);
        throw new SerializationException(ex);
    }
}
 

private static ByteBuf encodeSet(ByteBufAllocator alloc, Iterator<? extends CharSequence> iter, CodecContext context) {
    Charset charset = context.getClientCollation().getCharset();
    ByteBuf content = alloc.buffer();

    try {
        // Max size of var int, fill zero to protect memory data.
        VarIntUtils.reserveVarInt(content);

        CharSequence name = iter.next();
        int size = content.writeCharSequence(name, charset);

        while (iter.hasNext()) {
            name = iter.next();
            size += content.writeByte(',').writeCharSequence(name, charset) + 1;
        }

        return VarIntUtils.setReservedVarInt(content, size);
    } catch (Throwable e) {
        content.release();
        throw e;
    }
}
 

@Override
protected ByteBuf metadata(ByteBufAllocator alloc) {
    byte[] pathBytes = fileId.pathBytes();
    int metaLen = 3 * FastdfsConstants.FDFS_PROTO_PKG_LEN_SIZE + pathBytes.length;
    ByteBuf buf = alloc.buffer(metaLen);
    buf.writeLong(pathBytes.length);
    buf.writeLong(offset);
    buf.writeLong(size());
    buf.writeBytes(pathBytes);
    return buf;
}
 

static ByteBuf encodeBinary(ByteBufAllocator alloc, LocalDateTime value) {
    LocalTime time = value.toLocalTime();

    if (LocalTime.MIDNIGHT.equals(time)) {
        return LocalDateCodec.encodeDate(alloc, value.toLocalDate());
    }

    int nano = time.getNano();
    int bytes = nano > 0 ? DateTimes.MICRO_DATETIME_SIZE : DateTimes.DATETIME_SIZE;
    ByteBuf buf = alloc.buffer(Byte.BYTES + bytes);

    try {
        buf.writeByte(bytes)
            .writeShortLE(value.getYear())
            .writeByte(value.getMonthValue())
            .writeByte(value.getDayOfMonth())
            .writeByte(time.getHour())
            .writeByte(time.getMinute())
            .writeByte(time.getSecond());

        if (nano > 0) {
            return buf.writeIntLE(nano / DateTimes.NANOS_OF_MICRO);
        }

        return buf;
    } catch (Throwable e) {
        buf.release();
        throw e;
    }
}
 

private static ByteBuf encodePublishMessage(
        ByteBufAllocator byteBufAllocator,
        MqttPublishMessage message) {
    MqttFixedHeader mqttFixedHeader = message.fixedHeader();
    MqttPublishVariableHeader variableHeader = message.variableHeader();
    ByteBuf payload = message.payload().duplicate();

    String topicName = variableHeader.topicName();
    byte[] topicNameBytes = encodeStringUtf8(topicName);

    int variableHeaderBufferSize = 2 + topicNameBytes.length +
            (mqttFixedHeader.qos().value() > 0 ? 2 : 0);
    int payloadBufferSize = payload.readableBytes();
    int variablePartSize = variableHeaderBufferSize + payloadBufferSize;
    int fixedHeaderBufferSize = 1 + getVariableLengthInt(variablePartSize);

    ByteBuf buf = byteBufAllocator.buffer(fixedHeaderBufferSize + variablePartSize);
    buf.writeByte(getFixedHeaderByte1(mqttFixedHeader));
    writeVariableLengthInt(buf, variablePartSize);
    buf.writeShort(topicNameBytes.length);
    buf.writeBytes(topicNameBytes);
    if (mqttFixedHeader.qos().value() > 0) {
        buf.writeShort(variableHeader.packetId());
    }
    buf.writeBytes(payload);

    return buf;
}
 

@Override
public List<Object> encode(ByteBufAllocator alloc) {
    byte cmd = cmd();
    int length = FDFS_GROUP_LEN + fileId.pathBytes().length;
    ByteBuf buf = alloc.buffer(length + FDFS_HEAD_LEN);
    buf.writeLong(length);
    buf.writeByte(cmd);
    buf.writeByte(ERRNO_OK);
    writeFixLength(buf, fileId.group(), FDFS_GROUP_LEN);
    ByteBufUtil.writeUtf8(buf, fileId.path());
    return Collections.singletonList(buf);
}
 

/**
 * Encode a custom mime type and a metadata value length into a newly allocated {@link ByteBuf}.
 *
 * <p>This larger representation encodes the mime type representation's length on a single byte,
 * then the representation itself, then the unsigned metadata value length on 3 additional bytes.
 *
 * @param allocator the {@link ByteBufAllocator} to use to create the buffer.
 * @param customMime a custom mime type to encode.
 * @param metadataLength the metadata length to append to the buffer as an unsigned 24 bits
 *     integer.
 * @return the encoded mime and metadata length information
 */
static ByteBuf encodeMetadataHeader(
    ByteBufAllocator allocator, String customMime, int metadataLength) {
  ByteBuf metadataHeader = allocator.buffer(4 + customMime.length());
  // reserve 1 byte for the customMime length
  // /!\ careful not to read that first byte, which is random at this point
  int writerIndexInitial = metadataHeader.writerIndex();
  metadataHeader.writerIndex(writerIndexInitial + 1);

  // write the custom mime in UTF8 but validate it is all ASCII-compatible
  // (which produces the right result since ASCII chars are still encoded on 1 byte in UTF8)
  int customMimeLength = ByteBufUtil.writeUtf8(metadataHeader, customMime);
  if (!ByteBufUtil.isText(
      metadataHeader, metadataHeader.readerIndex() + 1, customMimeLength, CharsetUtil.US_ASCII)) {
    metadataHeader.release();
    throw new IllegalArgumentException("custom mime type must be US_ASCII characters only");
  }
  if (customMimeLength < 1 || customMimeLength > 128) {
    metadataHeader.release();
    throw new IllegalArgumentException(
        "custom mime type must have a strictly positive length that fits on 7 unsigned bits, ie 1-128");
  }
  metadataHeader.markWriterIndex();

  // go back to beginning and write the length
  // encoded length is one less than actual length, since 0 is never a valid length, which gives
  // wider representation range
  metadataHeader.writerIndex(writerIndexInitial);
  metadataHeader.writeByte(customMimeLength - 1);

  // go back to post-mime type and write the metadata content length
  metadataHeader.resetWriterIndex();
  NumberUtils.encodeUnsignedMedium(metadataHeader, metadataLength);

  return metadataHeader;
}
 

@Test
public void copiesAndReleasesPooledByteBuf() {
    ByteBufAllocator pooledAllocator = PooledByteBufAllocator.DEFAULT;
    ByteBufAllocator unpooledAllocator = UnpooledByteBufAllocator.DEFAULT;
    CopyByteBufHandler handler = new CopyByteBufHandler(unpooledAllocator);

    ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
    ArgumentCaptor<ByteBuf> valueCapture = ArgumentCaptor.forClass(ByteBuf.class);
    doReturn(ctx).when(ctx).fireChannelRead(valueCapture.capture());

    ByteBuf pooledBuf = pooledAllocator.buffer(4);
    assertThat(pooledBuf.alloc(), is(pooledAllocator));
    try {
        assertThat(writeAscii(pooledBuf, "test"), is(4));
        handler.channelRead(ctx, pooledBuf);
        assertThat(pooledBuf.refCnt(), is(0));

        ByteBuf unpooledBuf = valueCapture.getValue();
        assertThat(unpooledBuf, is(not(sameInstance(pooledBuf))));
        assertThat(unpooledBuf.alloc(), is(unpooledAllocator));
        assertThat(unpooledBuf.toString(US_ASCII), equalTo("test"));
    } finally {
        if (pooledBuf.refCnt() > 0) {
            pooledBuf.release();
        }
    }
}
 

private static ByteBuf encodePublishMessage(
        ByteBufAllocator byteBufAllocator,
        MqttPublishMessage message) {
    MqttFixedHeader mqttFixedHeader = message.fixedHeader();
    MqttPublishVariableHeader variableHeader = message.variableHeader();
    ByteBuf payload = message.payload().duplicate();

    String topicName = variableHeader.topicName();
    byte[] topicNameBytes = encodeStringUtf8(topicName);

    int variableHeaderBufferSize = 2 + topicNameBytes.length +
            (mqttFixedHeader.qosLevel().value() > 0 ? 2 : 0);
    int payloadBufferSize = payload.readableBytes();
    int variablePartSize = variableHeaderBufferSize + payloadBufferSize;
    int fixedHeaderBufferSize = 1 + getVariableLengthInt(variablePartSize);

    ByteBuf buf = byteBufAllocator.buffer(fixedHeaderBufferSize + variablePartSize);
    buf.writeByte(getFixedHeaderByte1(mqttFixedHeader));
    writeVariableLengthInt(buf, variablePartSize);
    buf.writeShort(topicNameBytes.length);
    buf.writeBytes(topicNameBytes);
    if (mqttFixedHeader.qosLevel().value() > 0) {
        buf.writeShort(variableHeader.messageId());
    }
    buf.writeBytes(payload);

    return buf;
}
 

@Override
public ByteBuf encodeHeader(ByteBufAllocator allocator) {
    ByteBuf buffer = allocator.buffer(1024);

    boolean release = true;
    try {
        // flags:1
        buffer.writeByte(getFlags());

        // code:1
        buffer.writeByte(getResponseCode().byteValue());

        // tracing:25
        CodecUtils.encodeTrace(getTracing(), buffer);

        // headers -> nh:1 (hk~1 hv~1){nh}
        CodecUtils.encodeSmallHeaders(getHeaders(), buffer);

        // csumtype:1
        buffer.writeByte(getChecksumType().byteValue());

        // (csum:4){0,1}
        CodecUtils.encodeChecksum(getChecksum(), getChecksumType(), buffer);
        release = false;
    } finally {
        if (release) {
            buffer.release();
        }
    }
    return buffer;
}
 

private static ByteBuf encodeConnAckMessage(
        ByteBufAllocator byteBufAllocator,
        MqttConnAckMessage message) {
    ByteBuf buf = byteBufAllocator.buffer(4);
    buf.writeByte(getFixedHeaderByte1(message.fixedHeader()));
    buf.writeByte(2);
    buf.writeByte(message.variableHeader().sessionPresent() ? 0x01 : 0x00);
    buf.writeByte(message.variableHeader().returnCode().byteValue());

    return buf;
}
 

@Test
public void simpleDiscardRequestChannelTest() {
  ByteBufAllocator allocator = rule.alloc();

  rule.setAcceptingSocket(
      new RSocket() {
        @Override
        public Flux<Payload> requestChannel(Publisher<Payload> payloads) {
          return (Flux<Payload>) payloads;
        }
      },
      1);

  rule.sendRequest(1, REQUEST_STREAM);

  ByteBuf cancelFrame = CancelFrameCodec.encode(allocator, 1);

  ByteBuf metadata1 = allocator.buffer();
  metadata1.writeCharSequence("abc1", CharsetUtil.UTF_8);
  ByteBuf data1 = allocator.buffer();
  data1.writeCharSequence("def1", CharsetUtil.UTF_8);
  ByteBuf nextFrame1 =
      PayloadFrameCodec.encode(allocator, 1, false, false, true, metadata1, data1);

  ByteBuf metadata2 = allocator.buffer();
  metadata2.writeCharSequence("abc2", CharsetUtil.UTF_8);
  ByteBuf data2 = allocator.buffer();
  data2.writeCharSequence("def2", CharsetUtil.UTF_8);
  ByteBuf nextFrame2 =
      PayloadFrameCodec.encode(allocator, 1, false, false, true, metadata2, data2);

  ByteBuf metadata3 = allocator.buffer();
  metadata3.writeCharSequence("abc3", CharsetUtil.UTF_8);
  ByteBuf data3 = allocator.buffer();
  data3.writeCharSequence("de3", CharsetUtil.UTF_8);
  ByteBuf nextFrame3 =
      PayloadFrameCodec.encode(allocator, 1, false, false, true, metadata3, data3);
  rule.connection.addToReceivedBuffer(nextFrame1, nextFrame2, nextFrame3);

  rule.connection.addToReceivedBuffer(cancelFrame);

  Assertions.assertThat(rule.connection.getSent()).allMatch(ReferenceCounted::release);

  rule.assertHasNoLeaks();
}
 

static ByteBuf encode(
    ByteBufAllocator allocator,
    boolean isEmpty,
    long traceIdHigh,
    long traceId,
    boolean extendedTraceId,
    long spanId,
    long parentId,
    boolean includesParent,
    Flags flag) {
  int size =
      1
          + (isEmpty
              ? 0
              : (Long.BYTES
                  + Long.BYTES
                  + (extendedTraceId ? Long.BYTES : 0)
                  + (includesParent ? Long.BYTES : 0)));
  final ByteBuf buffer = allocator.buffer(size);

  int byteFlags = 0;
  switch (flag) {
    case NOT_SAMPLE:
      byteFlags |= FLAG_NOT_SAMPLED;
      break;
    case SAMPLE:
      byteFlags |= FLAG_SAMPLED;
      break;
    case DEBUG:
      byteFlags |= FLAG_DEBUG;
      break;
  }

  if (isEmpty) {
    return buffer.writeByte(byteFlags);
  }

  byteFlags |= FLAG_IDS_SET;

  if (extendedTraceId) {
    byteFlags |= FLAG_EXTENDED_TRACE_ID_SIZE;
  }

  if (includesParent) {
    byteFlags |= FLAG_INCLUDE_PARENT_ID;
  }

  buffer.writeByte(byteFlags);

  if (extendedTraceId) {
    buffer.writeLong(traceIdHigh);
  }

  buffer.writeLong(traceId).writeLong(spanId);

  if (includesParent) {
    buffer.writeLong(parentId);
  }

  return buffer;
}