io.aeron.logbuffer.FragmentHandler Java Examples

private void waitUntilFullConnectivity()
{
    while (!spy.isConnected() || !subscription.isConnected() || !publication.isConnected())
    {
        Tests.yield();
    }

    // send initial message to ensure connectivity
    while (publication.offer(buffer, 0, buffer.capacity()) < 0L)
    {
        Tests.yield();
    }

    final FragmentHandler mockFragmentHandler = mock(FragmentHandler.class);

    while (spy.poll(mockFragmentHandler, 1) == 0)
    {
        Tests.yield();
    }

    while (subscription.poll(mockFragmentHandler, 1) == 0)
    {
        Tests.yield();
    }
}
 

private static void consume(
    final Subscription subscription, final int startIndex, final int count, final String prefix)
{
    final MutableInteger received = new MutableInteger(startIndex);

    final FragmentHandler fragmentHandler = new FragmentAssembler(
        (buffer, offset, length, header) ->
        {
            final String expected = prefix + received.value;
            final String actual = buffer.getStringWithoutLengthAscii(offset, length);

            assertEquals(expected, actual);

            received.value++;
        });

    while (received.value < (startIndex + count))
    {
        if (0 == subscription.poll(fragmentHandler, FRAGMENT_LIMIT))
        {
            Tests.yield();
        }
    }

    assertEquals(startIndex + count, received.get());
}
 

static void consume(final Subscription subscription, final int count, final String prefix)
{
    final MutableInteger received = new MutableInteger(0);

    final FragmentHandler fragmentHandler = new FragmentAssembler(
        (buffer, offset, length, header) ->
        {
            final String expected = prefix + received.value;
            final String actual = buffer.getStringWithoutLengthAscii(offset, length);

            assertEquals(expected, actual);

            received.value++;
        });

    while (received.value < count)
    {
        if (0 == subscription.poll(fragmentHandler, FRAGMENT_LIMIT))
        {
            Tests.yield();
        }
    }

    assertEquals(count, received.get());
}
 

private void verifyData(final UnsafeBuffer srcBuffer, final FragmentHandler mockFragmentHandler)
{
    final ArgumentCaptor<DirectBuffer> bufferArg = ArgumentCaptor.forClass(DirectBuffer.class);
    final ArgumentCaptor<Integer> offsetArg = ArgumentCaptor.forClass(Integer.class);

    verify(mockFragmentHandler, times(1)).onFragment(
        bufferArg.capture(), offsetArg.capture(), eq(srcBuffer.capacity()), any(Header.class));

    final DirectBuffer capturedBuffer = bufferArg.getValue();
    final int offset = offsetArg.getValue();
    for (int i = 0; i < srcBuffer.capacity(); i++)
    {
        final int index = offset + i;
        assertEquals(srcBuffer.getByte(i), capturedBuffer.getByte(index), "same at " + index);
    }
}
 

public void onStart(final Cluster cluster, final Image snapshotImage)
{
    this.cluster = cluster;
    this.idleStrategy = cluster.idleStrategy();

    if (null != snapshotImage)
    {
        System.out.println("onStart load snapshot");
        final FragmentHandler fragmentHandler =
            (buffer, offset, length, header) -> messageCount = buffer.getInt(offset);

        while (snapshotImage.poll(fragmentHandler, 1) <= 0)
        {
            idleStrategy.idle();
        }

        System.out.println("snapshot messageCount=" + messageCount);
    }
}
 

/**
 * Return a reusable, parametrised event loop that calls and idler when no messages are received.
 *
 * @param fragmentHandler to be called back for each message.
 * @param limit           passed to {@link Subscription#poll(FragmentHandler, int)}.
 * @param running         indication for loop.
 * @param idleStrategy    to use for loop.
 * @return loop function.
 */
