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

Example 66 with ByteBufferInputStream

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

the class GetBlobOperationTest method getBlobBufferFromServer.

/**
 * @param server the server hosting the blob.
 * @return the ByteBuffer for the blob contents on the specified server.
 * @throws IOException
 */
private ByteBuffer getBlobBufferFromServer(MockServer server) throws IOException {
    PartitionRequestInfo requestInfo = new PartitionRequestInfo(blobId.getPartition(), Collections.singletonList(blobId));
    GetRequest getRequest = new GetRequest(1, "assertBlobReadSuccess", MessageFormatFlags.All, Collections.singletonList(requestInfo), GetOption.None);
    GetResponse getResponse = server.makeGetResponse(getRequest, ServerErrorCode.No_Error);
    getRequest.release();
    // simulating server sending response over the wire
    ByteBufferChannel channel = new ByteBufferChannel(ByteBuffer.allocate((int) getResponse.sizeInBytes()));
    getResponse.writeTo(channel);
    getResponse.release();
    // simulating the Router receiving the data from the wire
    ByteBuffer data = channel.getBuffer();
    data.flip();
    DataInputStream stream = new DataInputStream(new ByteBufferInputStream(data));
    // read off the size because GetResponse.readFrom() expects that this be read off
    stream.readLong();
    // construct a GetResponse as the Router would have
    getResponse = GetResponse.readFrom(stream, mockClusterMap);
    byte[] blobData = Utils.readBytesFromStream(getResponse.getInputStream(), (int) getResponse.getPartitionResponseInfoList().get(0).getMessageInfoList().get(0).getSize());
    // set the put content buf to the data in the stream
    return ByteBuffer.wrap(blobData);
}

Example 67 with ByteBufferInputStream

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

the class StoredBlob method makeGetResponse.

/**
 * Make a {@link GetResponse} for the given {@link GetRequest} for which the given {@link ServerErrorCode} was
 * encountered. The request could be for BlobInfo or for Blob (the only two options that the router would request
 * for).
 * @param getRequest the {@link GetRequest} for which the response is being constructed.
 * @param getError the {@link ServerErrorCode} that was encountered.
 * @return the constructed {@link GetResponse}
 * @throws IOException if there was an error constructing the response.
 */
