Examples with ByteBufferOutputStream com.github.ambry.utils.ByteBufferOutputStream used on opensource projects

Example 16 with ByteBufferOutputStream

use of com.github.ambry.utils.ByteBufferOutputStream in project ambry by linkedin.

the class CompositeSendTest method testCompositeSendWithoutByteBuf.

@Test
public void testCompositeSendWithoutByteBuf() throws IOException {
    byte[] buf1 = new byte[1024];
    byte[] buf2 = new byte[2048];
    byte[] buf3 = new byte[4096];
    new Random().nextBytes(buf1);
    new Random().nextBytes(buf2);
    new Random().nextBytes(buf3);
    ByteArraySend byteArraySend1 = new ByteArraySend(buf1, false);
    ByteArraySend byteArraySend2 = new ByteArraySend(buf2, false);
    ByteArraySend byteArraySend3 = new ByteArraySend(buf3, false);
    List<Send> listToSend = new ArrayList<Send>(3);
    listToSend.add(byteArraySend1);
    listToSend.add(byteArraySend2);
    listToSend.add(byteArraySend3);
    CompositeSend compositeSend = new CompositeSend(listToSend);
    ByteBuffer bufferToWrite = ByteBuffer.allocate(1024 + 2048 + 4096);
    ByteBufferOutputStream bufferToWriteStream = new ByteBufferOutputStream(bufferToWrite);
    WritableByteChannel writableByteChannel = Channels.newChannel(bufferToWriteStream);
    while (!compositeSend.isSendComplete()) {
        compositeSend.writeTo(writableByteChannel);
    }
    bufferToWrite.flip();
    for (int i = 0; i < 1024; i++) {
        Assert.assertEquals(buf1[i], bufferToWrite.get(i));
    }
    for (int i = 0; i < 2048; i++) {
        Assert.assertEquals(buf2[i], bufferToWrite.get(1024 + i));
    }
    for (int i = 0; i < 4096; i++) {
        Assert.assertEquals(buf3[i], bufferToWrite.get(1024 + 2048 + i));
    }
}

Example 17 with ByteBufferOutputStream

use of com.github.ambry.utils.ByteBufferOutputStream in project ambry by linkedin.

the class BlobStoreTest method checkStoreInfo.

/**
 * Verifies the provided {@code storeInfo} for correctness of the {@link MessageInfo}. Also reads the blob as
 * described by the {@link MessageReadSet} inside {@code storeInfo} to verify that it matches the reference.
 * @param storeInfo the {@link StoreInfo} to verify.
 * @param expectedKeys all the {@link MockId}s that are expected to be found in {@code storeInfo}.
 * @param lifeVersion the expected life version for all the keys.
 * @throws IOException
 */
private void checkStoreInfo(StoreInfo storeInfo, Set<MockId> expectedKeys, short lifeVersion) throws IOException {
    List<MessageInfo> messageInfos = storeInfo.getMessageReadSetInfo();
    MessageReadSet readSet = storeInfo.getMessageReadSet();
    assertEquals("ReadSet contains an unexpected number of messages", expectedKeys.size(), readSet.count());
    Set<MockId> examinedKeys = new HashSet<>();
    for (int i = 0; i < messageInfos.size(); i++) {
        MessageInfo messageInfo = messageInfos.get(i);
        MockId id = (MockId) messageInfo.getStoreKey();
        MessageInfo expectedInfo = allKeys.get(id).getFirst();
        assertEquals("Unexpected size in MessageInfo", expectedInfo.getSize(), messageInfo.getSize());
        assertEquals("AccountId mismatch", expectedInfo.getAccountId(), messageInfo.getAccountId());
        assertEquals("ContainerId mismatch", expectedInfo.getContainerId(), messageInfo.getContainerId());
        assertEquals("OperationTime mismatch", expectedInfo.getOperationTimeMs(), messageInfo.getOperationTimeMs());
        assertEquals("isTTLUpdated not as expected", ttlUpdatedKeys.contains(id), messageInfo.isTtlUpdated());
        assertEquals("isDeleted not as expected", deletedKeys.contains(id), messageInfo.isDeleted());
        assertEquals("isUndeleted not as expected", undeletedKeys.contains(id), messageInfo.isUndeleted());
        if (IndexValue.hasLifeVersion(lifeVersion)) {
            assertEquals("lifeVersion not as expected", lifeVersion, messageInfo.getLifeVersion());
        }
        long expiresAtMs = ttlUpdatedKeys.contains(id) ? Utils.Infinite_Time : expectedInfo.getExpirationTimeInMs();
        expiresAtMs = Utils.getTimeInMsToTheNearestSec(expiresAtMs);
        assertEquals("Unexpected expiresAtMs in MessageInfo", expiresAtMs, messageInfo.getExpirationTimeInMs());
        assertEquals("Unexpected key in readSet", id, readSet.getKeyAt(i));
        assertEquals("Unexpected size in ReadSet", expectedInfo.getSize(), readSet.sizeInBytes(i));
        ByteBuffer readBuf = ByteBuffer.allocate((int) expectedInfo.getSize());
        ByteBufferOutputStream stream = new ByteBufferOutputStream(readBuf);
        WritableByteChannel channel = Channels.newChannel(stream);
        readSet.writeTo(i, channel, 0, expectedInfo.getSize());
        ByteBuffer expectedData = allKeys.get(id).getSecond();
        assertArrayEquals("Data obtained from reset does not match original", expectedData.array(), readBuf.array());
        examinedKeys.add(id);
    }
    assertEquals("Expected and examined keys do not match", expectedKeys, examinedKeys);
}