public static Consumer<Subscription> subscriberLoop(
    final FragmentHandler fragmentHandler,
    final int limit,
    final AtomicBoolean running,
    final IdleStrategy idleStrategy)
{
    return
        (subscription) ->
        {
            final FragmentAssembler assembler = new FragmentAssembler(fragmentHandler);
            while (running.get())
            {
                final int fragmentsRead = subscription.poll(assembler, limit);
                idleStrategy.idle(fragmentsRead);
            }
        };
}
 

/**
 * Poll the {@link Image}s under the subscription for available message fragments.
 * <p>
 * Each fragment read will be a whole message if it is under MTU length. If larger than MTU then it will come
 * as a series of fragments ordered within a session.
 * <p>
 * To assemble messages that span multiple fragments then use {@link FragmentAssembler}.
 *
 * @param fragmentHandler callback for handling each message fragment as it is read.
 * @param fragmentLimit   number of message fragments to limit when polling across multiple {@link Image}s.
 * @return the number of fragments received
 */
public int poll(final FragmentHandler fragmentHandler, final int fragmentLimit)
{
    final Image[] images = this.images;
    final int length = images.length;
    int fragmentsRead = 0;

    int startingIndex = roundRobinIndex++;
    if (startingIndex >= length)
    {
        roundRobinIndex = startingIndex = 0;
    }

    for (int i = startingIndex; i < length && fragmentsRead < fragmentLimit; i++)
    {
        fragmentsRead += images[i].poll(fragmentHandler, fragmentLimit - fragmentsRead);
    }

    for (int i = 0; i < startingIndex && fragmentsRead < fragmentLimit; i++)
    {
        fragmentsRead += images[i].poll(fragmentHandler, fragmentLimit - fragmentsRead);
    }

    return fragmentsRead;
}
 

@Test
public void shouldReadData()
{
    subscription.addImage(imageOneMock);

    when(imageOneMock.poll(any(FragmentHandler.class), anyInt())).then(
        (invocation) ->
        {
            final FragmentHandler handler = (FragmentHandler)invocation.getArguments()[0];
            handler.onFragment(atomicReadBuffer, HEADER_LENGTH, READ_BUFFER_CAPACITY - HEADER_LENGTH, header);

            return 1;
        });

    assertEquals(1, subscription.poll(fragmentHandler, FRAGMENT_COUNT_LIMIT));
    verify(fragmentHandler).onFragment(
        eq(atomicReadBuffer),
        eq(HEADER_LENGTH),
        eq(READ_BUFFER_CAPACITY - HEADER_LENGTH),
        any(Header.class));
}
 

/**
 * A reusable, parameterised {@link FragmentHandler} that prints to stdout for the second stream (STREAM + 1).
 *
 * @param streamId to show when printing.
 * @return subscription data handler function that prints the message contents.
 */
public static FragmentHandler reassembledStringMessage2(final int streamId)
{
    return (buffer, offset, length, header) ->
    {
        final byte[] data = new byte[length];
        buffer.getBytes(offset, data);

        System.out.format(
            "message to stream %d from session %x term id %x term offset %d (%d@%d)%n",
            streamId, header.sessionId(), header.termId(), header.termOffset(), length, offset);

        if (length != 9000)
        {
            System.out.format(
                "Received message was not assembled properly; received length was %d, but was expecting 9000%n",
                length);
        }
    };
}
 

/**
 * Return a reusable, parameterized {@link FragmentHandler} that prints to stdout for the first stream(STREAM).
 *
 * @param streamId to show when printing.
 * @return subscription data handler function that prints the message contents.
 */
public static FragmentHandler reassembledStringMessage1(final int streamId)
{
    return (buffer, offset, length, header) ->
    {
        final byte[] data = new byte[length];
        buffer.getBytes(offset, data);

        System.out.format(
            "message to stream %d from session %x term id %x term offset %d (%d@%d)%n",
            streamId, header.sessionId(), header.termId(), header.termOffset(), length, offset);

        if (length != 10000)
        {
            System.out.format(
                "Received message was not assembled properly;" +
                " received length was %d, but was expecting 10000%n",
                length);
        }
    };
}
 