GetResponse makeGetResponse(GetRequest getRequest, ServerErrorCode getError) throws IOException {
    GetResponse getResponse;
    if (getError == ServerErrorCode.No_Error) {
        List<PartitionRequestInfo> infos = getRequest.getPartitionInfoList();
        if (infos.size() != 1 || infos.get(0).getBlobIds().size() != 1) {
            getError = ServerErrorCode.Unknown_Error;
        }
    }
    ServerErrorCode serverError;
    ServerErrorCode partitionError;
    boolean isDataBlob = false;
    try {
        String id = getRequest.getPartitionInfoList().get(0).getBlobIds().get(0).getID();
        isDataBlob = blobs.get(id).type == BlobType.DataBlob;
    } catch (Exception ignored) {
    }
    if (!getErrorOnDataBlobOnly || isDataBlob) {
        // set it in the partitionResponseInfo
        if (getError == ServerErrorCode.No_Error || getError == ServerErrorCode.Blob_Expired || getError == ServerErrorCode.Blob_Deleted || getError == ServerErrorCode.Blob_Not_Found || getError == ServerErrorCode.Blob_Authorization_Failure || getError == ServerErrorCode.Disk_Unavailable) {
            partitionError = getError;
            serverError = ServerErrorCode.No_Error;
        } else {
            serverError = getError;
            // does not matter - this will not be checked if serverError is not No_Error.
            partitionError = ServerErrorCode.No_Error;
        }
    } else {
        serverError = ServerErrorCode.No_Error;
        partitionError = ServerErrorCode.No_Error;
    }
    if (serverError == ServerErrorCode.No_Error) {
        int byteBufferSize;
        ByteBuffer byteBuffer;
        StoreKey key = getRequest.getPartitionInfoList().get(0).getBlobIds().get(0);
        short accountId = Account.UNKNOWN_ACCOUNT_ID;
        short containerId = Container.UNKNOWN_CONTAINER_ID;
        long operationTimeMs = Utils.Infinite_Time;
        StoredBlob blob = blobs.get(key.getID());
        ServerErrorCode processedError = errorForGet(key.getID(), blob, getRequest);
        MessageMetadata msgMetadata = null;
        if (processedError == ServerErrorCode.No_Error) {
            ByteBuffer buf = blobs.get(key.getID()).serializedSentPutRequest.duplicate();
            // read off the size
            buf.getLong();
            // read off the type.
            buf.getShort();
            PutRequest originalBlobPutReq = PutRequest.readFrom(new DataInputStream(new ByteBufferInputStream(buf)), clusterMap);
            switch(getRequest.getMessageFormatFlag()) {
                case BlobInfo:
                    BlobProperties blobProperties = originalBlobPutReq.getBlobProperties();
                    accountId = blobProperties.getAccountId();
                    containerId = blobProperties.getContainerId();
                    operationTimeMs = blobProperties.getCreationTimeInMs();
                    ByteBuffer userMetadata = originalBlobPutReq.getUsermetadata();
                    byteBufferSize = MessageFormatRecord.BlobProperties_Format_V1.getBlobPropertiesRecordSize(blobProperties) + MessageFormatRecord.UserMetadata_Format_V1.getUserMetadataSize(userMetadata);
                    byteBuffer = ByteBuffer.allocate(byteBufferSize);
                    if (originalBlobPutReq.getBlobEncryptionKey() != null) {
                        msgMetadata = new MessageMetadata(originalBlobPutReq.getBlobEncryptionKey().duplicate());
                    }
                    MessageFormatRecord.BlobProperties_Format_V1.serializeBlobPropertiesRecord(byteBuffer, blobProperties);
                    MessageFormatRecord.UserMetadata_Format_V1.serializeUserMetadataRecord(byteBuffer, userMetadata);
                    break;
                case Blob:
                    switch(blobFormatVersion) {
                        case MessageFormatRecord.Blob_Version_V2:
                            if (originalBlobPutReq.getBlobEncryptionKey() != null) {
                                msgMetadata = new MessageMetadata(originalBlobPutReq.getBlobEncryptionKey().duplicate());
                            }
                            byteBufferSize = (int) MessageFormatRecord.Blob_Format_V2.getBlobRecordSize((int) originalBlobPutReq.getBlobSize());
                            byteBuffer = ByteBuffer.allocate(byteBufferSize);
                            MessageFormatRecord.Blob_Format_V2.serializePartialBlobRecord(byteBuffer, (int) originalBlobPutReq.getBlobSize(), originalBlobPutReq.getBlobType());
                            break;
                        case MessageFormatRecord.Blob_Version_V1:
                            byteBufferSize = (int) MessageFormatRecord.Blob_Format_V1.getBlobRecordSize((int) originalBlobPutReq.getBlobSize());
                            byteBuffer = ByteBuffer.allocate(byteBufferSize);
                            MessageFormatRecord.Blob_Format_V1.serializePartialBlobRecord(byteBuffer, (int) originalBlobPutReq.getBlobSize());
                            break;
                        default:
                            throw new IllegalStateException("Blob format version " + blobFormatVersion + " not supported.");
                    }
                    byteBuffer.put(Utils.readBytesFromStream(originalBlobPutReq.getBlobStream(), (int) originalBlobPutReq.getBlobSize()));
                    Crc32 crc = new Crc32();
                    crc.update(byteBuffer.array(), 0, byteBuffer.position());
                    byteBuffer.putLong(crc.getValue());
                    break;
                case All:
                    blobProperties = originalBlobPutReq.getBlobProperties();
                    accountId = blobProperties.getAccountId();
                    containerId = blobProperties.getContainerId();
                    userMetadata = originalBlobPutReq.getUsermetadata();
                    operationTimeMs = originalBlobPutReq.getBlobProperties().getCreationTimeInMs();
                    int blobHeaderSize = MessageFormatRecord.MessageHeader_Format_V2.getHeaderSize();
                    int blobEncryptionRecordSize = originalBlobPutReq.getBlobEncryptionKey() != null ? MessageFormatRecord.BlobEncryptionKey_Format_V1.getBlobEncryptionKeyRecordSize(originalBlobPutReq.getBlobEncryptionKey().duplicate()) : 0;
                    int blobPropertiesSize = MessageFormatRecord.BlobProperties_Format_V1.getBlobPropertiesRecordSize(blobProperties);
                    int userMetadataSize = MessageFormatRecord.UserMetadata_Format_V1.getUserMetadataSize(userMetadata);
                    int blobInfoSize = blobPropertiesSize + userMetadataSize;
                    int blobRecordSize;
                    switch(blobFormatVersion) {
                        case MessageFormatRecord.Blob_Version_V2:
                            blobRecordSize = (int) MessageFormatRecord.Blob_Format_V2.getBlobRecordSize((int) originalBlobPutReq.getBlobSize());
                            break;
                        case MessageFormatRecord.Blob_Version_V1:
                            blobRecordSize = (int) MessageFormatRecord.Blob_Format_V1.getBlobRecordSize((int) originalBlobPutReq.getBlobSize());
                            break;
                        default:
                            throw new IllegalStateException("Blob format version " + blobFormatVersion + " not supported.");
                    }
                    byteBufferSize = blobHeaderSize + key.sizeInBytes() + blobEncryptionRecordSize + blobInfoSize + blobRecordSize;
                    byteBuffer = ByteBuffer.allocate(byteBufferSize);
                    try {
                        MessageFormatRecord.MessageHeader_Format_V2.serializeHeader(byteBuffer, blobEncryptionRecordSize + blobInfoSize + blobRecordSize, originalBlobPutReq.getBlobEncryptionKey() == null ? Message_Header_Invalid_Relative_Offset : blobHeaderSize + key.sizeInBytes(), blobHeaderSize + key.sizeInBytes() + blobEncryptionRecordSize, Message_Header_Invalid_Relative_Offset, blobHeaderSize + key.sizeInBytes() + blobEncryptionRecordSize + blobPropertiesSize, blobHeaderSize + key.sizeInBytes() + blobEncryptionRecordSize + blobInfoSize);
                    } catch (MessageFormatException e) {
                        e.printStackTrace();
                    }
                    byteBuffer.put(key.toBytes());
                    if (originalBlobPutReq.getBlobEncryptionKey() != null) {
                        MessageFormatRecord.BlobEncryptionKey_Format_V1.serializeBlobEncryptionKeyRecord(byteBuffer, originalBlobPutReq.getBlobEncryptionKey().duplicate());
                        msgMetadata = new MessageMetadata(originalBlobPutReq.getBlobEncryptionKey().duplicate());
                    }
                    MessageFormatRecord.BlobProperties_Format_V1.serializeBlobPropertiesRecord(byteBuffer, blobProperties);
                    MessageFormatRecord.UserMetadata_Format_V1.serializeUserMetadataRecord(byteBuffer, userMetadata);
                    int blobRecordStart = byteBuffer.position();
                    switch(blobFormatVersion) {
                        case MessageFormatRecord.Blob_Version_V2:
                            MessageFormatRecord.Blob_Format_V2.serializePartialBlobRecord(byteBuffer, (int) originalBlobPutReq.getBlobSize(), originalBlobPutReq.getBlobType());
                            break;
                        case MessageFormatRecord.Blob_Version_V1:
                            MessageFormatRecord.Blob_Format_V1.serializePartialBlobRecord(byteBuffer, (int) originalBlobPutReq.getBlobSize());
                            break;
                        default:
                            throw new IllegalStateException("Blob format version " + blobFormatVersion + " not supported.");
                    }
                    byteBuffer.put(Utils.readBytesFromStream(originalBlobPutReq.getBlobStream(), (int) originalBlobPutReq.getBlobSize()));
                    crc = new Crc32();
                    crc.update(byteBuffer.array(), blobRecordStart, blobRecordSize - MessageFormatRecord.Crc_Size);
                    byteBuffer.putLong(crc.getValue());
                    break;
                default:
                    throw new IOException("GetRequest flag is not supported: " + getRequest.getMessageFormatFlag());
            }
        } else if (processedError == ServerErrorCode.Blob_Deleted) {
            if (partitionError == ServerErrorCode.No_Error) {
                partitionError = ServerErrorCode.Blob_Deleted;
            }
            byteBuffer = ByteBuffer.allocate(0);
            byteBufferSize = 0;
        } else if (processedError == ServerErrorCode.Blob_Expired) {
            if (partitionError == ServerErrorCode.No_Error) {
                partitionError = ServerErrorCode.Blob_Expired;
            }
            byteBuffer = ByteBuffer.allocate(0);
            byteBufferSize = 0;
        } else if (processedError == ServerErrorCode.Blob_Authorization_Failure) {
            if (partitionError == ServerErrorCode.No_Error) {
                partitionError = ServerErrorCode.Blob_Authorization_Failure;
            }
            byteBuffer = ByteBuffer.allocate(0);
            byteBufferSize = 0;
        } else {
            if (partitionError == ServerErrorCode.No_Error) {
                partitionError = ServerErrorCode.Blob_Not_Found;
            }
            byteBuffer = ByteBuffer.allocate(0);
            byteBufferSize = 0;
        }
        byteBuffer.flip();
        ByteBufferSend responseSend = new ByteBufferSend(byteBuffer);
        List<MessageInfo> messageInfoList = new ArrayList<>();
        List<MessageMetadata> messageMetadataList = new ArrayList<>();
        List<PartitionResponseInfo> partitionResponseInfoList = new ArrayList<PartitionResponseInfo>();
        if (partitionError == ServerErrorCode.No_Error) {
            messageInfoList.add(new MessageInfo(key, byteBufferSize, false, blob.isTtlUpdated(), blob.isUndeleted(), blob.expiresAt, null, accountId, containerId, operationTimeMs, blob.lifeVersion));
            messageMetadataList.add(msgMetadata);
        }
        PartitionResponseInfo partitionResponseInfo = partitionError == ServerErrorCode.No_Error ? new PartitionResponseInfo(getRequest.getPartitionInfoList().get(0).getPartition(), messageInfoList, messageMetadataList) : new PartitionResponseInfo(getRequest.getPartitionInfoList().get(0).getPartition(), partitionError);
        partitionResponseInfoList.add(partitionResponseInfo);
        getResponse = new GetResponse(getRequest.getCorrelationId(), getRequest.getClientId(), partitionResponseInfoList, responseSend, serverError);
    } else {
        getResponse = new GetResponse(getRequest.getCorrelationId(), getRequest.getClientId(), new ArrayList<PartitionResponseInfo>(), new ByteBufferSend(ByteBuffer.allocate(0)), serverError);
    }
    return getResponse;
}

