Examples with BlobStoreManager org.apache.samza.storage.blobstore.BlobStoreManager used on opensource projects

Example 1 with BlobStoreManager

use of org.apache.samza.storage.blobstore.BlobStoreManager in project samza by apache.

the class TestBlobStoreUtil method testCleanup.

@Test
public void testCleanup() throws IOException, ExecutionException, InterruptedException {
    BlobStoreManager blobStoreManager = mock(BlobStoreManager.class);
    // File, dir and recursive dir added, retained and removed in local
    // Using unique file names since test util uses only the file name (leaf node)
    // as the mock blob id, not the full file path.
    String local = "[a, c, z/1, y/2, p/m/3, q/n/4]";
    String remote = "[a, b, z/1, x/5, p/m/3, r/o/6]";
    String expectedRemoved = "[b, 5, 6]";
    // keep only the last character (the file name).
    SortedSet<String> expectedRemovedFiles = BlobStoreTestUtil.getExpected(expectedRemoved);
    // Set up environment
    Path localSnapshotDir = BlobStoreTestUtil.createLocalDir(local);
    String basePath = localSnapshotDir.toAbsolutePath().toString();
    DirIndex remoteSnapshotDir = BlobStoreTestUtil.createDirIndex(remote);
    SnapshotMetadata snapshotMetadata = new SnapshotMetadata(checkpointId, jobName, jobId, taskName, storeName);
    DirDiff dirDiff = DirDiffUtil.getDirDiff(localSnapshotDir.toFile(), remoteSnapshotDir, (localFile, remoteFile) -> localFile.getName().equals(remoteFile.getFileName()));
    BlobStoreUtil blobStoreUtil = new BlobStoreUtil(blobStoreManager, EXECUTOR, null, null);
    when(blobStoreManager.put(any(InputStream.class), any(Metadata.class))).thenReturn(CompletableFuture.completedFuture("blobId"));
    CompletionStage<DirIndex> dirIndexFuture = blobStoreUtil.putDir(dirDiff, snapshotMetadata);
    DirIndex dirIndex = null;
    try {
        // should be already complete. if not, future composition in putDir is broken.
        dirIndex = dirIndexFuture.toCompletableFuture().get(0, TimeUnit.MILLISECONDS);
    } catch (TimeoutException e) {
        fail("Future returned from putDir should be already complete.");
    }
    // Set up mocks
    SortedSet<String> allDeleted = new TreeSet<>();
    when(blobStoreManager.delete(anyString(), any(Metadata.class))).thenAnswer((Answer<CompletableFuture<Void>>) invocation -> {
        String blobId = invocation.getArgumentAt(0, String.class);
        allDeleted.add(blobId);
        return CompletableFuture.completedFuture(null);
    });
    // Execute
    CompletionStage<Void> cleanUpFuture = blobStoreUtil.cleanUpDir(dirIndex, metadata);
    try {
        // should be already complete. if not, future composition in putDir is broken.
        cleanUpFuture.toCompletableFuture().get(0, TimeUnit.MILLISECONDS);
    } catch (TimeoutException e) {
        fail("Future returned from putDir should be already complete.");
    }
    // Assert
    assertEquals(expectedRemovedFiles, allDeleted);
}

Example 2 with BlobStoreManager

use of org.apache.samza.storage.blobstore.BlobStoreManager in project samza by apache.

the class TestBlobStoreUtil method testPutDirFailsIfAnyFileUploadFails.

@Test
public void testPutDirFailsIfAnyFileUploadFails() throws IOException, TimeoutException, InterruptedException {
    BlobStoreManager blobStoreManager = mock(BlobStoreManager.class);
    // File, dir and recursive dir added, retained and removed in local
    String local = "[a, b]";
    String remote = "[]";
    // Set up environment
    Path localSnapshotDir = BlobStoreTestUtil.createLocalDir(local);
    String basePath = localSnapshotDir.toAbsolutePath().toString();
    DirIndex remoteSnapshotDir = BlobStoreTestUtil.createDirIndex(remote);
    SnapshotMetadata snapshotMetadata = new SnapshotMetadata(checkpointId, jobName, jobId, taskName, storeName);
    DirDiff dirDiff = DirDiffUtil.getDirDiff(localSnapshotDir.toFile(), remoteSnapshotDir, (localFile, remoteFile) -> localFile.getName().equals(remoteFile.getFileName()));
    // Set up mocks
    SamzaException exception = new SamzaException("Error uploading file");
    CompletableFuture<String> failedFuture = new CompletableFuture<>();
    failedFuture.completeExceptionally(exception);
    when(blobStoreManager.put(any(InputStream.class), any(Metadata.class))).thenAnswer((Answer<CompletableFuture<String>>) invocation -> {
        Metadata metadata = invocation.getArgumentAt(1, Metadata.class);
        String path = metadata.getPayloadPath();
        if (path.endsWith("a")) {
            return CompletableFuture.completedFuture("aBlobId");
        } else {
            return failedFuture;
        }
    });
    // Execute
    BlobStoreUtil blobStoreUtil = new BlobStoreUtil(blobStoreManager, EXECUTOR, null, null);
    CompletionStage<DirIndex> dirIndexFuture = blobStoreUtil.putDir(dirDiff, snapshotMetadata);
    try {
        // should be already complete. if not, future composition in putDir is broken.
        dirIndexFuture.toCompletableFuture().get(0, TimeUnit.MILLISECONDS);
    } catch (ExecutionException e) {
        Throwable cause = e.getCause();
        // Assert that the result future fails and that the cause is propagated correctly
        assertEquals(exception, cause);
        return;
    }
    fail("DirIndex future should have been completed with an exception");
}