Example 18 with ByteBufferOutputStream

use of com.github.ambry.utils.ByteBufferOutputStream in project ambry by linkedin.

the class StoreMessageReadSetTest method storeMessageReadSetTest.

/**
 * Primarily tests {@link StoreMessageReadSet} and its APIs but also checks the the {@link Comparable} APIs of
 * {@link BlobReadOptions}.
 * @throws IOException
 */
@Test
public void storeMessageReadSetTest() throws IOException, StoreException {
    int logCapacity = 2000;
    int segCapacity = 1000;
    Log log = new Log(tempDir.getAbsolutePath(), logCapacity, StoreTestUtils.DEFAULT_DISK_SPACE_ALLOCATOR, createStoreConfig(segCapacity, setFilePermissionEnabled), metrics, null);
    try {
        LogSegment firstSegment = log.getFirstSegment();
        int availableSegCapacity = (int) (segCapacity - firstSegment.getStartOffset());
        byte[] srcOfTruth = TestUtils.getRandomBytes(2 * availableSegCapacity);
        ReadableByteChannel dataChannel = Channels.newChannel(new ByteBufferInputStream(ByteBuffer.wrap(srcOfTruth)));
        log.appendFrom(dataChannel, availableSegCapacity);
        log.appendFrom(dataChannel, availableSegCapacity);
        LogSegment secondSegment = log.getNextSegment(firstSegment);
        Offset firstSegOffset1 = new Offset(firstSegment.getName(), firstSegment.getStartOffset());
        Offset firstSegOffset2 = new Offset(firstSegment.getName(), firstSegment.getStartOffset() + availableSegCapacity / 2);
        Offset secondSegOffset1 = new Offset(secondSegment.getName(), secondSegment.getStartOffset());
        Offset secondSegOffset2 = new Offset(secondSegment.getName(), secondSegment.getStartOffset() + availableSegCapacity / 2);
        List<MockId> mockIdList = new ArrayList<>();
        MockId mockId = new MockId("id1");
        mockIdList.add(mockId);
        BlobReadOptions ro1 = new BlobReadOptions(log, firstSegOffset2, new MessageInfo(mockId, availableSegCapacity / 3, 1, Utils.getRandomShort(TestUtils.RANDOM), Utils.getRandomShort(TestUtils.RANDOM), System.currentTimeMillis() + TestUtils.RANDOM.nextInt(10000)));
        mockId = new MockId("id2");
        mockIdList.add(mockId);
        BlobReadOptions ro2 = new BlobReadOptions(log, secondSegOffset1, new MessageInfo(mockId, availableSegCapacity / 4, 1, Utils.getRandomShort(TestUtils.RANDOM), Utils.getRandomShort(TestUtils.RANDOM), System.currentTimeMillis() + TestUtils.RANDOM.nextInt(10000)));
        mockId = new MockId("id3");
        mockIdList.add(mockId);
        BlobReadOptions ro3 = new BlobReadOptions(log, secondSegOffset2, new MessageInfo(mockId, availableSegCapacity / 2, 1, Utils.getRandomShort(TestUtils.RANDOM), Utils.getRandomShort(TestUtils.RANDOM), System.currentTimeMillis() + TestUtils.RANDOM.nextInt(10000)));
        mockId = new MockId("id4");
        mockIdList.add(mockId);
        BlobReadOptions ro4 = new BlobReadOptions(log, firstSegOffset1, new MessageInfo(mockId, availableSegCapacity / 5, 1, Utils.getRandomShort(TestUtils.RANDOM), Utils.getRandomShort(TestUtils.RANDOM), System.currentTimeMillis() + TestUtils.RANDOM.nextInt(10000)));
        // to test equality in the compareTo() of BlobReadOptions
        mockId = new MockId("id5");
        mockIdList.add(mockId);
        BlobReadOptions ro5 = new BlobReadOptions(log, firstSegOffset2, new MessageInfo(mockId, availableSegCapacity / 6, 1, Utils.getRandomShort(TestUtils.RANDOM), Utils.getRandomShort(TestUtils.RANDOM), System.currentTimeMillis() + TestUtils.RANDOM.nextInt(10000)));
        List<BlobReadOptions> options = new ArrayList<>(Arrays.asList(ro1, ro2, ro3, ro4, ro5));
        MessageReadSet readSet = new StoreMessageReadSet(options);
        assertEquals(readSet.count(), options.size());
        // options should get sorted by offsets in the constructor
        assertEquals(readSet.getKeyAt(0), mockIdList.get(3));
        assertEquals(readSet.sizeInBytes(0), availableSegCapacity / 5);
        assertEquals(readSet.getKeyAt(1), mockIdList.get(0));
        assertEquals(readSet.sizeInBytes(1), availableSegCapacity / 3);
        assertEquals(readSet.getKeyAt(2), mockIdList.get(4));
        assertEquals(readSet.sizeInBytes(2), availableSegCapacity / 6);
        assertEquals(readSet.getKeyAt(3), mockIdList.get(1));
        assertEquals(readSet.sizeInBytes(3), availableSegCapacity / 4);
        assertEquals(readSet.getKeyAt(4), mockIdList.get(2));
        assertEquals(readSet.sizeInBytes(4), availableSegCapacity / 2);
        ByteBuffer readBuf = ByteBuffer.allocate(availableSegCapacity / 5);
        ByteBufferOutputStream stream = new ByteBufferOutputStream(readBuf);
        // read the first one all at once
        if (doDataPrefetch) {
            readSet.doPrefetch(0, 0, Long.MAX_VALUE);
        }
        long written = readSet.writeTo(0, Channels.newChannel(stream), 0, Long.MAX_VALUE);
        assertEquals("Return value from writeTo() is incorrect", availableSegCapacity / 5, written);
        assertArrayEquals(readBuf.array(), Arrays.copyOfRange(srcOfTruth, 0, availableSegCapacity / 5));
        // read the second one byte by byte
        readBuf = ByteBuffer.allocate(availableSegCapacity / 3);
        stream = new ByteBufferOutputStream(readBuf);
        WritableByteChannel channel = Channels.newChannel(stream);
        long currentReadOffset = 0;
        if (doDataPrefetch) {
            readSet.doPrefetch(1, currentReadOffset, availableSegCapacity / 3);
        }
        while (currentReadOffset < availableSegCapacity / 3) {
            written = readSet.writeTo(1, channel, currentReadOffset, 1);
            assertEquals("Return value from writeTo() is incorrect", 1, written);
            currentReadOffset++;
        }
        long startOffset = availableSegCapacity / 2;
        long endOffset = availableSegCapacity / 2 + availableSegCapacity / 3;
        assertArrayEquals(readBuf.array(), Arrays.copyOfRange(srcOfTruth, (int) startOffset, (int) endOffset));
        // read the last one in multiple stages
        readBuf = ByteBuffer.allocate(availableSegCapacity / 2);
        stream = new ByteBufferOutputStream(readBuf);
        channel = Channels.newChannel(stream);
        currentReadOffset = 0;
        if (doDataPrefetch) {
            readSet.doPrefetch(4, currentReadOffset, availableSegCapacity / 2);
        }
        while (currentReadOffset < availableSegCapacity / 2) {
            written = readSet.writeTo(4, channel, currentReadOffset, availableSegCapacity / 6);
            long expectedWritten = Math.min(availableSegCapacity / 2 - currentReadOffset, availableSegCapacity / 6);
            assertEquals("Return value from writeTo() is incorrect", expectedWritten, written);
            currentReadOffset += availableSegCapacity / 6;
        }
        startOffset = availableSegCapacity + availableSegCapacity / 2;
        endOffset = startOffset + availableSegCapacity / 2;
        assertArrayEquals(readBuf.array(), Arrays.copyOfRange(srcOfTruth, (int) startOffset, (int) endOffset));
        // should not write anything if relative offset is at the size
        readBuf = ByteBuffer.allocate(1);
        stream = new ByteBufferOutputStream(readBuf);
        channel = Channels.newChannel(stream);
        written = readSet.writeTo(0, channel, readSet.sizeInBytes(0), 1);
        assertEquals("No data should have been written", 0, written);
        try {
            readSet.sizeInBytes(options.size());
            fail("Reading should have failed because index is out of bounds");
        } catch (IndexOutOfBoundsException e) {
        // expected. Nothing to do.
        }
        try {
            readSet.writeTo(options.size(), channel, 100, 100);
            fail("Reading should have failed because index is out of bounds");
        } catch (IndexOutOfBoundsException e) {
        // expected. Nothing to do.
        }
        try {
            readSet.getKeyAt(options.size());
            fail("Getting key should have failed because index is out of bounds");
        } catch (IndexOutOfBoundsException e) {
        // expected. Nothing to do.
        }
    } finally {
        log.close(false);
    }
}

