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

Example 36 with ByteBufferInputStream

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

the class AmbryServerRequestsTest method verifyTtlUpdate.

/**
 * Verifies that the TTL update request was delivered to the {@link Store} correctly.
 * @param blobId the {@link BlobId} that was updated
 * @param expiresAtMs the new expire time (in ms)
 * @param opTimeMs the op time (in ms)
 * @param messageWriteSet the {@link MessageWriteSet} received at the {@link Store}
 * @throws IOException
 * @throws MessageFormatException
 * @throws StoreException
 */
private void verifyTtlUpdate(BlobId blobId, long expiresAtMs, long opTimeMs, MessageWriteSet messageWriteSet) throws IOException, MessageFormatException, StoreException {
    BlobId key = (BlobId) conversionMap.getOrDefault(blobId, blobId);
    assertEquals("There should be one message in the write set", 1, messageWriteSet.getMessageSetInfo().size());
    MessageInfo info = messageWriteSet.getMessageSetInfo().get(0);
    // verify stream
    ByteBuffer record = getDataInWriteSet(messageWriteSet, (int) info.getSize());
    int expectedSize = record.remaining();
    InputStream stream = new ByteBufferInputStream(record);
    // verify stream
    MessageFormatInputStreamTest.checkTtlUpdateMessage(stream, null, key, key.getAccountId(), key.getContainerId(), expiresAtMs, opTimeMs, (short) 0);
    // verify MessageInfo
    // since the record has been verified, the buffer size before verification is a good indicator of size
    MessageInfoTest.checkGetters(info, key, expectedSize, false, true, false, expiresAtMs, null, key.getAccountId(), key.getContainerId(), opTimeMs, (short) 0);
}

Example 37 with ByteBufferInputStream

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

the class CloudBlobStore method uploadWithRetries.

/**
 * Upload the supplied message buffer to a blob in the cloud destination.
 * @param blobId the {@link BlobId}.
 * @param messageBuf the byte buffer to upload.
 * @param bufferSize the size of the buffer.
 * @param blobMetadata the {@link CloudBlobMetadata} for the blob.
 * @return boolean indicating if the upload was completed.
 * @throws CloudStorageException if the upload failed.
 */
private boolean uploadWithRetries(BlobId blobId, ByteBuffer messageBuf, long bufferSize, CloudBlobMetadata blobMetadata) throws CloudStorageException {
    return requestAgent.doWithRetries(() -> {
        // Note: reset buffer and input stream each time through the loop
        messageBuf.rewind();
        InputStream uploadInputStream = new ByteBufferInputStream(messageBuf);
        return cloudDestination.uploadBlob(blobId, bufferSize, blobMetadata, uploadInputStream);
    }, "Upload", partitionId.toPathString());
}

Example 38 with ByteBufferInputStream

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

the class HelixClusterManagerTest method testInvalidPartitionId.

/**
 * Test that invalid partition id deserialization fails as expected.
 */
private void testInvalidPartitionId() {
    PartitionId partition = clusterManager.getWritablePartitionIds(null).get(0);
    try {
        byte[] fakePartition = Arrays.copyOf(partition.getBytes(), partition.getBytes().length);
        for (int i = fakePartition.length; i > fakePartition.length - Long.SIZE / Byte.SIZE; i--) {
            fakePartition[i - 1] = (byte) 0xff;
        }
        InputStream partitionStream = new ByteBufferInputStream(ByteBuffer.allocate(fakePartition.length));
        clusterManager.getPartitionIdFromStream(partitionStream);
        fail("partition id deserialization should have failed");
    } catch (IOException e) {
    // OK
    }
}

Example 39 with ByteBufferInputStream

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

the class MockRouterCallback method testErrorPrecedence.

/**
 * Help test error precedence.
 * @param serverErrorCodesInOrder the list of error codes to set the mock servers with.
 * @param expectedErrorCode the expected router error code for the operation.
 * @throws Exception
 */
