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

@Override
protected Object filterOutboundMessage(Object msg) {
    if (msg instanceof UkcpPacket) {
        UkcpPacket p = (UkcpPacket) msg;
        ByteBuf content = p.content();
        if (isSingleDirectBuffer(content)) {
            return p;
        }
        content.retain(); // newDirectBuffer method call release method of content
        UkcpPacket np = UkcpPacket.newInstance(newDirectBuffer(content), p.remoteAddress());
        p.release();
        return np;
    }

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

/**
 * Publish a message to the given payload, using the given qos and optional retain
 *
 * @param topic   The topic to publish to
 * @param payload The payload to send
 * @param qos     The qos to use while publishing
 * @param retain  true if you want to retain the message on the server, false otherwise
 * @return A future which will be completed when the message is delivered to the server
 */
@Override
public Future<Void> publish(String topic, ByteBuf payload, MqttQoS qos, boolean retain) {
    Promise<Void> future = new DefaultPromise<>(this.eventLoop.next());
    MqttFixedHeader fixedHeader = new MqttFixedHeader(MqttMessageType.PUBLISH, false, qos, retain, 0);
    MqttPublishVariableHeader variableHeader = new MqttPublishVariableHeader(topic, getNewMessageId().messageId());
    MqttPublishMessage message = new MqttPublishMessage(fixedHeader, variableHeader, payload);
    MqttPendingPublish pendingPublish = new MqttPendingPublish(variableHeader.packetId(), future, payload.retain(), message, qos);
    ChannelFuture channelFuture = this.sendAndFlushPacket(message);

    if (channelFuture != null) {
        pendingPublish.setSent(true);
        if (channelFuture.cause() != null) {
            future.setFailure(channelFuture.cause());
            return future;
        }
    }
    if (pendingPublish.isSent() && pendingPublish.getQos() == MqttQoS.AT_MOST_ONCE) {
        pendingPublish.getFuture().setSuccess(null); //We don't get an ACK for QOS 0
    } else if (pendingPublish.isSent()) {
        this.pendingPublishes.put(pendingPublish.getMessageId(), pendingPublish);
        pendingPublish.startPublishRetransmissionTimer(this.eventLoop.next(), this::sendAndFlushPacket);
    }
    return future;
}
 

public CompletableFuture<TxnID> newTransactionAsync(long timeout, TimeUnit unit) {
    if (LOG.isDebugEnabled()) {
        LOG.debug("New transaction with timeout in ms {}", unit.toMillis(timeout));
    }
    CompletableFuture<TxnID> callback = new CompletableFuture<>();

    if (!canSendRequest(callback)) {
        return callback;
    }
    long requestId = client.newRequestId();
    ByteBuf cmd = Commands.newTxn(transactionCoordinatorId, requestId, unit.toMillis(timeout));
    OpForTxnIdCallBack op = OpForTxnIdCallBack.create(cmd, callback);
    pendingRequests.put(requestId, op);
    timeoutQueue.add(new RequestTime(System.currentTimeMillis(), requestId));
    cmd.retain();
    cnx().ctx().writeAndFlush(cmd, cnx().ctx().voidPromise());
    return callback;
}
 

public void dataFromScreenArrived(QueryData header, ByteBuf data, ResponseSender sender) {
  if(data != null){
    // we're going to send this some place, we need increment to ensure this is around long enough to send.
    data.retain();
    observer.execDataArrived(new ScreenShuttle(sender), new QueryWritableBatch(header, data));
  } else {
    observer.execDataArrived(new ScreenShuttle(sender), new QueryWritableBatch(header));
  }
}
 

@Override
protected void decode(ChannelHandlerContext ctx, SctpMessage msg, List<Object> out) throws Exception {
    final ByteBuf byteBuf = msg.content();
    final int protocolIdentifier = msg.protocolIdentifier();
    final int streamIdentifier = msg.streamIdentifier();
    final boolean isComplete = msg.isComplete();

    ByteBuf frag;
    if (fragments.containsKey(streamIdentifier)) {
        frag = fragments.remove(streamIdentifier);
    } else {
        frag = Unpooled.EMPTY_BUFFER;
    }

    if (isComplete && !frag.isReadable()) {
        //data chunk is not fragmented
        out.add(msg);
    } else if (!isComplete && frag.isReadable()) {
        //more message to complete
        fragments.put(streamIdentifier, Unpooled.wrappedBuffer(frag, byteBuf));
    } else if (isComplete && frag.isReadable()) {
        //last message to complete
        fragments.remove(streamIdentifier);
        SctpMessage assembledMsg = new SctpMessage(
                protocolIdentifier,
                streamIdentifier,
                Unpooled.wrappedBuffer(frag, byteBuf));
        out.add(assembledMsg);
    } else {
        //first incomplete message
        fragments.put(streamIdentifier, byteBuf);
    }
    byteBuf.retain();
}
 

/**
 * Ensures the {@code buf} will work with the specified {@code nativeStuff}. After this function
 * is called, you should decrement the reference count on the {@code buf} with
 * {@link ByteBuf#release()}.
 *
 * @param alloc the {@link ByteBufAllocator} to use
 * @param nativeStuff the native we are working with
 * @param buf the buffer we are working with
 * @return a buffer compatible with the native
 */
public static ByteBuf ensureCompatible(ByteBufAllocator alloc, Native nativeStuff, ByteBuf buf) {
  if (!nativeStuff.isNative() || buf.hasMemoryAddress()) {
    // Will always work in either case. JNI code demands a memory address, and if we have a Java
    // fallback, it uses byte arrays in all cases.
    return buf.retain();
  }

  // It's not, so we must make a direct copy.
  ByteBuf newBuf = alloc.directBuffer(buf.readableBytes());
  newBuf.writeBytes(buf);
  return newBuf;
}
 

@Override
protected void encode(ChannelHandlerContext ctx, ByteBuf msg, List<Object> out) throws Exception {

  if (RpcConstants.EXTRA_DEBUGGING) {
    logger.debug("ChunkCreationHandler called with msg {} of size {} with chunkSize {}",
        msg, msg.readableBytes(), chunkSize);
  }

  if (!ctx.channel().isOpen()) {
    logger.debug("Channel closed, skipping encode inside {}.", RpcConstants.CHUNK_CREATION_HANDLER);
    msg.release();
    return;
  }

  // Calculate the number of chunks based on configured chunk size and input msg size
  int numChunks = (int) Math.ceil((double) msg.readableBytes() / chunkSize);

  // Initialize a composite buffer to hold numChunks chunk.
  final CompositeByteBuf cbb = ctx.alloc().compositeBuffer(numChunks);

  int cbbWriteIndex = 0;
  int currentChunkLen = min(msg.readableBytes(), chunkSize);

  // Create slices of chunkSize from input msg and add it to the composite buffer.
  while (numChunks > 0) {
    final ByteBuf chunkBuf = msg.slice(msg.readerIndex(), currentChunkLen);
    chunkBuf.retain();
    cbb.addComponent(chunkBuf);
    cbbWriteIndex += currentChunkLen;
    msg.skipBytes(currentChunkLen);
    --numChunks;
    currentChunkLen = min(msg.readableBytes(), chunkSize);
  }

  // Update the writerIndex of composite byte buffer. Netty doesn't do it automatically.
  cbb.writerIndex(cbbWriteIndex);

  // Add the final composite bytebuf into output buffer.
  out.add(cbb);
}
 

private static FullBinaryMemcacheRequest toFullRequest(BinaryMemcacheRequest request, ByteBuf content) {
    ByteBuf extras = request.getExtras();
    if (extras != null) {
        extras = extras.retain();
    }
    ByteBuf framingExtras = request.getFramingExtras();
    if (framingExtras != null) {
        framingExtras = framingExtras.retain();
    }

    FullBinaryMemcacheRequest fullRequest = new DefaultFullBinaryMemcacheRequest(request.getKey(),
        extras, content);

    fullRequest.setMagic(request.getMagic());
    fullRequest.setOpcode(request.getOpcode());
    fullRequest.setKeyLength(request.getKeyLength());
    fullRequest.setExtrasLength(request.getExtrasLength());
    fullRequest.setDataType(request.getDataType());
    fullRequest.setTotalBodyLength(request.getTotalBodyLength());
    fullRequest.setOpaque(request.getOpaque());
    fullRequest.setCAS(request.getCAS());
    fullRequest.setReserved(request.getReserved());
    fullRequest.setFramingExtras(framingExtras);
    fullRequest.setFramingExtrasLength(request.getFramingExtrasLength());

    return fullRequest;
}
 

@Override
protected final void decode(ChannelHandlerContext ctx, ByteBuf buffer, List<Object> out)
{
    int frameOffset = buffer.readerIndex();
    TChannelBufferInputTransport transport = new TChannelBufferInputTransport(buffer.retain());
    try {
        TProtocolReader protocolReader = protocol.createProtocol(transport);

        TMessage message = protocolReader.readMessageBegin();
        TProtocolUtil.skip(protocolReader, TType.STRUCT);
        protocolReader.readMessageEnd();

        int frameLength = buffer.readerIndex() - frameOffset;
        if (frameLength > maxFrameSize) {
            FrameInfo frameInfo = new FrameInfo(message.getName(), message.getType(), message.getSequenceId(), UNFRAMED, protocol, assumeClientsSupportOutOfOrderResponses);
            ctx.fireExceptionCaught(new FrameTooLargeException(
                    Optional.of(frameInfo),
                    frameLength,
                    maxFrameSize));
        }

        out.add(buffer.slice(frameOffset, frameLength).retain());
    }
    catch (Throwable th) {
        buffer.readerIndex(frameOffset);
    }
    finally {
        transport.release();
    }
}
 

@DisplayName("fragments and reassembles data")
@Test
void fragmentAndReassembleData() {
  ByteBuf frame =
      PayloadFrameCodec.encodeNextCompleteReleasingPayload(
          allocator, 2, DefaultPayload.create(data));
  System.out.println(FrameUtil.toString(frame));

  frame.retain();

  Publisher<ByteBuf> fragments =
      FrameFragmenter.fragmentFrame(allocator, 64, frame, FrameHeaderCodec.frameType(frame));

  FrameReassembler reassembler = new FrameReassembler(allocator);

  ByteBuf assembled =
      Flux.from(fragments)
          .doOnNext(byteBuf -> System.out.println(FrameUtil.toString(byteBuf)))
          .handle(reassembler::reassembleFrame)
          .blockLast();

  System.out.println("assembled");
  String s = FrameUtil.toString(assembled);
  System.out.println(s);

  Assert.assertEquals(FrameHeaderCodec.frameType(frame), FrameHeaderCodec.frameType(assembled));
  Assert.assertEquals(frame.readableBytes(), assembled.readableBytes());
  Assert.assertEquals(PayloadFrameCodec.data(frame), PayloadFrameCodec.data(assembled));
}
 

private ByteBuf handleUnprotected(List<ByteBuf> unprotectedBufs, ByteBufAllocator alloc)
    throws GeneralSecurityException {
  long unprotectedBytes = 0;
  for (ByteBuf buf : unprotectedBufs) {
    unprotectedBytes += buf.readableBytes();
  }
  // Empty plaintext not allowed since this should be handled as no-op in layer above.
  checkArgument(unprotectedBytes > 0);

  // Compute number of frames and allocate a single buffer for all frames.
  long frameNum = unprotectedBytes / maxUnprotectedBytesPerFrame + 1;
  int lastFrameUnprotectedBytes = (int) (unprotectedBytes % maxUnprotectedBytesPerFrame);
  if (lastFrameUnprotectedBytes == 0) {
    frameNum--;
    lastFrameUnprotectedBytes = maxUnprotectedBytesPerFrame;
  }
  long protectedBytes = frameNum * (HEADER_BYTES + suffixBytes) + unprotectedBytes;

  ByteBuf protectedBuf = alloc.directBuffer(Ints.checkedCast(protectedBytes));
  try {
    int bufferIdx = 0;
    for (int frameIdx = 0; frameIdx < frameNum; ++frameIdx) {
      int unprotectedBytesLeft =
          (frameIdx == frameNum - 1) ? lastFrameUnprotectedBytes : maxUnprotectedBytesPerFrame;
      // Write header (at most LIMIT_MAX_ALLOWED_FRAME_BYTES).
      protectedBuf.writeIntLE(unprotectedBytesLeft + HEADER_TYPE_FIELD_BYTES + suffixBytes);
      protectedBuf.writeIntLE(HEADER_TYPE_DEFAULT);

      // Ownership of the backing buffer remains with protectedBuf.
      ByteBuf frameOut = writeSlice(protectedBuf, unprotectedBytesLeft + suffixBytes);
      List<ByteBuf> framePlain = new ArrayList<>();
      while (unprotectedBytesLeft > 0) {
        // Ownership of the buffer backing in remains with unprotectedBufs.
        ByteBuf in = unprotectedBufs.get(bufferIdx);
        if (in.readableBytes() <= unprotectedBytesLeft) {
          // The complete buffer belongs to this frame.
          framePlain.add(in);
          unprotectedBytesLeft -= in.readableBytes();
          bufferIdx++;
        } else {
          // The remainder of in will be part of the next frame.
          framePlain.add(in.readSlice(unprotectedBytesLeft));
          unprotectedBytesLeft = 0;
        }
      }
      crypter.encrypt(frameOut, framePlain);
      verify(!frameOut.isWritable());
    }
    protectedBuf.readerIndex(0);
    protectedBuf.writerIndex(protectedBuf.capacity());
    return protectedBuf.retain();
  } finally {
    protectedBuf.release();
  }
}
 

public void sendMessage(Entry entry) {
    if ((STATE_UPDATER.get(this) == State.Started) && isWritable()) {

        int length = entry.getLength();
        ByteBuf headersAndPayload = entry.getDataBuffer();
        MessageImpl msg;
        try {
            msg = MessageImpl.deserialize(headersAndPayload);
        } catch (Throwable t) {
            log.error("[{}][{} -> {}] Failed to deserialize message at {} (buffer size: {}): {}", topicName,
                    localCluster, remoteCluster, entry.getPosition(), length, t.getMessage(), t);
            entry.release();
            return;
        }

        if (msg.isReplicated()) {
            // Discard messages that were already replicated into this region
            entry.release();
            msg.recycle();
            return;
        }

        if (msg.hasReplicateTo() && !msg.getReplicateTo().contains(remoteCluster)) {
            if (log.isDebugEnabled()) {
                log.debug("[{}][{} -> {}] Skipping message at {} / msg-id: {}: replicateTo {}", topicName,
                        localCluster, remoteCluster, entry.getPosition(), msg.getMessageId(), msg.getReplicateTo());
            }
            entry.release();
            msg.recycle();
            return;
        }

        msgOut.recordEvent(headersAndPayload.readableBytes());

        msg.setReplicatedFrom(localCluster);

        headersAndPayload.retain();

        producer.sendAsync(msg, ProducerSendCallback.create(this, entry, msg));

    } else {
        if (log.isDebugEnabled()) {
            log.debug("[{}][{} -> {}] dropping message because replicator producer is not started/writable",
                    topicName, localCluster, remoteCluster);
        }
        msgDrop.recordEvent();
        entry.release();
    }
}
 

public void setBuffer(ByteBuf buffer) {
  if (buffer != null) {
    buffer.retain();
    this.buffer = buffer;
  }
}
 

@Override
public ByteBuf encode(ByteBuf raw) {
    // Provides an encoder that simply returns the same uncompressed buffer
    return raw.retain();
}
 

public MyBsonContext(@Loose @Retains @ConservesIndexes ByteBuf byteBuf) {
  this.byteBuf = byteBuf.retain();
}
 

@Override
protected final void decode(ChannelHandlerContext ctx, ByteBuf buf, List<Object> out) throws Exception {
    buf.retain();
    out.add(new MessageImpl(buf, null));
}
 

@Override
protected void channelRead0(ChannelHandlerContext ctx, ByteBuf msg) throws Exception {
    msg.retain();
    dataQueue.add(msg);
    inputStream.notifyData();
}
 

@Override
protected void channelRead0(ChannelHandlerContext ctx, ByteBuf byteBuf) throws Exception {
    byteBuf.retain();
    cryptor.decrypt(byteBuf);
    ctx.fireChannelRead(byteBuf);
}
 

void inboundDataReceived(ByteBuf frame, boolean endOfStream) {
  super.inboundDataReceived(new NettyReadableBuffer(frame.retain()), endOfStream);
}
 

@Test
public void testGrpcServer() throws Exception {
  final Http2Headers cannedHeaders = new DefaultHttp2Headers();
  cannedHeaders
      .status("200")
      .add("content-type", "application/grpc")
      .add("grpc-encoding", "identity")
      .add("grpc-accept-encoding", "gzip");

  final Http2Headers cannedTrailers = new DefaultHttp2Headers().add("grpc-status", "0");

  ByteBuf buf =
      Unpooled.copiedBuffer(ByteBufUtil.decodeHexDump("000000000d0a0b48656c6c6f20776f726c64"));
  final Http2DataFrame cannedData = new DefaultHttp2DataFrame(buf.retain(), false);

  XioServerBootstrap bootstrap =
      XioServerBootstrap.fromConfig("xio.testGrpcServer")
          .addToPipeline(
              new SmartHttpPipeline() {
                @Override
                public ChannelHandler getApplicationRouter() {
                  return new PipelineRouter(
                      ImmutableMap.of(),
                      new PipelineRequestHandler() {
                        @Override
                        public void handle(
                            ChannelHandlerContext ctx, Request request, RouteState route) {
                          if (request instanceof SegmentedRequestData) {
                            SegmentedRequestData streaming = (SegmentedRequestData) request;

                            if (streaming.endOfMessage()) {
                              ctx.write(Http2Response.build(request.streamId(), cannedHeaders));
                              ctx.write(
                                  Http2Response.build(request.streamId(), cannedData, false));
                              ctx.write(
                                  Http2Response.build(request.streamId(), cannedTrailers, true));
                            }
                          }
                        }
                      });
                }
              });

  XioServer xioServer = bootstrap.build();
  GrpcClient client = GrpcClient.run(xioServer.getPort());

  HelloReply response = client.greet("world");

  assertEquals("Hello world", response.getMessage());

  client.shutdown();
  xioServer.close();
}