Example 68 with ByteBufferInputStream

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

the class ServerTestUtil method getPutMessageInputStreamForBlob.

/**
 * Create {@link PutMessageFormatInputStream} for a blob with given {@link BlobId} and update {@code blobIdToSizeMap}.
 * @param blobId {@link BlobId} object.
 * @param blobSize size of blob.
 * @param blobIdToSizeMap {@link Map} of {@link BlobId} to size of blob uploaded.
 * @return {@link PutMessageFormatInputStream} object.
 * @throws Exception
 */
static PutMessageFormatInputStream getPutMessageInputStreamForBlob(BlobId blobId, int blobSize, Map<BlobId, Integer> blobIdToSizeMap, short accountId, short containerId) throws Exception {
    int userMetaDataSize = 100;
    byte[] userMetadata = new byte[userMetaDataSize];
    TestUtils.RANDOM.nextBytes(userMetadata);
    byte[] data = new byte[blobSize];
    BlobProperties blobProperties = new BlobProperties(blobSize, "serviceid1", null, null, false, Utils.Infinite_Time, accountId, containerId, false, null, null, null);
    TestUtils.RANDOM.nextBytes(data);
    blobIdToSizeMap.put(blobId, blobSize);
    return new PutMessageFormatInputStream(blobId, null, blobProperties, ByteBuffer.wrap(userMetadata), new ByteBufferInputStream(ByteBuffer.wrap(data)), blobSize);
}

