TransactionalStateTaskRestoreManager implements TaskRestoreManager { @VisibleForTesting static void setupStoreDirs( TaskModel taskModel, Map<String, StorageEngine> storeEngines, StoreActions storeActions, StorageManagerUtil storageManagerUtil, FileUtil fileUtil, File loggedStoreBaseDirectory, File nonLoggedStoreBaseDirectory) { TaskName taskName = taskModel.getTaskName(); TaskMode taskMode = taskModel.getTaskMode(); ListMultimap<String, File> storeDirsToDelete = storeActions.storeDirsToDelete; Map<String, File> storeDirsToRetain = storeActions.storeDirsToRetain; storeDirsToDelete.entries().forEach(entry -> { String storeName = entry.getKey(); File storeDirToDelete = entry.getValue(); LOG.info("Deleting persistent store directory: {} for store: {} in task: {}", storeDirToDelete, storeName, taskName); fileUtil.rm(storeDirToDelete); }); storeDirsToRetain.forEach((storeName, storeDirToRetain) -> { File currentDir = storageManagerUtil.getTaskStoreDir( loggedStoreBaseDirectory, storeName, taskName, taskMode); LOG.info("Moving logged store checkpoint directory: {} for store: {} in task: {} to current directory: {}", storeDirsToRetain.toString(), storeName, taskName, currentDir); storageManagerUtil.restoreCheckpointFiles(storeDirToRetain, currentDir); }); storeEngines.forEach((storeName, storageEngine) -> { if (storageEngine.getStoreProperties().isPersistedToDisk()) { File currentDir; if (storageEngine.getStoreProperties().isLoggedStore()) { currentDir = storageManagerUtil.getTaskStoreDir( loggedStoreBaseDirectory, storeName, taskName, taskMode); } else { currentDir = storageManagerUtil.getTaskStoreDir( nonLoggedStoreBaseDirectory, storeName, taskName, taskMode); } try { if (!fileUtil.exists(currentDir.toPath())) { LOG.info("Creating missing persistent store current directory: {} for store: {} in task: {}", currentDir, storeName, taskName); fileUtil.createDirectories(currentDir.toPath()); } } catch (Exception e) { throw new SamzaException(String.format("Error setting up current directory for store: %s", storeName), e); } } }); } TransactionalStateTaskRestoreManager(
      TaskModel taskModel,
      Map<String, StorageEngine> storeEngines,
      Map<String, SystemStream> storeChangelogs,
      SystemAdmins systemAdmins,
      Map<String, SystemConsumer> storeConsumers,
      SSPMetadataCache sspMetadataCache,
      File loggedStoreBaseDirectory,
      File nonLoggedStoreBaseDirectory,
      Config config,
      Clock clock); @Override void init(Map<SystemStreamPartition, String> checkpointedChangelogOffsets); @Override void restore(); void stopPersistentStores();  }
@Test public void testSetupStoreDirs() { TaskModel mockTaskModel = mock(TaskModel.class); TaskName taskName = new TaskName("Partition 0"); when(mockTaskModel.getTaskName()).thenReturn(taskName); Partition taskChangelogPartition = new Partition(0); when(mockTaskModel.getChangelogPartition()).thenReturn(taskChangelogPartition); String store1Name = "store1"; StorageEngine store1Engine = mock(StorageEngine.class); StoreProperties mockStore1Properties = mock(StoreProperties.class); when(store1Engine.getStoreProperties()).thenReturn(mockStore1Properties); when(mockStore1Properties.isLoggedStore()).thenReturn(true); when(mockStore1Properties.isPersistedToDisk()).thenReturn(true); String store2Name = "store2"; StorageEngine store2Engine = mock(StorageEngine.class); StoreProperties mockStore2Properties = mock(StoreProperties.class); when(store2Engine.getStoreProperties()).thenReturn(mockStore2Properties); when(mockStore2Properties.isLoggedStore()).thenReturn(false); when(mockStore2Properties.isPersistedToDisk()).thenReturn(true); Map<String, StorageEngine> mockStoreEngines = ImmutableMap.of( store1Name, store1Engine, store2Name, store2Engine); File mockStore1DirToRetain = mock(File.class); ImmutableMap<String, File> storeDirsToRetain = ImmutableMap.of(store1Name, mockStore1DirToRetain); ListMultimap<String, File> storeDirsToDelete = ArrayListMultimap.create(); File mockStore1CurrentDir = mock(File.class); Path mockStore1CurrentDirPath = mock(Path.class); when(mockStore1CurrentDir.toPath()).thenReturn(mockStore1CurrentDirPath); File mockStore2CurrentDir = mock(File.class); Path mockStore2CurrentDirPath = mock(Path.class); when(mockStore2CurrentDir.toPath()).thenReturn(mockStore2CurrentDirPath); File mockStore1DirToDelete1 = mock(File.class); File mockStore1DirToDelete2 = mock(File.class); File mockStore2DirToDelete1 = mock(File.class); File mockStore2DirToDelete2 = mock(File.class); storeDirsToDelete.put(store1Name, mockStore1CurrentDir); storeDirsToDelete.put(store1Name, mockStore1DirToDelete1); storeDirsToDelete.put(store1Name, mockStore1DirToDelete2); storeDirsToDelete.put(store2Name, mockStore2CurrentDir); storeDirsToDelete.put(store2Name, mockStore2DirToDelete1); storeDirsToDelete.put(store2Name, mockStore2DirToDelete2); StoreActions storeActions = new StoreActions(storeDirsToRetain, storeDirsToDelete, ImmutableMap.of()); StorageManagerUtil mockStorageManagerUtil = mock(StorageManagerUtil.class); FileUtil mockFileUtil = mock(FileUtil.class); File mockLoggedStoreBaseDir = mock(File.class); File mockNonLoggedStoreBaseDir = mock(File.class); when(mockStorageManagerUtil.getTaskStoreDir(eq(mockLoggedStoreBaseDir), eq(store1Name), any(), any())) .thenReturn(mockStore1CurrentDir); when(mockStorageManagerUtil.getTaskStoreDir(eq(mockNonLoggedStoreBaseDir), eq(store2Name), any(), any())) .thenReturn(mockStore2CurrentDir); when(mockFileUtil.exists(eq(mockStore1CurrentDirPath))) .thenReturn(false); when(mockFileUtil.exists(eq(mockStore2CurrentDirPath))) .thenReturn(true); TransactionalStateTaskRestoreManager.setupStoreDirs(mockTaskModel, mockStoreEngines, storeActions, mockStorageManagerUtil, mockFileUtil, mockLoggedStoreBaseDir, mockNonLoggedStoreBaseDir); verify(mockFileUtil, times(1)).rm(mockStore1CurrentDir); verify(mockFileUtil, times(1)).rm(mockStore1DirToDelete1); verify(mockFileUtil, times(1)).rm(mockStore1DirToDelete2); verify(mockFileUtil, times(1)).rm(mockStore2CurrentDir); verify(mockFileUtil, times(1)).rm(mockStore2DirToDelete1); verify(mockFileUtil, times(1)).rm(mockStore2DirToDelete2); verify(mockStorageManagerUtil, times(1)).restoreCheckpointFiles(any(), any()); verify(mockFileUtil, times(1)).exists(mockStore1CurrentDirPath); verify(mockFileUtil, times(1)).createDirectories(mockStore1CurrentDirPath); verify(mockFileUtil, times(1)).exists(mockStore2CurrentDirPath); verify(mockFileUtil, never()).createDirectories(mockStore2CurrentDirPath); verifyNoMoreInteractions(mockFileUtil); }