Examples with BlobMetadataContext org.apache.samza.system.azureblob.utils.BlobMetadataContext used on opensource projects

Example 1 with BlobMetadataContext

use of org.apache.samza.system.azureblob.utils.BlobMetadataContext in project samza by apache.

the class TestAzureBlobOutputStream method setup.

@Before
public void setup() throws Exception {
    threadPool = new ThreadPoolExecutor(1, 1, 60, TimeUnit.SECONDS, new LinkedBlockingDeque<Runnable>());
    mockByteArrayOutputStream = spy(new ByteArrayOutputStream(THRESHOLD));
    mockBlobAsyncClient = PowerMockito.mock(BlockBlobAsyncClient.class);
    when(mockBlobAsyncClient.getBlobUrl()).thenReturn("https://samza.blob.core.windows.net/fake-blob-url");
    mockMetrics = mock(AzureBlobWriterMetrics.class);
    mockCompression = mock(Compression.class);
    doReturn(COMPRESSED_BYTES).when(mockCompression).compress(BYTES);
    doAnswer(invocation -> {
        BlobMetadataContext blobMetadataContext = invocation.getArgumentAt(0, BlobMetadataContext.class);
        String streamName = blobMetadataContext.getStreamName();
        Long blobSize = blobMetadataContext.getBlobSize();
        Long numberOfRecords = blobMetadataContext.getNumberOfMessagesInBlob();
        Map<String, String> metadataProperties = new HashMap<>();
        metadataProperties.put(BLOB_STREAM_NAME_METADATA, streamName);
        metadataProperties.put(BLOB_RAW_SIZE_BYTES_METADATA, Long.toString(blobSize));
        metadataProperties.put(BLOB_RECORD_NUMBER_METADATA, Long.toString(numberOfRecords));
        return metadataProperties;
    }).when(mockBlobMetadataGenerator).getBlobMetadata(anyObject());
    azureBlobOutputStream = spy(new AzureBlobOutputStream(mockBlobAsyncClient, threadPool, mockMetrics, blobMetadataGeneratorFactory, blobMetadataGeneratorConfig, FAKE_STREAM, 60000, THRESHOLD, mockByteArrayOutputStream, mockCompression));
    doNothing().when(azureBlobOutputStream).commitBlob(any(ArrayList.class), anyMap());
    doNothing().when(azureBlobOutputStream).stageBlock(anyString(), any(ByteBuffer.class), anyInt());
    doNothing().when(azureBlobOutputStream).clearAndMarkClosed();
    doReturn(mockBlobMetadataGenerator).when(azureBlobOutputStream).getBlobMetadataGenerator();
}

Example 2 with BlobMetadataContext

use of org.apache.samza.system.azureblob.utils.BlobMetadataContext in project samza by apache.

the class AzureBlobOutputStream method close.

/**
 * This api waits for all pending upload (stageBlock task) futures to finish.
 * It then synchronously commits the list of blocks to persist the actual blob on storage.
 * Note: this method does not invoke flush and flush has to be explicitly called before close.
 * @throws IllegalStateException when
 *       - when closing an already closed stream
 * @throws RuntimeException when
 *       - byteArrayOutputStream.close fails or
 *       - any of the pending uploads fails or
 *       - blob's commitBlockList fails
 * throws ClassNotFoundException or IllegalAccessException or InstantiationException
 *       - while creating an instance of BlobMetadataGenerator
 */
@Override
public synchronized void close() {
    if (isClosed) {
        LOG.info("{}: already closed", blobAsyncClient.getBlobUrl().toString());
        return;
    }
    LOG.info("{}: Close", blobAsyncClient.getBlobUrl().toString());
    try {
        if (byteArrayOutputStream.isPresent()) {
            byteArrayOutputStream.get().close();
        }
        if (blockList.size() == 0) {
            return;
        }
        CompletableFuture<Void> future = CompletableFuture.allOf(pendingUpload.toArray(new CompletableFuture[0]));
        LOG.info("Closing blob: {} PendingUpload:{} ", blobAsyncClient.getBlobUrl().toString(), pendingUpload.size());
        future.get((long) flushTimeoutMs, TimeUnit.MILLISECONDS);
        LOG.info("For blob: {} committing blockList size:{}", blobAsyncClient.getBlobUrl().toString(), blockList.size());
        metrics.updateAzureCommitMetrics();
        BlobMetadataGenerator blobMetadataGenerator = getBlobMetadataGenerator();
        commitBlob(blockList, blobMetadataGenerator.getBlobMetadata(new BlobMetadataContext(streamName, totalUploadedBlockSize, totalNumberOfRecordsInBlob)));
    } catch (Exception e) {
        String msg = String.format("Close blob %s failed with exception. Total pending sends %d", blobAsyncClient.getBlobUrl().toString(), pendingUpload.size());
        throw new AzureException(msg, e);
    } finally {
        clearAndMarkClosed();
    }
}