private static void roundTripMessages(
    final FragmentHandler fragmentHandler,
    final Publication publication,
    final Subscription subscription,
    final long count)
{
    while (!subscription.isConnected())
    {
        Thread.yield();
    }

    final Image image = subscription.imageAtIndex(0);

    for (long i = 0; i < count; i++)
    {
        long offeredPosition;

        do
        {
            OFFER_BUFFER.putLong(0, System.nanoTime());
        }
        while ((offeredPosition = publication.offer(OFFER_BUFFER, 0, MESSAGE_LENGTH, null)) < 0L);

        POLLING_IDLE_STRATEGY.reset();

        do
        {
            while (image.poll(fragmentHandler, FRAGMENT_COUNT_LIMIT) <= 0)
            {
                POLLING_IDLE_STRATEGY.idle();
            }
        }
        while (image.position() < offeredPosition);
    }
}
 

private void verifyFragments(final FragmentHandler fragmentHandler, final int numMessagesToSend)
{
    verify(fragmentHandler, times(numMessagesToSend)).onFragment(
        any(DirectBuffer.class),
        anyInt(),
        eq(MESSAGE_LENGTH),
        any(Header.class));
}
 

public static int pollForFragments(
    final Subscription subscription,
    final FragmentHandler handler,
    final int minFragments,
    final long timeoutNs)
{
    final long startNs = System.nanoTime();

    long nowNs = startNs;
    int totalFragments = 0;
    do
    {
        final int numFragments = subscription.poll(handler, 10);
        if (numFragments <= 0)
        {
            Thread.yield();
            Tests.checkInterruptStatus();
            nowNs = System.nanoTime();
        }
        else
        {
            totalFragments += numFragments;
        }
    }
    while (totalFragments < minFragments && ((nowNs - startNs) < timeoutNs));

    return totalFragments;
}
 

public void onStart(final Cluster cluster, final Image snapshotImage)
{
    super.onStart(cluster, snapshotImage);

    if (null != snapshotImage)
    {
        activeSessionCount = cluster.clientSessions().size();

        final FragmentHandler handler =
            (buffer, offset, length, header) -> messageCount = buffer.getInt(offset);

        int fragmentCount = 0;
        while (true)
        {
            final int fragments = snapshotImage.poll(handler, 10);
            fragmentCount += fragments;

            if (snapshotImage.isClosed() || snapshotImage.isEndOfStream())
            {
                break;
            }

            idleStrategy.idle(fragments);
        }

        if (fragmentCount != SNAPSHOT_FRAGMENT_COUNT)
        {
            throw new AgentTerminationException(
                "unexpected snapshot length: expected=" + SNAPSHOT_FRAGMENT_COUNT + " actual=" + fragmentCount);
        }

        wasSnapshotLoaded = true;
    }
}
 

public static void main(final String[] args)
{
    System.out.println("Subscribing to " + CHANNEL + " on stream id " + STREAM_ID);

    final FragmentHandler fragmentHandler = SamplesUtil.printStringMessage(STREAM_ID);
    final AtomicBoolean running = new AtomicBoolean(true);

    SigInt.register(() -> running.set(false));

    // Create a unique response stream id so not to clash with other archive clients.
    final AeronArchive.Context archiveCtx = new AeronArchive.Context()
        .controlResponseStreamId(AeronArchive.Configuration.controlResponseStreamId() + 2);

    try (AeronArchive archive = AeronArchive.connect(archiveCtx))
    {
        final long recordingId = findLatestRecording(archive);
        final long position = 0L;
        final long length = Long.MAX_VALUE;

        final long sessionId = archive.startReplay(recordingId, position, length, CHANNEL, REPLAY_STREAM_ID);
        final String channel = ChannelUri.addSessionId(CHANNEL, (int)sessionId);

        try (Subscription subscription = archive.context().aeron().addSubscription(channel, REPLAY_STREAM_ID))
        {
            SamplesUtil.subscriberLoop(fragmentHandler, FRAGMENT_COUNT_LIMIT, running).accept(subscription);
            System.out.println("Shutting down...");
        }
    }
}
 