private void testErrorPrecedence(ServerErrorCode[] serverErrorCodesInOrder, RouterErrorCode expectedErrorCode) throws Exception {
    NonBlockingRouter.currentOperationsCount.incrementAndGet();
    GetBlobInfoOperation op = new GetBlobInfoOperation(routerConfig, routerMetrics, mockClusterMap, responseHandler, blobId, options, null, routerCallback, kms, cryptoService, cryptoJobHandler, time);
    ArrayList<RequestInfo> requestListToFill = new ArrayList<>();
    requestRegistrationCallback.requestListToFill = requestListToFill;
    int i = 0;
    for (MockServer server : mockServerLayout.getMockServers()) {
        server.setServerErrorForAllRequests(serverErrorCodesInOrder[i++]);
    }
    while (!op.isOperationComplete()) {
        op.poll(requestRegistrationCallback);
        List<ResponseInfo> responses = sendAndWaitForResponses(requestListToFill);
        for (ResponseInfo responseInfo : responses) {
            GetResponse getResponse = responseInfo.getError() == null ? GetResponse.readFrom(new DataInputStream(new ByteBufferInputStream(responseInfo.getResponse())), mockClusterMap) : null;
            op.handleResponse(responseInfo, getResponse);
            if (op.isOperationComplete()) {
                break;
            }
        }
    }
    RouterException routerException = (RouterException) op.getOperationException();
    Assert.assertEquals(expectedErrorCode, routerException.getErrorCode());
}

Example 40 with ByteBufferInputStream

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

the class MessageFormatInputStreamTest method messageFormatRecordsTest.