Example 19 with ByteBufferOutputStream

use of com.github.ambry.utils.ByteBufferOutputStream in project ambry by linkedin.

the class CompositeSendTest method testCompositeSend.

@Test
public void testCompositeSend() throws IOException {
    byte[] buf1 = new byte[1024];
    byte[] buf2 = new byte[2048];
    byte[] buf3 = new byte[4096];
    new Random().nextBytes(buf1);
    new Random().nextBytes(buf2);
    new Random().nextBytes(buf3);
    ByteArraySend byteArraySend1 = new ByteArraySend(buf1);
    ByteArraySend byteArraySend2 = new ByteArraySend(buf2);
    ByteArraySend byteArraySend3 = new ByteArraySend(buf3);
    List<Send> listToSend = new ArrayList<Send>(3);
    listToSend.add(byteArraySend1);
    listToSend.add(byteArraySend2);
    listToSend.add(byteArraySend3);
    CompositeSend compositeSend = new CompositeSend(listToSend);
    ByteBuffer bufferToWrite = ByteBuffer.allocate(1024 + 2048 + 4096);
    ByteBufferOutputStream bufferToWriteStream = new ByteBufferOutputStream(bufferToWrite);
    WritableByteChannel writableByteChannel = Channels.newChannel(bufferToWriteStream);
    while (!compositeSend.isSendComplete()) {
        compositeSend.writeTo(writableByteChannel);
    }
    bufferToWrite.flip();
    for (int i = 0; i < 1024; i++) {
        Assert.assertEquals(buf1[i], bufferToWrite.get(i));
    }
    for (int i = 0; i < 2048; i++) {
        Assert.assertEquals(buf2[i], bufferToWrite.get(1024 + i));
    }
    for (int i = 0; i < 4096; i++) {
        Assert.assertEquals(buf3[i], bufferToWrite.get(1024 + 2048 + i));
    }
}