Example 3 with BlobStoreManager

use of org.apache.samza.storage.blobstore.BlobStoreManager in project samza by apache.

the class TestBlobStoreUtil method testCleanUpFailsIfAnyFileDeleteFails.

@Test
public void testCleanUpFailsIfAnyFileDeleteFails() throws IOException, TimeoutException, InterruptedException, ExecutionException {
    BlobStoreManager blobStoreManager = mock(BlobStoreManager.class);
    // File, dir and recursive dir added, retained and removed in local
    // Using unique file names since test util uses only the file name (leaf node)
    // as the mock blob id, not the full file path.
    String local = "[a, b]";
    String remote = "[c, d]";
    // Set up environment
    Path localSnapshotDir = BlobStoreTestUtil.createLocalDir(local);
    String basePath = localSnapshotDir.toAbsolutePath().toString();
    DirIndex remoteSnapshotDir = BlobStoreTestUtil.createDirIndex(remote);
    SnapshotMetadata snapshotMetadata = new SnapshotMetadata(checkpointId, jobName, jobId, taskName, storeName);
    DirDiff dirDiff = DirDiffUtil.getDirDiff(localSnapshotDir.toFile(), remoteSnapshotDir, (localFile, remoteFile) -> localFile.getName().equals(remoteFile.getFileName()));
    BlobStoreUtil blobStoreUtil = new BlobStoreUtil(blobStoreManager, EXECUTOR, null, null);
    when(blobStoreManager.put(any(InputStream.class), any(Metadata.class))).thenReturn(CompletableFuture.completedFuture("blobId"));
    CompletionStage<DirIndex> dirIndexFuture = blobStoreUtil.putDir(dirDiff, snapshotMetadata);
    DirIndex dirIndex = null;
    try {
        // should be already complete. if not, future composition in putDir is broken.
        dirIndex = dirIndexFuture.toCompletableFuture().get(0, TimeUnit.MILLISECONDS);
    } catch (TimeoutException e) {
        fail("Future returned from putDir should be already complete.");
    }
    // Set up mocks
    SamzaException exception = new SamzaException("Error deleting file");
    CompletableFuture<Void> failedFuture = new CompletableFuture<>();
    failedFuture.completeExceptionally(exception);
    when(blobStoreManager.delete(anyString(), any(Metadata.class))).thenAnswer((Answer<CompletableFuture<Void>>) invocation -> {
        String blobId = invocation.getArgumentAt(0, String.class);
        if (blobId.equals("c")) {
            return CompletableFuture.completedFuture(null);
        } else {
            return failedFuture;
        }
    });
    // Execute
    CompletionStage<Void> cleanUpFuture = blobStoreUtil.cleanUpDir(dirIndex, metadata);
    try {
        // should be already complete. if not, future composition in putDir is broken.
        cleanUpFuture.toCompletableFuture().get(0, TimeUnit.MILLISECONDS);
    } catch (ExecutionException e) {
        Throwable cause = e.getCause();
        // Assert that the result future fails and that the cause is propagated correctly
        assertEquals(exception, cause);
        return;
    }
    fail("Clean up future should have been completed with an exception");
}

Example 4 with BlobStoreManager

use of org.apache.samza.storage.blobstore.BlobStoreManager in project samza by apache.

the class TestBlobStoreUtil method testRestoreDirRestoresMultiPartFilesCorrectly.