private void messageFormatRecordsTest(short blobVersion, BlobType blobType, boolean useV2Header) throws IOException, MessageFormatException {
    StoreKey key = new MockId("id1");
    StoreKeyFactory keyFactory = new MockIdFactory();
    short accountId = Utils.getRandomShort(TestUtils.RANDOM);
    short containerId = Utils.getRandomShort(TestUtils.RANDOM);
    BlobProperties prop = new BlobProperties(10, "servid", accountId, containerId, false);
    byte[] encryptionKey = new byte[100];
    new Random().nextBytes(encryptionKey);
    byte[] usermetadata = new byte[1000];
    new Random().nextBytes(usermetadata);
    int blobContentSize = 2000;
    byte[] data = new byte[blobContentSize];
    new Random().nextBytes(data);
    long blobSize = -1;
    MessageFormatRecord.headerVersionToUse = useV2Header ? MessageFormatRecord.Message_Header_Version_V2 : MessageFormatRecord.Message_Header_Version_V1;
    if (blobVersion == MessageFormatRecord.Blob_Version_V1) {
        blobSize = MessageFormatRecord.Blob_Format_V1.getBlobRecordSize(blobContentSize);
    } else if (blobVersion == MessageFormatRecord.Blob_Version_V2 && blobType == BlobType.DataBlob) {
        blobSize = (int) MessageFormatRecord.Blob_Format_V2.getBlobRecordSize(blobContentSize);
    } else if (blobVersion == MessageFormatRecord.Blob_Version_V2 && blobType == BlobType.MetadataBlob) {
        ByteBuffer byteBufferBlob = MessageFormatTestUtils.getBlobContentForMetadataBlob(blobContentSize);
        data = byteBufferBlob.array();
        blobContentSize = data.length;
        blobSize = (int) MessageFormatRecord.Blob_Format_V2.getBlobRecordSize(blobContentSize);
    }
    ByteBufferInputStream stream = new ByteBufferInputStream(ByteBuffer.wrap(data));
    MessageFormatInputStream messageFormatStream = (blobVersion == MessageFormatRecord.Blob_Version_V2) ? new PutMessageFormatInputStream(key, ByteBuffer.wrap(encryptionKey), prop, ByteBuffer.wrap(usermetadata), stream, blobContentSize, blobType) : new PutMessageFormatBlobV1InputStream(key, prop, ByteBuffer.wrap(usermetadata), stream, blobContentSize, blobType);
    int headerSize = MessageFormatRecord.getHeaderSizeForVersion(useV2Header ? MessageFormatRecord.Message_Header_Version_V2 : MessageFormatRecord.Message_Header_Version_V1);
    int blobEncryptionKeySize = useV2Header ? MessageFormatRecord.BlobEncryptionKey_Format_V1.getBlobEncryptionKeyRecordSize(ByteBuffer.wrap(encryptionKey)) : 0;
    int blobPropertiesRecordSize = MessageFormatRecord.BlobProperties_Format_V1.getBlobPropertiesRecordSize(prop);
    int userMetadataSize = MessageFormatRecord.UserMetadata_Format_V1.getUserMetadataSize(ByteBuffer.wrap(usermetadata));
    Assert.assertEquals(messageFormatStream.getSize(), headerSize + blobEncryptionKeySize + blobPropertiesRecordSize + userMetadataSize + blobSize + key.sizeInBytes());
    // verify header
    byte[] headerOutput = new byte[headerSize];
    messageFormatStream.read(headerOutput);
    ByteBuffer headerBuf = ByteBuffer.wrap(headerOutput);
    Assert.assertEquals(useV2Header ? MessageFormatRecord.Message_Header_Version_V2 : MessageFormatRecord.Message_Header_Version_V1, headerBuf.getShort());
    Assert.assertEquals(blobEncryptionKeySize + blobPropertiesRecordSize + userMetadataSize + blobSize, headerBuf.getLong());
    if (useV2Header) {
        Assert.assertEquals(headerSize + key.sizeInBytes(), headerBuf.getInt());
        Assert.assertEquals(headerSize + key.sizeInBytes() + blobEncryptionKeySize, headerBuf.getInt());
        Assert.assertEquals(MessageFormatRecord.Message_Header_Invalid_Relative_Offset, headerBuf.getInt());
        Assert.assertEquals(headerSize + key.sizeInBytes() + blobEncryptionKeySize + blobPropertiesRecordSize, headerBuf.getInt());
        Assert.assertEquals(headerSize + key.sizeInBytes() + blobEncryptionKeySize + blobPropertiesRecordSize + userMetadataSize, headerBuf.getInt());
    } else {
        Assert.assertEquals(headerSize + key.sizeInBytes(), headerBuf.getInt());
        Assert.assertEquals(MessageFormatRecord.Message_Header_Invalid_Relative_Offset, headerBuf.getInt());
        Assert.assertEquals(headerSize + key.sizeInBytes() + blobPropertiesRecordSize, headerBuf.getInt());
        Assert.assertEquals(headerSize + key.sizeInBytes() + blobPropertiesRecordSize + userMetadataSize, headerBuf.getInt());
    }
    Crc32 crc = new Crc32();
    crc.update(headerOutput, 0, headerSize - MessageFormatRecord.Crc_Size);
    Assert.assertEquals(crc.getValue(), headerBuf.getLong());
    // verify handle
    byte[] handleOutput = new byte[key.sizeInBytes()];
    ByteBuffer handleOutputBuf = ByteBuffer.wrap(handleOutput);
    messageFormatStream.read(handleOutput);
    byte[] dest = new byte[key.sizeInBytes()];
    handleOutputBuf.get(dest);
    Assert.assertArrayEquals(dest, key.toBytes());
    // verify encryption key
    if (useV2Header) {
        byte[] blobEncryptionKeyOutput = new byte[blobEncryptionKeySize];
        ByteBuffer blobEncryptionKeyBuf = ByteBuffer.wrap(blobEncryptionKeyOutput);
        messageFormatStream.read(blobEncryptionKeyOutput);
        Assert.assertEquals(blobEncryptionKeyBuf.getShort(), MessageFormatRecord.Blob_Encryption_Key_V1);
        Assert.assertEquals(blobEncryptionKeyBuf.getInt(), 100);
        dest = new byte[100];
        blobEncryptionKeyBuf.get(dest);
        Assert.assertArrayEquals(dest, encryptionKey);
        crc = new Crc32();
        crc.update(blobEncryptionKeyOutput, 0, blobEncryptionKeySize - MessageFormatRecord.Crc_Size);
        Assert.assertEquals(crc.getValue(), blobEncryptionKeyBuf.getLong());
    }
    // verify blob properties
    byte[] blobPropertiesOutput = new byte[blobPropertiesRecordSize];
    ByteBuffer blobPropertiesBuf = ByteBuffer.wrap(blobPropertiesOutput);
    messageFormatStream.read(blobPropertiesOutput);
    Assert.assertEquals(blobPropertiesBuf.getShort(), 1);
    BlobProperties propOutput = BlobPropertiesSerDe.getBlobPropertiesFromStream(new DataInputStream(new ByteBufferInputStream(blobPropertiesBuf)));
    Assert.assertEquals(10, propOutput.getBlobSize());
    Assert.assertEquals("servid", propOutput.getServiceId());
    Assert.assertEquals("AccountId mismatch", accountId, propOutput.getAccountId());
    Assert.assertEquals("ContainerId mismatch", containerId, propOutput.getContainerId());
    crc = new Crc32();
    crc.update(blobPropertiesOutput, 0, blobPropertiesRecordSize - MessageFormatRecord.Crc_Size);
    Assert.assertEquals(crc.getValue(), blobPropertiesBuf.getLong());
    // verify user metadata
    byte[] userMetadataOutput = new byte[userMetadataSize];
    ByteBuffer userMetadataBuf = ByteBuffer.wrap(userMetadataOutput);
    messageFormatStream.read(userMetadataOutput);
    Assert.assertEquals(userMetadataBuf.getShort(), 1);
    Assert.assertEquals(userMetadataBuf.getInt(), 1000);
    dest = new byte[1000];
    userMetadataBuf.get(dest);
    Assert.assertArrayEquals(dest, usermetadata);
    crc = new Crc32();
    crc.update(userMetadataOutput, 0, userMetadataSize - MessageFormatRecord.Crc_Size);
    Assert.assertEquals(crc.getValue(), userMetadataBuf.getLong());
    // verify blob
    CrcInputStream crcstream = new CrcInputStream(messageFormatStream);
    DataInputStream streamData = new DataInputStream(crcstream);
    Assert.assertEquals(streamData.readShort(), blobVersion);
    if (blobVersion == MessageFormatRecord.Blob_Version_V2) {
        Assert.assertEquals(streamData.readShort(), blobType.ordinal());
    }
    Assert.assertEquals(streamData.readLong(), blobContentSize);
    for (int i = 0; i < blobContentSize; i++) {
        Assert.assertEquals((byte) streamData.read(), data[i]);
    }
    long crcVal = crcstream.getValue();
    Assert.assertEquals(crcVal, streamData.readLong());
    // Verify Blob All
    stream = new ByteBufferInputStream(ByteBuffer.wrap(data));
    messageFormatStream = (blobVersion == MessageFormatRecord.Blob_Version_V2) ? new PutMessageFormatInputStream(key, ByteBuffer.wrap(encryptionKey), prop, ByteBuffer.wrap(usermetadata), stream, blobContentSize, blobType) : new PutMessageFormatBlobV1InputStream(key, prop, ByteBuffer.wrap(usermetadata), stream, blobContentSize, blobType);
    int totalSize;
    if (useV2Header) {
        totalSize = headerSize + key.sizeInBytes() + blobEncryptionKeySize + blobPropertiesRecordSize + userMetadataSize + (int) blobSize;
    } else {
        totalSize = headerSize + key.sizeInBytes() + blobPropertiesRecordSize + userMetadataSize + (int) blobSize;
    }
    ByteBuffer allBuf = ByteBuffer.allocate(totalSize);
    messageFormatStream.read(allBuf.array());
    BlobAll blobAll = MessageFormatRecord.deserializeBlobAll(new ByteBufferInputStream(allBuf), keyFactory);
    Assert.assertEquals(key, blobAll.getStoreKey());
    Assert.assertArrayEquals(usermetadata, blobAll.getBlobInfo().getUserMetadata());
    Assert.assertEquals(blobContentSize, blobAll.getBlobData().getSize());
    Assert.assertEquals(blobType, blobAll.getBlobData().getBlobType());
    if (useV2Header) {
        Assert.assertEquals(ByteBuffer.wrap(encryptionKey), blobAll.getBlobEncryptionKey());
    } else {
        Assert.assertEquals(null, blobAll.getBlobEncryptionKey());
    }
    Assert.assertEquals(ByteBuffer.wrap(data), blobAll.getBlobData().getStream().getByteBuffer());
}