Example 69 with ByteBufferInputStream

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

the class IndexSegment method getKeyAt.

private StoreKey getKeyAt(ByteBuffer mmap, int index) throws StoreException {
    StoreKey storeKey = null;
    try {
        mmap.position(firstKeyRelativeOffset + index * persistedEntrySize);
        storeKey = factory.getStoreKey(new DataInputStream(new ByteBufferInputStream(mmap)));
    } catch (InternalError e) {
        StoreErrorCodes errorCode = StoreException.resolveErrorCode(e);
        throw new StoreException("Internal " + errorCode.toString() + " while trying to get store key", e, errorCode);
    } catch (IOException e) {
        StoreErrorCodes errorCode = StoreException.resolveErrorCode(e);
        throw new StoreException(errorCode.toString() + " while trying to get store key", e, errorCode);
    } catch (Throwable t) {
        throw new StoreException("Unknown error while trying to get store key ", t, StoreErrorCodes.Unknown_Error);
    }
    return storeKey;
}

Example 70 with ByteBufferInputStream

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

the class IndexSegment method map.

/**
 * Maps the segment of index either as a memory map or a in memory buffer depending on config.
 * @throws StoreException if there are problems with the index
 */
private void map() throws StoreException {
    try (RandomAccessFile raf = new RandomAccessFile(indexFile, "r")) {
        switch(config.storeIndexMemState) {
            case IN_DIRECT_MEM:
                serEntries = readFileIntoBuffer(raf, true);
                break;
            case IN_HEAP_MEM:
                serEntries = readFileIntoBuffer(raf, false);
                break;
            default:
                MappedByteBuffer buf = raf.getChannel().map(FileChannel.MapMode.READ_ONLY, 0, indexFile.length());
                if (config.storeIndexMemState == IndexMemState.MMAP_WITH_FORCE_LOAD) {
                    buf.load();
                }
                serEntries = buf;
                break;
        }
        checkDataIntegrity();
        serEntries.position(0);
        setVersion(serEntries.getShort());
        StoreKey storeKey;
        int keySize;
        PersistentIndex.IndexEntryType resetKeyType;
        short resetKeyLifeVersion = UNINITIALIZED_RESET_KEY_VERSION;
        switch(getVersion()) {
            case PersistentIndex.VERSION_0:
                indexSizeExcludingEntries = VERSION_FIELD_LENGTH + KEY_OR_ENTRY_SIZE_FIELD_LENGTH + VALUE_SIZE_FIELD_LENGTH + LOG_END_OFFSET_FIELD_LENGTH + CRC_FIELD_LENGTH;
                keySize = serEntries.getInt();
                valueSize = serEntries.getInt();
                persistedEntrySize = keySize + valueSize;
                endOffset.set(new Offset(startOffset.getName(), serEntries.getLong()));
                lastModifiedTimeSec.set(indexFile.lastModified() / 1000);
                firstKeyRelativeOffset = indexSizeExcludingEntries - CRC_FIELD_LENGTH;
                break;
            case PersistentIndex.VERSION_1:
                keySize = serEntries.getInt();
                valueSize = serEntries.getInt();
                persistedEntrySize = keySize + valueSize;
                endOffset.set(new Offset(startOffset.getName(), serEntries.getLong()));
                lastModifiedTimeSec.set(serEntries.getLong());
                storeKey = factory.getStoreKey(new DataInputStream(new ByteBufferInputStream(serEntries)));
                resetKeyType = PersistentIndex.IndexEntryType.values()[serEntries.getShort()];
                resetKeyInfo = new ResetKeyInfo(storeKey, resetKeyType, resetKeyLifeVersion);
                indexSizeExcludingEntries = VERSION_FIELD_LENGTH + KEY_OR_ENTRY_SIZE_FIELD_LENGTH + VALUE_SIZE_FIELD_LENGTH + LOG_END_OFFSET_FIELD_LENGTH + CRC_FIELD_LENGTH + LAST_MODIFIED_TIME_FIELD_LENGTH + storeKey.sizeInBytes() + RESET_KEY_TYPE_FIELD_LENGTH;
                firstKeyRelativeOffset = indexSizeExcludingEntries - CRC_FIELD_LENGTH;
                break;
            case PersistentIndex.VERSION_2:
            case PersistentIndex.VERSION_3:
                persistedEntrySize = serEntries.getInt();
                valueSize = serEntries.getInt();
                endOffset.set(new Offset(startOffset.getName(), serEntries.getLong()));
                lastModifiedTimeSec.set(serEntries.getLong());
                storeKey = factory.getStoreKey(new DataInputStream(new ByteBufferInputStream(serEntries)));
                resetKeyType = PersistentIndex.IndexEntryType.values()[serEntries.getShort()];
                resetKeyInfo = new ResetKeyInfo(storeKey, resetKeyType, resetKeyLifeVersion);
                indexSizeExcludingEntries = VERSION_FIELD_LENGTH + KEY_OR_ENTRY_SIZE_FIELD_LENGTH + VALUE_SIZE_FIELD_LENGTH + LOG_END_OFFSET_FIELD_LENGTH + CRC_FIELD_LENGTH + LAST_MODIFIED_TIME_FIELD_LENGTH + storeKey.sizeInBytes() + RESET_KEY_TYPE_FIELD_LENGTH;
                firstKeyRelativeOffset = indexSizeExcludingEntries - CRC_FIELD_LENGTH;
                break;
            case PersistentIndex.VERSION_4:
                persistedEntrySize = serEntries.getInt();
                valueSize = serEntries.getInt();
                endOffset.set(new Offset(startOffset.getName(), serEntries.getLong()));
                lastModifiedTimeSec.set(serEntries.getLong());
                storeKey = factory.getStoreKey(new DataInputStream(new ByteBufferInputStream(serEntries)));
                resetKeyType = PersistentIndex.IndexEntryType.values()[serEntries.getShort()];
                resetKeyLifeVersion = serEntries.getShort();
                resetKeyInfo = new ResetKeyInfo(storeKey, resetKeyType, resetKeyLifeVersion);
                indexSizeExcludingEntries = VERSION_FIELD_LENGTH + KEY_OR_ENTRY_SIZE_FIELD_LENGTH + VALUE_SIZE_FIELD_LENGTH + LOG_END_OFFSET_FIELD_LENGTH + CRC_FIELD_LENGTH + LAST_MODIFIED_TIME_FIELD_LENGTH + storeKey.sizeInBytes() + RESET_KEY_TYPE_FIELD_LENGTH + RESET_KEY_VERSION_FIELD_LENGTH;
                firstKeyRelativeOffset = indexSizeExcludingEntries - CRC_FIELD_LENGTH;
                break;
            default:
                throw new StoreException("IndexSegment : " + indexFile.getAbsolutePath() + " unknown version in index file", StoreErrorCodes.Index_Version_Error);
        }
        index = null;
    } catch (FileNotFoundException e) {
        throw new StoreException("File not found while mapping the segment of index", e, StoreErrorCodes.File_Not_Found);
    } catch (IOException e) {
        StoreErrorCodes errorCode = StoreException.resolveErrorCode(e);
        throw new StoreException(errorCode.toString() + " while mapping the segment of index", e, errorCode);
    }
}