@Test
public void testRestoreDirRestoresMultiPartFilesCorrectly() throws IOException {
    Path restoreDirBasePath = Files.createTempDirectory(BlobStoreTestUtil.TEMP_DIR_PREFIX);
    // remote file == 26 blobs, blob ids from a to z, blob contents from a to z, offsets 0 to 25.
    DirIndex mockDirIndex = mock(DirIndex.class);
    when(mockDirIndex.getDirName()).thenReturn(DirIndex.ROOT_DIR_NAME);
    FileIndex mockFileIndex = mock(FileIndex.class);
    when(mockFileIndex.getFileName()).thenReturn("1.sst");
    // setup mock file attributes. create a temp file to get current user/group/permissions so that they
    // match with restored files.
    File tmpFile = Paths.get(restoreDirBasePath.toString(), "tempfile-" + new Random().nextInt()).toFile();
    tmpFile.createNewFile();
    PosixFileAttributes attrs = Files.readAttributes(tmpFile.toPath(), PosixFileAttributes.class);
    FileMetadata fileMetadata = new // ctime mtime does not matter. size == 26
    FileMetadata(// ctime mtime does not matter. size == 26
    1234L, // ctime mtime does not matter. size == 26
    1243L, // ctime mtime does not matter. size == 26
    26, attrs.owner().getName(), attrs.group().getName(), PosixFilePermissions.toString(attrs.permissions()));
    when(mockFileIndex.getFileMetadata()).thenReturn(fileMetadata);
    // delete so that it doesn't show up in restored dir contents.
    Files.delete(tmpFile.toPath());
    List<FileBlob> mockFileBlobs = new ArrayList<>();
    StringBuilder fileContents = new StringBuilder();
    for (int i = 0; i < 26; i++) {
        FileBlob mockFileBlob = mock(FileBlob.class);
        char c = (char) ('a' + i);
        // blob contents == blobId
        fileContents.append(c);
        when(mockFileBlob.getBlobId()).thenReturn(String.valueOf(c));
        when(mockFileBlob.getOffset()).thenReturn(i);
        mockFileBlobs.add(mockFileBlob);
    }
    when(mockFileIndex.getBlobs()).thenReturn(mockFileBlobs);
    CRC32 checksum = new CRC32();
    checksum.update(fileContents.toString().getBytes());
    when(mockFileIndex.getChecksum()).thenReturn(checksum.getValue());
    when(mockDirIndex.getFilesPresent()).thenReturn(ImmutableList.of(mockFileIndex));
    BlobStoreManager mockBlobStoreManager = mock(BlobStoreManager.class);
    when(mockBlobStoreManager.get(anyString(), any(OutputStream.class), any(Metadata.class))).thenAnswer((Answer<CompletionStage<Void>>) invocationOnMock -> {
        String blobId = invocationOnMock.getArgumentAt(0, String.class);
        OutputStream outputStream = invocationOnMock.getArgumentAt(1, OutputStream.class);
        outputStream.write(blobId.getBytes());
        ((FileOutputStream) outputStream).getFD().sync();
        return CompletableFuture.completedFuture(null);
    });
    BlobStoreUtil blobStoreUtil = new BlobStoreUtil(mockBlobStoreManager, EXECUTOR, null, null);
    blobStoreUtil.restoreDir(restoreDirBasePath.toFile(), mockDirIndex, metadata).join();
    assertTrue(new DirDiffUtil().areSameDir(Collections.emptySet(), false).test(restoreDirBasePath.toFile(), mockDirIndex));
}

Example 5 with BlobStoreManager

use of org.apache.samza.storage.blobstore.BlobStoreManager in project samza by apache.

the class TestBlobStoreUtil method testPutFileChecksumAndMetadata.

@Test
public void testPutFileChecksumAndMetadata() throws IOException, ExecutionException, InterruptedException {
    // Setup
    SnapshotMetadata snapshotMetadata = new SnapshotMetadata(checkpointId, jobName, jobId, taskName, storeName);
    Path path = Files.createTempFile("samza-testPutFileChecksum-", ".tmp");
    FileUtil fileUtil = new FileUtil();
    fileUtil.writeToTextFile(path.toFile(), RandomStringUtils.random(1000), false);
    long expectedChecksum = FileUtils.checksumCRC32(path.toFile());
    BlobStoreManager blobStoreManager = mock(BlobStoreManager.class);
    ArgumentCaptor<Metadata> argumentCaptor = ArgumentCaptor.forClass(Metadata.class);
    when(blobStoreManager.put(any(InputStream.class), argumentCaptor.capture())).thenAnswer((Answer<CompletionStage<String>>) invocation -> {
        InputStream inputStream = invocation.getArgumentAt(0, InputStream.class);
        IOUtils.copy(inputStream, NullOutputStream.NULL_OUTPUT_STREAM);
        return CompletableFuture.completedFuture("blobId");
    });
    BlobStoreUtil blobStoreUtil = new BlobStoreUtil(blobStoreManager, EXECUTOR, null, null);
    CompletionStage<FileIndex> fileIndexFuture = blobStoreUtil.putFile(path.toFile(), snapshotMetadata);
    FileIndex fileIndex = null;
    try {
        // should be already complete. if not, future composition in putFile is broken.
        fileIndex = fileIndexFuture.toCompletableFuture().get(0, TimeUnit.MILLISECONDS);
    } catch (TimeoutException e) {
        fail("Future returned from putFile should be already complete.");
    }
    // Assert
    Metadata metadata = (Metadata) argumentCaptor.getValue();
    assertEquals(path.toAbsolutePath().toString(), metadata.getPayloadPath());
    assertEquals(path.toFile().length(), Long.valueOf(metadata.getPayloadSize()).longValue());
    assertEquals(expectedChecksum, fileIndex.getChecksum());
}