| /* |
| * Licensed to the Apache Software Foundation (ASF) under one or more |
| * contributor license agreements. See the NOTICE file distributed with |
| * this work for additional information regarding copyright ownership. |
| * The ASF licenses this file to You under the Apache License, Version 2.0 |
| * (the "License"); you may not use this file except in compliance with |
| * the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| package org.apache.nifi.controller.repository; |
| |
| import org.apache.commons.lang3.SystemUtils; |
| import org.apache.nifi.connectable.Connectable; |
| import org.apache.nifi.connectable.Connection; |
| import org.apache.nifi.controller.queue.FlowFileQueue; |
| import org.apache.nifi.controller.queue.NopConnectionEventListener; |
| import org.apache.nifi.controller.queue.QueueSize; |
| import org.apache.nifi.controller.queue.StandardFlowFileQueue; |
| import org.apache.nifi.controller.repository.claim.ContentClaim; |
| import org.apache.nifi.controller.repository.claim.ResourceClaim; |
| import org.apache.nifi.controller.repository.claim.ResourceClaimManager; |
| import org.apache.nifi.controller.repository.claim.StandardContentClaim; |
| import org.apache.nifi.controller.repository.claim.StandardResourceClaimManager; |
| import org.apache.nifi.controller.swap.StandardSwapContents; |
| import org.apache.nifi.controller.swap.StandardSwapSummary; |
| import org.apache.nifi.rocksdb.RocksDBMetronome; |
| import org.apache.nifi.util.NiFiProperties; |
| import org.apache.nifi.util.file.FileUtils; |
| import org.junit.Assume; |
| import org.junit.Before; |
| import org.junit.BeforeClass; |
| import org.junit.Rule; |
| import org.junit.Test; |
| import org.junit.rules.TemporaryFolder; |
| import org.junit.rules.TestName; |
| import org.mockito.Mockito; |
| import org.mockito.stubbing.Answer; |
| import org.slf4j.Logger; |
| import org.slf4j.LoggerFactory; |
| |
| import java.io.File; |
| import java.io.IOException; |
| import java.nio.charset.StandardCharsets; |
| import java.nio.file.Files; |
| import java.nio.file.Path; |
| import java.nio.file.Paths; |
| import java.nio.file.StandardCopyOption; |
| import java.util.ArrayList; |
| import java.util.Collection; |
| import java.util.Collections; |
| import java.util.HashMap; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.Set; |
| import java.util.UUID; |
| import java.util.concurrent.ConcurrentSkipListSet; |
| import java.util.function.Function; |
| import java.util.stream.Collectors; |
| |
| import static org.junit.Assert.assertEquals; |
| import static org.junit.Assert.assertFalse; |
| import static org.junit.Assert.assertNotNull; |
| import static org.junit.Assert.assertNull; |
| import static org.junit.Assert.assertTrue; |
| import static org.junit.Assert.fail; |
| import static org.mockito.ArgumentMatchers.any; |
| import static org.mockito.Mockito.doAnswer; |
| import static org.mockito.Mockito.when; |
| |
| public class TestRocksDBFlowFileRepository { |
| |
| private static final Logger logger = LoggerFactory.getLogger(TestRocksDBFlowFileRepository.class); |
| |
| private final Map<String, String> additionalProperties = new HashMap<>(); |
| private String nifiPropertiesPath; |
| |
| @Rule |
| public TestName testName = new TestName(); |
| |
| @Rule |
| public TemporaryFolder temporaryFolder = new TemporaryFolder(); |
| |
| @BeforeClass |
| public static void setupClass() { |
| Assume.assumeTrue("Test only runs on *nix", !SystemUtils.IS_OS_WINDOWS); |
| } |
| |
| @Before |
| public void before() throws IOException { |
| File testRepoDir = temporaryFolder.newFolder(testName.getMethodName()); |
| additionalProperties.put(NiFiProperties.FLOWFILE_REPOSITORY_DIRECTORY, testRepoDir.getAbsolutePath()); |
| |
| File properties = temporaryFolder.newFile(); |
| Files.copy(Paths.get("src/test/resources/conf/nifi.properties"), properties.toPath(), StandardCopyOption.REPLACE_EXISTING); |
| nifiPropertiesPath = properties.getAbsolutePath(); |
| |
| logger.info("Running test: {}", testName.getMethodName()); |
| } |
| |
| @Test |
| public void testNormalizeSwapLocation() { |
| assertEquals("/", RocksDBFlowFileRepository.normalizeSwapLocation("/")); |
| assertEquals("", RocksDBFlowFileRepository.normalizeSwapLocation("")); |
| assertNull(RocksDBFlowFileRepository.normalizeSwapLocation(null)); |
| assertEquals("test", RocksDBFlowFileRepository.normalizeSwapLocation("test.txt")); |
| assertEquals("test", RocksDBFlowFileRepository.normalizeSwapLocation("/test.txt")); |
| assertEquals("test", RocksDBFlowFileRepository.normalizeSwapLocation("/tmp/test.txt")); |
| assertEquals("test", RocksDBFlowFileRepository.normalizeSwapLocation("//test.txt")); |
| assertEquals("test", RocksDBFlowFileRepository.normalizeSwapLocation("/path/to/other/file/repository/test.txt")); |
| assertEquals("test", RocksDBFlowFileRepository.normalizeSwapLocation("test.txt/")); |
| assertEquals("test", RocksDBFlowFileRepository.normalizeSwapLocation("/path/to/test.txt/")); |
| assertEquals("test", WriteAheadFlowFileRepository.normalizeSwapLocation(WriteAheadFlowFileRepository.normalizeSwapLocation("/path/to/test.txt/"))); |
| } |
| |
| @Test |
| public void testSwapLocationsRestored() throws IOException { |
| final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties)); |
| repo.initialize(new StandardResourceClaimManager()); |
| |
| final TestRocksDBFlowFileRepository.TestQueueProvider queueProvider = new TestRocksDBFlowFileRepository.TestQueueProvider(); |
| repo.loadFlowFiles(queueProvider); |
| |
| final Connection connection = Mockito.mock(Connection.class); |
| when(connection.getIdentifier()).thenReturn("1234"); |
| |
| final FlowFileQueue queue = Mockito.mock(FlowFileQueue.class); |
| when(queue.getIdentifier()).thenReturn("1234"); |
| when(connection.getFlowFileQueue()).thenReturn(queue); |
| |
| queueProvider.addConnection(connection); |
| |
| StandardFlowFileRecord.Builder ffBuilder = new StandardFlowFileRecord.Builder(); |
| ffBuilder.id(1L); |
| ffBuilder.size(0L); |
| final FlowFileRecord flowFileRecord = ffBuilder.build(); |
| |
| final List<RepositoryRecord> records = new ArrayList<>(); |
| final StandardRepositoryRecord record = new StandardRepositoryRecord(queue, flowFileRecord, "swap123"); |
| record.setDestination(queue); |
| records.add(record); |
| |
| repo.updateRepository(records); |
| repo.close(); |
| |
| // restore |
| final RocksDBFlowFileRepository repo2 = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties)); |
| repo2.initialize(new StandardResourceClaimManager()); |
| repo2.loadFlowFiles(queueProvider); |
| assertTrue(repo2.isValidSwapLocationSuffix("swap123")); |
| assertFalse(repo2.isValidSwapLocationSuffix("other")); |
| repo2.close(); |
| } |
| |
| @Test |
| public void testSwapLocationsUpdatedOnRepoUpdate() throws IOException { |
| final Path path = Paths.get("target/test-swap-repo"); |
| if (Files.exists(path)) { |
| FileUtils.deleteFile(path.toFile(), true); |
| } |
| |
| final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties)); |
| repo.initialize(new StandardResourceClaimManager()); |
| |
| final TestRocksDBFlowFileRepository.TestQueueProvider queueProvider = new TestRocksDBFlowFileRepository.TestQueueProvider(); |
| repo.loadFlowFiles(queueProvider); |
| |
| final Connection connection = Mockito.mock(Connection.class); |
| when(connection.getIdentifier()).thenReturn("1234"); |
| |
| final FlowFileQueue queue = Mockito.mock(FlowFileQueue.class); |
| when(queue.getIdentifier()).thenReturn("1234"); |
| when(connection.getFlowFileQueue()).thenReturn(queue); |
| |
| queueProvider.addConnection(connection); |
| |
| StandardFlowFileRecord.Builder ffBuilder = new StandardFlowFileRecord.Builder(); |
| ffBuilder.id(1L); |
| ffBuilder.size(0L); |
| final FlowFileRecord flowFileRecord = ffBuilder.build(); |
| |
| final List<RepositoryRecord> records = new ArrayList<>(); |
| final StandardRepositoryRecord record = new StandardRepositoryRecord(queue, flowFileRecord, "/tmp/swap123"); |
| record.setDestination(queue); |
| records.add(record); |
| |
| assertFalse(repo.isValidSwapLocationSuffix("swap123")); |
| repo.updateRepository(records); |
| assertTrue(repo.isValidSwapLocationSuffix("swap123")); |
| repo.close(); |
| } |
| |
| @Test |
| public void testResourceClaimsIncremented() throws IOException { |
| final ResourceClaimManager claimManager = new StandardResourceClaimManager(); |
| |
| final TestRocksDBFlowFileRepository.TestQueueProvider queueProvider = new TestRocksDBFlowFileRepository.TestQueueProvider(); |
| final Connection connection = Mockito.mock(Connection.class); |
| when(connection.getIdentifier()).thenReturn("1234"); |
| when(connection.getDestination()).thenReturn(Mockito.mock(Connectable.class)); |
| |
| final FlowFileSwapManager swapMgr = new TestRocksDBFlowFileRepository.MockFlowFileSwapManager(); |
| final FlowFileQueue queue = new StandardFlowFileQueue("1234", new NopConnectionEventListener(), null, null, claimManager, null, swapMgr, null, 10000, "0 sec", 0L, "0 B"); |
| |
| when(connection.getFlowFileQueue()).thenReturn(queue); |
| queueProvider.addConnection(connection); |
| |
| final ResourceClaim resourceClaim1 = claimManager.newResourceClaim("container", "section", "1", false, false); |
| final ContentClaim claim1 = new StandardContentClaim(resourceClaim1, 0L); |
| |
| final ResourceClaim resourceClaim2 = claimManager.newResourceClaim("container", "section", "2", false, false); |
| final ContentClaim claim2 = new StandardContentClaim(resourceClaim2, 0L); |
| |
| // Create a flowfile repo, update it once with a FlowFile that points to one resource claim. Then, |
| // indicate that a FlowFile was swapped out. We should then be able to recover these FlowFiles and the |
| // resource claims' counts should be updated for both the swapped out FlowFile and the non-swapped out FlowFile |
| try (final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties))) { |
| repo.initialize(claimManager); |
| repo.loadFlowFiles(queueProvider); |
| |
| // Create a Repository Record that indicates that a FlowFile was created |
| final FlowFileRecord flowFile1 = new StandardFlowFileRecord.Builder() |
| .id(1L) |
| .addAttribute("uuid", "11111111-1111-1111-1111-111111111111") |
| .contentClaim(claim1) |
| .build(); |
| final StandardRepositoryRecord rec1 = new StandardRepositoryRecord(queue); |
| rec1.setWorking(flowFile1, false); |
| rec1.setDestination(queue); |
| |
| // Create a Record that we can swap out |
| final FlowFileRecord flowFile2 = new StandardFlowFileRecord.Builder() |
| .id(2L) |
| .addAttribute("uuid", "11111111-1111-1111-1111-111111111112") |
| .contentClaim(claim2) |
| .build(); |
| |
| final StandardRepositoryRecord rec2 = new StandardRepositoryRecord(queue); |
| rec2.setWorking(flowFile2, false); |
| rec2.setDestination(queue); |
| |
| final List<RepositoryRecord> records = new ArrayList<>(); |
| records.add(rec1); |
| records.add(rec2); |
| repo.updateRepository(records); |
| |
| final String swapLocation = swapMgr.swapOut(Collections.singletonList(flowFile2), queue, null); |
| repo.swapFlowFilesOut(Collections.singletonList(flowFile2), queue, swapLocation); |
| } |
| |
| final ResourceClaimManager recoveryClaimManager = new StandardResourceClaimManager(); |
| try (final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties))) { |
| repo.initialize(recoveryClaimManager); |
| final long largestId = repo.loadFlowFiles(queueProvider); |
| |
| // largest ID known is 1 because this doesn't take into account the FlowFiles that have been swapped out |
| assertEquals(1, largestId); |
| } |
| |
| // resource claim 1 will have a single claimant count while resource claim 2 will have no claimant counts |
| // because resource claim 2 is referenced only by flowfiles that are swapped out. |
| assertEquals(1, recoveryClaimManager.getClaimantCount(resourceClaim1)); |
| assertEquals(0, recoveryClaimManager.getClaimantCount(resourceClaim2)); |
| |
| final SwapSummary summary = queue.recoverSwappedFlowFiles(); |
| assertNotNull(summary); |
| assertEquals(2, summary.getMaxFlowFileId().intValue()); |
| assertEquals(new QueueSize(1, 0L), summary.getQueueSize()); |
| |
| final List<ResourceClaim> swappedOutClaims = summary.getResourceClaims(); |
| assertNotNull(swappedOutClaims); |
| assertEquals(1, swappedOutClaims.size()); |
| assertEquals(claim2.getResourceClaim(), swappedOutClaims.get(0)); |
| } |
| |
| @Test |
| public void testRestartWithOneRecord() throws IOException { |
| final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties)); |
| repo.initialize(new StandardResourceClaimManager()); |
| |
| final TestRocksDBFlowFileRepository.TestQueueProvider queueProvider = new TestRocksDBFlowFileRepository.TestQueueProvider(); |
| repo.loadFlowFiles(queueProvider); |
| |
| final List<FlowFileRecord> flowFileCollection = new ArrayList<>(); |
| |
| final Connection connection = Mockito.mock(Connection.class); |
| when(connection.getIdentifier()).thenReturn("1234"); |
| |
| final FlowFileQueue queue = Mockito.mock(FlowFileQueue.class); |
| when(queue.getIdentifier()).thenReturn("1234"); |
| doAnswer((Answer<Object>) invocation -> { |
| flowFileCollection.add((FlowFileRecord) invocation.getArguments()[0]); |
| return null; |
| }).when(queue).put(any(FlowFileRecord.class)); |
| |
| when(connection.getFlowFileQueue()).thenReturn(queue); |
| |
| queueProvider.addConnection(connection); |
| |
| StandardFlowFileRecord.Builder ffBuilder = new StandardFlowFileRecord.Builder(); |
| ffBuilder.id(1L); |
| ffBuilder.addAttribute("abc", "xyz"); |
| ffBuilder.size(0L); |
| final FlowFileRecord flowFileRecord = ffBuilder.build(); |
| |
| final List<RepositoryRecord> records = new ArrayList<>(); |
| final StandardRepositoryRecord record = new StandardRepositoryRecord(null); |
| record.setWorking(flowFileRecord, false); |
| record.setDestination(connection.getFlowFileQueue()); |
| records.add(record); |
| |
| repo.updateRepository(records); |
| |
| // update to add new attribute |
| ffBuilder = new StandardFlowFileRecord.Builder().fromFlowFile(flowFileRecord).addAttribute("hello", "world"); |
| final FlowFileRecord flowFileRecord2 = ffBuilder.build(); |
| record.setWorking(flowFileRecord2, false); |
| repo.updateRepository(records); |
| |
| // update size but no attribute |
| ffBuilder = new StandardFlowFileRecord.Builder().fromFlowFile(flowFileRecord2).size(40L); |
| final FlowFileRecord flowFileRecord3 = ffBuilder.build(); |
| record.setWorking(flowFileRecord3, false); |
| repo.updateRepository(records); |
| |
| repo.close(); |
| |
| // restore |
| final RocksDBFlowFileRepository repo2 = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties)); |
| repo2.initialize(new StandardResourceClaimManager()); |
| repo2.loadFlowFiles(queueProvider); |
| |
| assertEquals(1, flowFileCollection.size()); |
| final FlowFileRecord flowFile = flowFileCollection.get(0); |
| assertEquals(1L, flowFile.getId()); |
| assertEquals("xyz", flowFile.getAttribute("abc")); |
| assertEquals(40L, flowFile.getSize()); |
| assertEquals("world", flowFile.getAttribute("hello")); |
| |
| repo2.close(); |
| } |
| |
| |
| @Test |
| public void testDoNotRemoveOrphans() throws Exception { |
| final TestQueue testQueue = new TestQueue(); |
| |
| try (final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties))) { |
| |
| repo.initialize(new StandardResourceClaimManager()); |
| repo.loadFlowFiles(testQueue.provider); |
| |
| repo.updateRepository(testQueue.getRepositoryRecord(new StandardFlowFileRecord.Builder() |
| .id(1L) |
| .addAttribute("abc", "xyz") |
| .size(0L) |
| .build() |
| )); |
| } |
| |
| // restore (& confirm present) |
| try (final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties))) { |
| repo.initialize(new StandardResourceClaimManager()); |
| repo.loadFlowFiles(testQueue.provider); |
| assertEquals(1, repo.getInMemoryFlowFiles()); |
| } |
| |
| // restore with empty queue provider (should throw exception) |
| try (final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties))) { |
| repo.initialize(new StandardResourceClaimManager()); |
| repo.loadFlowFiles(new TestQueueProvider()); |
| fail(); |
| } catch (IOException expected) { |
| assertTrue(expected.getMessage().contains("Found FlowFile in repository without a corresponding queue")); |
| } |
| } |
| |
| @Test |
| public void testRemoveOrphans() throws Exception { |
| final TestQueue testQueue = new TestQueue(); |
| |
| additionalProperties.put(RocksDBFlowFileRepository.RocksDbProperty.REMOVE_ORPHANED_FLOWFILES.propertyName, "true"); |
| |
| try (final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties))) { |
| |
| repo.initialize(new StandardResourceClaimManager()); |
| repo.loadFlowFiles(testQueue.provider); |
| |
| repo.updateRepository(testQueue.getRepositoryRecord(new StandardFlowFileRecord.Builder() |
| .id(1L) |
| .addAttribute("abc", "xyz") |
| .size(0L) |
| .build() |
| )); |
| } |
| |
| // restore (& confirm present) |
| try (final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties))) { |
| repo.initialize(new StandardResourceClaimManager()); |
| repo.loadFlowFiles(testQueue.provider); |
| assertEquals(1, repo.getInMemoryFlowFiles()); |
| } |
| // restore with empty queue provider (should throw exception) |
| try (final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties))) { |
| repo.initialize(new StandardResourceClaimManager()); |
| repo.loadFlowFiles(new TestQueueProvider()); |
| assertEquals(0, repo.getInMemoryFlowFiles()); |
| } |
| } |
| |
| @Test |
| public void testKnownVersion() throws Exception { |
| final NiFiProperties niFiProperties = NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties); |
| |
| // create db with known version |
| try (RocksDBMetronome db = new RocksDBMetronome.Builder().setStoragePath(RocksDBFlowFileRepository.getFlowFileRepoPath(niFiProperties)).build()) { |
| db.initialize(); |
| db.putConfiguration(RocksDBFlowFileRepository.REPOSITORY_VERSION_KEY, RocksDBFlowFileRepository.VERSION_ONE_BYTES); |
| } |
| try (final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties))) { |
| repo.initialize(new StandardResourceClaimManager()); |
| } |
| } |
| |
| @Test(expected = IllegalStateException.class) |
| public void testUnknownVersion() throws Exception { |
| final NiFiProperties niFiProperties = NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties); |
| |
| // create db with known version |
| try (RocksDBMetronome db = new RocksDBMetronome.Builder().setStoragePath(RocksDBFlowFileRepository.getFlowFileRepoPath(niFiProperties)).build()) { |
| db.initialize(); |
| db.putConfiguration(RocksDBFlowFileRepository.REPOSITORY_VERSION_KEY, "UNKNOWN".getBytes(StandardCharsets.UTF_8)); |
| } |
| try (final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties))) { |
| repo.initialize(new StandardResourceClaimManager()); |
| } |
| } |
| |
| @Test |
| public void testRecoveryMode() throws Exception { |
| int totalFlowFiles = 50; |
| |
| final TestQueue testQueue = new TestQueue(); |
| |
| try (final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties))) { |
| |
| repo.initialize(new StandardResourceClaimManager()); |
| repo.loadFlowFiles(testQueue.provider); |
| |
| // add records to the repo |
| for (int i = 1; i <= totalFlowFiles; i++) { |
| repo.updateRepository(testQueue.getRepositoryRecord(new StandardFlowFileRecord.Builder() |
| .id(i) |
| .addAttribute("abc", "xyz") |
| .size(0L) |
| .build() |
| )); |
| } |
| assertEquals(totalFlowFiles, repo.getInMemoryFlowFiles()); |
| } |
| |
| // restore in recovery mode with varying limits |
| additionalProperties.put(RocksDBFlowFileRepository.RocksDbProperty.ENABLE_RECOVERY_MODE.propertyName, "true"); |
| for (int recoveryLimit = 0; recoveryLimit < totalFlowFiles; recoveryLimit += 10) { |
| |
| additionalProperties.put(RocksDBFlowFileRepository.RocksDbProperty.RECOVERY_MODE_FLOWFILE_LIMIT.propertyName, Integer.toString(recoveryLimit)); |
| try (final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties))) { |
| repo.initialize(new StandardResourceClaimManager()); |
| repo.loadFlowFiles(testQueue.provider); |
| assertEquals(recoveryLimit, repo.getInMemoryFlowFiles()); |
| } |
| } |
| |
| // restore in recovery mode with limit equal to available files |
| additionalProperties.put(RocksDBFlowFileRepository.RocksDbProperty.RECOVERY_MODE_FLOWFILE_LIMIT.propertyName, Integer.toString(totalFlowFiles)); |
| try (final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties))) { |
| repo.initialize(new StandardResourceClaimManager()); |
| repo.loadFlowFiles(testQueue.provider); |
| assertEquals(totalFlowFiles, repo.getInMemoryFlowFiles()); |
| } |
| |
| // restore in recovery mode with limit higher than available files |
| additionalProperties.put(RocksDBFlowFileRepository.RocksDbProperty.RECOVERY_MODE_FLOWFILE_LIMIT.propertyName, Integer.toString(Integer.MAX_VALUE)); |
| try (final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties))) { |
| repo.initialize(new StandardResourceClaimManager()); |
| repo.loadFlowFiles(testQueue.provider); |
| assertEquals(totalFlowFiles, repo.getInMemoryFlowFiles()); |
| } |
| |
| // restore in normal mode |
| additionalProperties.put(RocksDBFlowFileRepository.RocksDbProperty.ENABLE_RECOVERY_MODE.propertyName, "false"); |
| additionalProperties.put(RocksDBFlowFileRepository.RocksDbProperty.RECOVERY_MODE_FLOWFILE_LIMIT.propertyName, Integer.toString(0)); |
| |
| try (final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties))) { |
| repo.initialize(new StandardResourceClaimManager()); |
| repo.loadFlowFiles(testQueue.provider); |
| assertEquals(totalFlowFiles, repo.getInMemoryFlowFiles()); |
| } |
| } |
| |
| @Test |
| public void testRecoveryModeWithContinuedLoading() throws Exception { |
| additionalProperties.put(RocksDBFlowFileRepository.RocksDbProperty.CLAIM_CLEANUP_PERIOD.propertyName, "24 hours"); // "disable" the cleanup thread, let us manually force recovery |
| |
| int totalFlowFiles = 50; |
| int recoveryLimit = 10; |
| |
| final TestQueue testQueue = new TestQueue(); |
| |
| try (final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties))) { |
| |
| repo.initialize(new StandardResourceClaimManager()); |
| repo.loadFlowFiles(testQueue.provider); |
| |
| // add records to the repo |
| for (int i = 1; i <= totalFlowFiles; i++) { |
| repo.updateRepository(testQueue.getRepositoryRecord(new StandardFlowFileRecord.Builder() |
| .id(i) |
| .addAttribute("abc", "xyz") |
| .size(0L) |
| .build() |
| )); |
| } |
| assertEquals(totalFlowFiles, repo.getInMemoryFlowFiles()); |
| } |
| |
| // restore in recovery mode |
| additionalProperties.put(RocksDBFlowFileRepository.RocksDbProperty.ENABLE_RECOVERY_MODE.propertyName, "true"); |
| additionalProperties.put(RocksDBFlowFileRepository.RocksDbProperty.RECOVERY_MODE_FLOWFILE_LIMIT.propertyName, Integer.toString(recoveryLimit)); |
| |
| try (final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties))) { |
| repo.initialize(new StandardResourceClaimManager()); |
| repo.loadFlowFiles(testQueue.provider); |
| assertEquals(recoveryLimit, repo.getInMemoryFlowFiles()); |
| assertEquals(totalFlowFiles - recoveryLimit, repo.getRecordsToRestoreCount()); |
| |
| long flowFilesRecovered = repo.getInMemoryFlowFiles(); |
| |
| for (int i = 0; i < 4; i++) { |
| testQueue.deleteQueuedFlowFiles(repo); |
| assertEquals(0, repo.getInMemoryFlowFiles()); |
| |
| repo.doRecovery(); |
| assertEquals(recoveryLimit, repo.getInMemoryFlowFiles()); |
| |
| flowFilesRecovered += repo.getInMemoryFlowFiles(); |
| assertEquals((recoveryLimit * (i + 2)), flowFilesRecovered); |
| assertEquals(totalFlowFiles - flowFilesRecovered, repo.getRecordsToRestoreCount()); |
| } |
| |
| // should have restored all files |
| assertEquals(0, repo.getRecordsToRestoreCount()); |
| assertEquals(recoveryLimit, repo.getInMemoryFlowFiles()); |
| |
| // delete last files |
| testQueue.deleteQueuedFlowFiles(repo); |
| assertEquals(0, repo.getRecordsToRestoreCount()); |
| assertEquals(0, repo.getInMemoryFlowFiles()); |
| |
| repo.doRecovery(); |
| |
| // should have nothing left |
| assertEquals(0, repo.getRecordsToRestoreCount()); |
| assertEquals(0, repo.getInMemoryFlowFiles()); |
| } |
| |
| // restore in normal mode |
| additionalProperties.put(RocksDBFlowFileRepository.RocksDbProperty.ENABLE_RECOVERY_MODE.propertyName, "false"); |
| additionalProperties.put(RocksDBFlowFileRepository.RocksDbProperty.RECOVERY_MODE_FLOWFILE_LIMIT.propertyName, Integer.toString(1)); |
| |
| try (final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties))) { |
| repo.initialize(new StandardResourceClaimManager()); |
| repo.loadFlowFiles(testQueue.provider); |
| assertEquals(0, repo.getRecordsToRestoreCount()); |
| assertEquals(0, repo.getInMemoryFlowFiles()); |
| } |
| } |
| |
| @Test |
| public void testStallStop() throws IOException { |
| final TestQueue testQueue = new TestQueue(); |
| |
| // set stall & stop properties |
| additionalProperties.put(RocksDBFlowFileRepository.RocksDbProperty.ENABLE_STALL_STOP.propertyName, "true"); |
| additionalProperties.put(RocksDBFlowFileRepository.RocksDbProperty.STALL_FLOWFILE_COUNT.propertyName, "2"); |
| additionalProperties.put(RocksDBFlowFileRepository.RocksDbProperty.STOP_FLOWFILE_COUNT.propertyName, "3"); |
| |
| // take heap usage out of the calculation |
| additionalProperties.put(RocksDBFlowFileRepository.RocksDbProperty.STALL_HEAP_USAGE_PERCENT.propertyName, "100%"); |
| additionalProperties.put(RocksDBFlowFileRepository.RocksDbProperty.STOP_HEAP_USAGE_PERCENT.propertyName, "100%"); |
| |
| try (final RocksDBFlowFileRepository repo = new RocksDBFlowFileRepository(NiFiProperties.createBasicNiFiProperties(nifiPropertiesPath, additionalProperties))) { |
| |
| repo.initialize(new StandardResourceClaimManager()); |
| repo.loadFlowFiles(testQueue.provider); |
| |
| StandardFlowFileRecord.Builder ffBuilder = new StandardFlowFileRecord.Builder(); |
| ffBuilder.addAttribute("abc", "xyz"); |
| ffBuilder.size(0L); |
| |
| List<RepositoryRecord> record1 = testQueue.getRepositoryRecord(ffBuilder.id(1).build()); |
| List<RepositoryRecord> record2 = testQueue.getRepositoryRecord(ffBuilder.id(2).build()); |
| List<RepositoryRecord> record3 = testQueue.getRepositoryRecord(ffBuilder.id(3).build()); |
| |
| // CREATE one... should incur no penalty |
| repo.updateRepository(record1); |
| repo.updateStallStop(); |
| assertFalse(repo.stallNewFlowFiles); |
| assertFalse(repo.stopNewFlowFiles); |
| |
| // CREATE another... should stall |
| repo.updateRepository(record2); |
| repo.updateStallStop(); |
| assertTrue(repo.stallNewFlowFiles); |
| assertFalse(repo.stopNewFlowFiles); |
| |
| // CREATE another... should stop |
| repo.updateRepository(record3); |
| repo.updateStallStop(); |
| assertTrue(repo.stallNewFlowFiles); |
| assertTrue(repo.stopNewFlowFiles); |
| |
| // DELETE one... should be stalled but not stopped |
| ((StandardRepositoryRecord) record1.get(0)).markForDelete(); |
| repo.updateRepository(record1); |
| repo.updateStallStop(); |
| assertTrue(repo.stallNewFlowFiles); |
| assertFalse(repo.stopNewFlowFiles); |
| |
| // DELETE another... shouldn't be stalled or stopped |
| ((StandardRepositoryRecord) record2.get(0)).markForDelete(); |
| repo.updateRepository(record2); |
| repo.updateStallStop(); |
| assertFalse(repo.stallNewFlowFiles); |
| assertFalse(repo.stopNewFlowFiles); |
| } |
| } |
| |
| private class TestQueue { |
| private final TestQueueProvider provider; |
| private final Collection<FlowFileRecord> queuedFlowFiles; |
| private final Connection connection; |
| |
| TestQueue() { |
| provider = new TestQueueProvider(); |
| queuedFlowFiles = new ConcurrentSkipListSet<>(); // potentially accessed from multiple threads |
| |
| final FlowFileQueue queue = new StandardFlowFileQueue("1234", null, null, null, null, null, null, null, 0, "0 sec",0, "0 B") { |
| @Override |
| public void put(final FlowFileRecord file) { |
| queuedFlowFiles.add(file); |
| } |
| }; |
| |
| connection = Mockito.mock(Connection.class); |
| when(connection.getIdentifier()).thenReturn(queue.getIdentifier()); |
| when(connection.getFlowFileQueue()).thenReturn(queue); |
| |
| provider.addConnection(connection); |
| |
| } |
| |
| void deleteQueuedFlowFiles(RocksDBFlowFileRepository repo) throws IOException { |
| Collection<RepositoryRecord> recordsToDelete = queuedFlowFiles.stream().map((Function<FlowFileRecord, RepositoryRecord>) flowFileRecord -> { |
| StandardRepositoryRecord record = new StandardRepositoryRecord(null, flowFileRecord); |
| record.markForDelete(); |
| return record; |
| }).collect(Collectors.toSet()); |
| |
| repo.updateRepository(recordsToDelete); |
| queuedFlowFiles.clear(); |
| } |
| |
| private List<RepositoryRecord> getRepositoryRecord(final FlowFileRecord flowFileRecord) { |
| final StandardRepositoryRecord record = new StandardRepositoryRecord(null); |
| record.setWorking(flowFileRecord, false); |
| record.setDestination(connection.getFlowFileQueue()); |
| return Collections.singletonList(record); |
| } |
| } |
| |
| private static class TestQueueProvider implements QueueProvider { |
| |
| private List<Connection> connectionList = new ArrayList<>(); |
| |
| void addConnection(final Connection connection) { |
| this.connectionList.add(connection); |
| } |
| |
| @Override |
| public Collection<FlowFileQueue> getAllQueues() { |
| final List<FlowFileQueue> queueList = new ArrayList<>(); |
| for (final Connection conn : connectionList) { |
| queueList.add(conn.getFlowFileQueue()); |
| } |
| |
| return queueList; |
| } |
| } |
| |
| private static class MockFlowFileSwapManager implements FlowFileSwapManager { |
| |
| private final Map<FlowFileQueue, Map<String, List<FlowFileRecord>>> swappedRecords = new HashMap<>(); |
| |
| @Override |
| public void initialize(SwapManagerInitializationContext initializationContext) { |
| } |
| |
| @Override |
| public String swapOut(List<FlowFileRecord> flowFiles, FlowFileQueue flowFileQueue, final String partitionName) { |
| Map<String, List<FlowFileRecord>> swapMap = swappedRecords.computeIfAbsent(flowFileQueue, k -> new HashMap<>()); |
| |
| final String location = UUID.randomUUID().toString(); |
| swapMap.put(location, new ArrayList<>(flowFiles)); |
| return location; |
| } |
| |
| @Override |
| public SwapContents peek(String swapLocation, FlowFileQueue flowFileQueue) { |
| Map<String, List<FlowFileRecord>> swapMap = swappedRecords.get(flowFileQueue); |
| if (swapMap == null) { |
| return null; |
| } |
| |
| final List<FlowFileRecord> flowFiles = swapMap.get(swapLocation); |
| final SwapSummary summary = getSwapSummary(swapLocation); |
| return new StandardSwapContents(summary, flowFiles); |
| } |
| |
| @Override |
| public SwapContents swapIn(String swapLocation, FlowFileQueue flowFileQueue) { |
| Map<String, List<FlowFileRecord>> swapMap = swappedRecords.get(flowFileQueue); |
| if (swapMap == null) { |
| return null; |
| } |
| |
| final List<FlowFileRecord> flowFiles = swapMap.remove(swapLocation); |
| final SwapSummary summary = getSwapSummary(swapLocation); |
| return new StandardSwapContents(summary, flowFiles); |
| } |
| |
| @Override |
| public List<String> recoverSwapLocations(FlowFileQueue flowFileQueue, final String partitionName) { |
| Map<String, List<FlowFileRecord>> swapMap = swappedRecords.get(flowFileQueue); |
| if (swapMap == null) { |
| return null; |
| } |
| |
| return new ArrayList<>(swapMap.keySet()); |
| } |
| |
| @Override |
| public SwapSummary getSwapSummary(String swapLocation) { |
| List<FlowFileRecord> records = null; |
| for (final Map<String, List<FlowFileRecord>> swapMap : swappedRecords.values()) { |
| records = swapMap.get(swapLocation); |
| if (records != null) { |
| break; |
| } |
| } |
| |
| if (records == null) { |
| return null; |
| } |
| |
| final List<ResourceClaim> resourceClaims = new ArrayList<>(); |
| long size = 0L; |
| Long maxId = null; |
| for (final FlowFileRecord flowFile : records) { |
| size += flowFile.getSize(); |
| |
| if (maxId == null || flowFile.getId() > maxId) { |
| maxId = flowFile.getId(); |
| } |
| |
| final ContentClaim contentClaim = flowFile.getContentClaim(); |
| if (contentClaim != null) { |
| final ResourceClaim resourceClaim = contentClaim.getResourceClaim(); |
| resourceClaims.add(resourceClaim); |
| } |
| } |
| |
| return new StandardSwapSummary(new QueueSize(records.size(), size), maxId, resourceClaims, 0L, 0L); |
| } |
| |
| @Override |
| public void purge() { |
| this.swappedRecords.clear(); |
| } |
| |
| @Override |
| public String getQueueIdentifier(final String swapLocation) { |
| return null; |
| } |
| |
| @Override |
| public Set<String> getSwappedPartitionNames(FlowFileQueue queue) { |
| return Collections.emptySet(); |
| } |
| |
| @Override |
| public String changePartitionName(String swapLocation, String newPartitionName) { |
| return swapLocation; |
| } |
| } |
| |
| |
| } |