/**
 * Construct an adapter to reassemble message fragments and delegate on whole messages.
 *
 * @param delegate            onto which whole messages are forwarded.
 * @param initialBufferLength to be used for each session.
 * @param isDirectByteBuffer  is the underlying buffer to be a direct {@link java.nio.ByteBuffer}?
 */
public FragmentAssembler(
    final FragmentHandler delegate, final int initialBufferLength, final boolean isDirectByteBuffer)
{
    this.initialBufferLength = initialBufferLength;
    this.delegate = delegate;
    this.isDirectByteBuffer = isDirectByteBuffer;
}
 

public static void main(final String[] args)
{
    System.out.println("Subscribing to " + CHANNEL + " on stream id " + STREAM_ID);

    final MediaDriver driver = EMBEDDED_MEDIA_DRIVER ? MediaDriver.launchEmbedded() : null;
    final Aeron.Context ctx = new Aeron.Context()
        .availableImageHandler(SamplesUtil::printAvailableImage)
        .unavailableImageHandler(SamplesUtil::printUnavailableImage);

    if (EMBEDDED_MEDIA_DRIVER)
    {
        ctx.aeronDirectoryName(driver.aeronDirectoryName());
    }

    final FragmentHandler fragmentHandler = SamplesUtil.printStringMessage(STREAM_ID);
    final AtomicBoolean running = new AtomicBoolean(true);

    // Register a SIGINT handler for graceful shutdown.
    SigInt.register(() -> running.set(false));

    // Create an Aeron instance using the configured Context and create a
    // Subscription on that instance that subscribes to the configured
    // channel and stream ID.
    // The Aeron and Subscription classes implement "AutoCloseable" and will automatically
    // clean up resources when this try block is finished
    try (Aeron aeron = Aeron.connect(ctx);
        Subscription subscription = aeron.addSubscription(CHANNEL, STREAM_ID))
    {
        SamplesUtil.subscriberLoop(fragmentHandler, FRAGMENT_COUNT_LIMIT, running).accept(subscription);

        System.out.println("Shutting down...");
    }

    CloseHelper.quietClose(driver);
}
 

/**
 * Return a reusable, parametrised {@link FragmentHandler} that prints to stdout.
 *
 * @param streamId to show when printing.
 * @return subscription data handler function that prints the message contents.
 */
public static FragmentHandler printStringMessage(final int streamId)
{
    return (buffer, offset, length, header) ->
    {
        final byte[] data = new byte[length];
        buffer.getBytes(offset, data);

        System.out.println(String.format(
            "Message to stream %d from session %d (%d@%d) <<%s>>",
            streamId, header.sessionId(), length, offset, new String(data)));
    };
}
 

/**
 * Return a reusable, parameterized event
 * loop that calls and idler
 * when no messages are received
 *
 * @param fragmentHandler to be called back for each message.
 * @param limit           passed to {@link Subscription#poll(FragmentHandler, int)}
 * @param running         indication for loop
 * @param idleStrategy    to use for loop
 * @return loop function
 */
public static Consumer<Subscription> subscriberLoop(final FragmentHandler fragmentHandler, final int limit,
                final AtomicBoolean running, final IdleStrategy idleStrategy, final AtomicBoolean launched) {
    return (subscription) -> {
        try {
            while (running.get()) {
                idleStrategy.idle(subscription.poll(fragmentHandler, limit));
                launched.set(true);
            }
        } catch (final Exception ex) {
            LangUtil.rethrowUnchecked(ex);
        }
    };
}
 