Example 20 with ByteBufferOutputStream

use of com.github.ambry.utils.ByteBufferOutputStream in project ambry by linkedin.

the class HardDeleteRecoveryMetadata method getBlobEncryptionKeyRecord.

/**
 * Get the Blob Encryption Key Record from the given readSet
 * @param readSet the {@link MessageReadSet} from which to read.
 * @param readSetIndex the index of the message in the readSet.
 * @param relativeOffset the relative offset in the message from which to read.
 * @param blobEncryptionKeySize the size of the record to read (in this case the encryption key record).
 * @return returns the read encryption key.
 * @throws MessageFormatException
 * @throws IOException
 */
private ByteBuffer getBlobEncryptionKeyRecord(MessageReadSet readSet, int readSetIndex, long relativeOffset, long blobEncryptionKeySize) throws MessageFormatException, IOException {
    /* Read the field from the channel */
    ByteBuffer blobEncryptionKey = ByteBuffer.allocate((int) blobEncryptionKeySize);
    readSet.writeTo(readSetIndex, Channels.newChannel(new ByteBufferOutputStream(blobEncryptionKey)), relativeOffset, blobEncryptionKeySize);
    blobEncryptionKey.flip();
    return deserializeBlobEncryptionKey(new ByteBufferInputStream(blobEncryptionKey));
}