/**
 * Return a reusable, parameterized {@link FragmentHandler} that prints to stdout
 *
 * @param streamId to show when printing
 * @return subscription data handler function that prints the message contents
 */
public static FragmentHandler printStringMessage(final int streamId) {
    return (buffer, offset, length, header) -> {
        final byte[] data = new byte[length];
        buffer.getBytes(offset, data);

        System.out.println(String.format("Message to stream %d from session %d (%d@%d) <<%s>>", streamId,
                        header.sessionId(), length, offset, new String(data)));
    };
}
 

/**
 * Return a reusable, parameterized event
 * loop that calls and idler
 * when no messages are received
 *
 * @param fragmentHandler to be called back for each message.
 * @param limit           passed to {@link Subscription#poll(FragmentHandler, int)}
 * @param running         indication for loop
 * @param idleStrategy    to use for loop
 * @return loop function
 */
public static Consumer<Subscription> subscriberLoop(final FragmentHandler fragmentHandler, final int limit,
                final AtomicBoolean running, final IdleStrategy idleStrategy, final AtomicBoolean launched) {
    return (subscription) -> {
        try {
            while (running.get()) {
                idleStrategy.idle(subscription.poll(fragmentHandler, limit));
                launched.set(true);
            }
        } catch (final Exception ex) {
            LangUtil.rethrowUnchecked(ex);
        }
    };
}
 

/**
 * Return a reusable, parameterized {@link FragmentHandler} that prints to stdout
 *
 * @param streamId to show when printing
 * @return subscription data handler function that prints the message contents
 */
public static FragmentHandler printStringMessage(final int streamId) {
    return (buffer, offset, length, header) -> {
        final byte[] data = new byte[length];
        buffer.getBytes(offset, data);

        System.out.println(String.format("Message to stream %d from session %d (%d@%d) <<%s>>", streamId,
                        header.sessionId(), length, offset, new String(data)));
    };
}
 

static void pipeMessages(
    final Subscription subscription, final ExclusivePublication publication, final AtomicBoolean running)
{
    final IdleStrategy idleStrategy = idleStrategy();
    final BufferClaim bufferClaim = new BufferClaim();
    final FragmentHandler dataHandler =
        (buffer, offset, length, header) ->
        {
            long result;
            while ((result = publication.tryClaim(length, bufferClaim)) <= 0)
            {
                checkPublicationResult(result);
            }

            bufferClaim
                .flags(header.flags())
                .putBytes(buffer, offset, length)
                .commit();
        };

    final Image image = subscription.imageAtIndex(0);
    final int fragmentLimit = fragmentLimit();

    while (true)
    {
        final int fragmentsRead = image.poll(dataHandler, fragmentLimit);
        if (0 == fragmentsRead)
        {
            if (image.isClosed() || !running.get())
            {
                break;
            }
        }

        idleStrategy.idle(fragmentsRead);
    }
}
 

private int pollFragments(final Subscription subscription, final FragmentHandler fragmentHandler)
{
    final int fragmentsRead = subscription.poll(fragmentHandler, FRAGMENT_COUNT_LIMIT);
    if (0 == fragmentsRead)
    {
        Tests.yield();
    }

    return fragmentsRead;
}
 

@Test
@Timeout(10)
public void shouldTransferUnfragmentedTwoPartMessage()
{
    final UnsafeBuffer expectedBuffer = new UnsafeBuffer(new byte[256]);
    final UnsafeBuffer bufferOne = new UnsafeBuffer(expectedBuffer, 0, 32);
    final UnsafeBuffer bufferTwo = new UnsafeBuffer(expectedBuffer, 32, expectedBuffer.capacity() - 32);

    bufferOne.setMemory(0, bufferOne.capacity(), (byte)'a');
    bufferTwo.setMemory(0, bufferTwo.capacity(), (byte)'b');
    final String expectedMessage = expectedBuffer.getStringWithoutLengthAscii(0, expectedBuffer.capacity());

    final MutableReference<String> receivedMessage = new MutableReference<>();
    final FragmentHandler fragmentHandler = (buffer, offset, length, header) ->
        receivedMessage.set(buffer.getStringWithoutLengthAscii(offset, length));

    try (Subscription subscription = aeron.addSubscription(CHANNEL, STREAM_ID))
    {
        try (Publication publication = aeron.addPublication(CHANNEL, STREAM_ID))
        {
            publishMessage(bufferOne, bufferTwo, publication);
            pollForMessage(subscription, receivedMessage, fragmentHandler);

            assertEquals(expectedMessage, receivedMessage.get());
        }

        try (Publication publication = aeron.addExclusivePublication(CHANNEL, STREAM_ID))
        {
            publishMessage(bufferOne, bufferTwo, publication);
            pollForMessage(subscription, receivedMessage, fragmentHandler);

            assertEquals(expectedMessage, receivedMessage.get());
        }
    }
}
 

@Test
@Timeout(10)
public void shouldTransferFragmentedTwoPartMessage()
{
    final UnsafeBuffer expectedBuffer = new UnsafeBuffer(new byte[32 + driver.context().mtuLength()]);
    final UnsafeBuffer bufferOne = new UnsafeBuffer(expectedBuffer, 0, 32);
    final UnsafeBuffer bufferTwo = new UnsafeBuffer(expectedBuffer, 32, expectedBuffer.capacity() - 32);

    bufferOne.setMemory(0, bufferOne.capacity(), (byte)'a');
    bufferTwo.setMemory(0, bufferTwo.capacity(), (byte)'b');
    final String expectedMessage = expectedBuffer.getStringWithoutLengthAscii(0, expectedBuffer.capacity());

    final MutableReference<String> receivedMessage = new MutableReference<>();
    final FragmentHandler fragmentHandler = new FragmentAssembler((buffer, offset, length, header) ->
        receivedMessage.set(buffer.getStringWithoutLengthAscii(offset, length)));

    try (Subscription subscription = aeron.addSubscription(CHANNEL, STREAM_ID))
    {
        try (Publication publication = aeron.addPublication(CHANNEL, STREAM_ID))
        {
            publishMessage(bufferOne, bufferTwo, publication);
            pollForMessage(subscription, receivedMessage, fragmentHandler);

            assertEquals(expectedMessage, receivedMessage.get());
        }

        try (Publication publication = aeron.addExclusivePublication(CHANNEL, STREAM_ID))
        {
            publishMessage(bufferOne, bufferTwo, publication);
            pollForMessage(subscription, receivedMessage, fragmentHandler);

            assertEquals(expectedMessage, receivedMessage.get());
        }
    }
}
 

private void pollForMessage(
    final Subscription subscription, final MutableReference<String> receivedMessage, final FragmentHandler handler)
{
    receivedMessage.set(null);

    while (receivedMessage.get() == null)
    {
        final int fragments = subscription.poll(handler, FRAGMENT_COUNT_LIMIT);
        if (fragments == 0)
        {
            Tests.yield();
        }
    }
}
 

private int pollWithTimeout(
    final Subscription subscription,
    final FragmentHandler fragmentHandler,
    final int fragmentLimit,
    final Supplier<String> message)
{
    final int numFragments = subscription.poll(fragmentHandler, fragmentLimit);
    if (0 == numFragments)
    {
        Tests.yieldingWait(message);
    }

    return numFragments;
}
 

private void pollForFragment(final Subscription subscription, final FragmentHandler handler)
{
    while (0 == subscription.poll(handler, 1))
    {
        Tests.yield();
    }
}
 

private void verifyFragments(final FragmentHandler fragmentHandler, final int numMessagesToSend)
{
    verify(fragmentHandler, times(numMessagesToSend)).onFragment(
        any(DirectBuffer.class),
        anyInt(),
        eq(MESSAGE_LENGTH),
        any(Header.class));
}