KeyValueServiceMigrators { public static KeyValueServiceMigrator setupMigrator(MigratorSpec migratorSpec) { AtlasDbServices fromServices = migratorSpec.fromServices(); AtlasDbServices toServices = migratorSpec.toServices(); TimestampManagementService toTimestampManagementService = toServices.getManagedTimestampService(); toTimestampManagementService.fastForwardTimestamp( fromServices.getManagedTimestampService().getFreshTimestamp() + 1); long migrationStartTimestamp = toServices.getManagedTimestampService().getFreshTimestamp(); long migrationCommitTimestamp = toServices.getManagedTimestampService().getFreshTimestamp(); toServices.getTransactionService().putUnlessExists(migrationStartTimestamp, migrationCommitTimestamp); return new KeyValueServiceMigrator( CHECKPOINT_NAMESPACE, fromServices.getTransactionManager(), toServices.getTransactionManager(), fromServices.getKeyValueService(), toServices.getKeyValueService(), Suppliers.ofInstance(migrationStartTimestamp), migratorSpec.threads(), migratorSpec.batchSize(), ImmutableMap.of(), (String message, KeyValueServiceMigrator.KvsMigrationMessageLevel level) -> printer.info(level.toString() + ": " + message), new TaskProgress() { @Override public void beginTask(String message, int tasks) { printer.info(message); } @Override public void subTaskComplete() { } @Override public void taskComplete() { } }, ImmutableSet.of()); } private  KeyValueServiceMigrators(); static KeyValueServiceMigrator setupMigrator(MigratorSpec migratorSpec);  }
@Test public void setupMigratorFastForwardsTimestamp() { fromServices.getManagedTimestampService().fastForwardTimestamp(FUTURE_TIMESTAMP); assertThat(toServices.getManagedTimestampService().getFreshTimestamp()).isLessThan(FUTURE_TIMESTAMP); KeyValueServiceMigrators.setupMigrator(migratorSpec); assertThat(toServices.getManagedTimestampService().getFreshTimestamp()) .isGreaterThanOrEqualTo(FUTURE_TIMESTAMP); }
@Test public void setupMigratorCommitsOneTransaction() { KeyValueServiceMigrators.setupMigrator(migratorSpec); ArgumentCaptor<Long> startTimestampCaptor = ArgumentCaptor.forClass(Long.class); ArgumentCaptor<Long> commitTimestampCaptor = ArgumentCaptor.forClass(Long.class); verify(toServices.getTransactionService()).putUnlessExists( startTimestampCaptor.capture(), commitTimestampCaptor.capture()); assertThat(startTimestampCaptor.getValue()).isLessThan(commitTimestampCaptor.getValue()); assertThat(commitTimestampCaptor.getValue()) .isLessThan(toServices.getManagedTimestampService().getFreshTimestamp()); }
@Test public void setupClearsOutExistingTablesExceptAtomic() { toKvs.createTables(TEST_AND_CHECKPOINT_TABLES); fromKvs.dropTables(fromKvs.getAllTableNames()); KeyValueServiceMigrator migrator = KeyValueServiceMigrators.setupMigrator(migratorSpec); migrator.setup(); assertThat(toKvs.getAllTableNames()).containsExactlyInAnyOrder(TransactionConstants.TRANSACTION_TABLE, TransactionConstants.TRANSACTIONS2_TABLE, AtlasDbConstants.COORDINATION_TABLE); }
@Test public void setupCreatesAndResetsExistingTables() { fromKvs.createTables(TEST_AND_CHECKPOINT_TABLES); toKvs.createTable(TEST_TABLE, AtlasDbConstants.GENERIC_TABLE_METADATA); toKvs.put(TEST_TABLE, ImmutableMap.of(TEST_CELL, TEST_VALUE1), 1); assertThat(toKvs.get(TEST_TABLE, ImmutableMap.of(TEST_CELL, Long.MAX_VALUE))).isNotEmpty(); KeyValueServiceMigrator migrator = KeyValueServiceMigrators.setupMigrator(migratorSpec); migrator.setup(); assertThat(toKvs.getAllTableNames()).containsExactlyInAnyOrderElementsOf(fromKvs.getAllTableNames()); assertThat(toKvs.get(TEST_TABLE, ImmutableMap.of(TEST_CELL, Long.MAX_VALUE))).isEmpty(); }
@Test public void checkpointTableOnSourceKvsIsIgnored() { fromKvs.createTables(TEST_AND_CHECKPOINT_TABLES); KeyValueServiceMigrator migrator = KeyValueServiceMigrators.setupMigrator(migratorSpec); migrator.setup(); migrator.migrate(); verify(fromKvs, never()).get(eq(CHECKPOINT_TABLE), any()); verify(fromKvs, never()).getRange(eq(CHECKPOINT_TABLE), any(), anyLong()); verify(fromKvs, never()).getRows(eq(CHECKPOINT_TABLE), any(), any(), anyLong()); verify(fromKvs, never()).getRowsColumnRange(eq(CHECKPOINT_TABLE), any(), any(), anyLong()); verify(fromKvs, never()).getRowsColumnRange(eq(CHECKPOINT_TABLE), any(), any(), anyInt(), anyLong()); verify(fromKvs, never()).getFirstBatchForRanges(eq(CHECKPOINT_TABLE), any(), anyLong()); }
@Test public void setupCopiesTableMetadata() { toKvs.createTables(TEST_AND_CHECKPOINT_TABLES); KeyValueServiceMigrator migrator = KeyValueServiceMigrators.setupMigrator(migratorSpec); migrator.setup(); assertThat(fromKvs.getMetadataForTable(TEST_TABLE)).isEqualTo(toKvs.getMetadataForTable(TEST_TABLE)); }
@Test public void migrateRevertsUncommittedWritesAndMigratesMostRecentlyCommitted() { fromKvs.createTables(TEST_AND_CHECKPOINT_TABLES); fromTxManager.runTaskWithRetry(tx -> { tx.put(TEST_TABLE, ImmutableMap.of(TEST_CELL, TEST_VALUE1)); return tx.getTimestamp(); }); long uncommittedTs = fromServices.getManagedTimestampService().getFreshTimestamp(); fromKvs.put(TEST_TABLE, ImmutableMap.of(TEST_CELL, TEST_VALUE2), uncommittedTs); KeyValueServiceMigrator migrator = KeyValueServiceMigrators.setupMigrator(migratorSpec); migrator.setup(); migrator.migrate(); assertThat(fromServices.getTransactionService().get(uncommittedTs)) .isEqualTo(TransactionConstants.FAILED_COMMIT_TS); assertThat(toKvs.get(TEST_TABLE, ImmutableMap.of(TEST_CELL, Long.MAX_VALUE)).get(TEST_CELL).getContents()) .containsExactly(TEST_VALUE1); }
@Test public void migrateOnlyMigratesMostRecentVersions() { fromKvs.createTables(TEST_AND_CHECKPOINT_TABLES); fromTxManager.runTaskWithRetry(tx -> { tx.put(TEST_TABLE, ImmutableMap.of(TEST_CELL, TEST_VALUE1)); tx.put(TEST_TABLE, ImmutableMap.of(TEST_CELL2, TEST_VALUE1)); return tx.getTimestamp(); }); fromTxManager.runTaskWithRetry(tx -> { tx.put(TEST_TABLE, ImmutableMap.of(TEST_CELL, TEST_VALUE2)); return tx.getTimestamp(); }); KeyValueServiceMigrator migrator = KeyValueServiceMigrators.setupMigrator(migratorSpec); migrator.setup(); migrator.migrate(); toTxManager.runTaskThrowOnConflict(tx -> { Map<Cell, byte[]> values = tx.get(TEST_TABLE, ImmutableSet.of(TEST_CELL, TEST_CELL2)); assertThat(values.get(TEST_CELL)).isEqualTo(TEST_VALUE2); assertThat(values.get(TEST_CELL2)).isEqualTo(TEST_VALUE1); return null; }); assertThat(toKvs.getAllTimestamps(TEST_TABLE, ImmutableSet.of(TEST_CELL), Long.MAX_VALUE).size()).isEqualTo(1); }
@Test public void deletedEntriesAreNotMigrated() { fromKvs.createTables(TEST_AND_CHECKPOINT_TABLES); fromTxManager.runTaskWithRetry(tx -> { tx.put(TEST_TABLE, ImmutableMap.of(TEST_CELL, TEST_VALUE1)); return tx.getTimestamp(); }); fromTxManager.runTaskWithRetry(tx -> { tx.delete(TEST_TABLE, ImmutableSet.of(TEST_CELL)); return tx.getTimestamp(); }); KeyValueServiceMigrator migrator = KeyValueServiceMigrators.setupMigrator(migratorSpec); migrator.setup(); migrator.migrate(); assertThat(toKvs.get(TEST_TABLE, ImmutableMap.of(TEST_CELL, Long.MAX_VALUE))).isEmpty(); }
@Test public void cleanupDropsCheckpointTable() { fromKvs.createTables(TEST_AND_CHECKPOINT_TABLES); fromTxManager.runTaskWithRetry(tx -> { tx.put(TEST_TABLE, ImmutableMap.of(TEST_CELL, TEST_VALUE1)); return tx.getTimestamp(); }); KeyValueServiceMigrator migrator = KeyValueServiceMigrators.setupMigrator(migratorSpec); migrator.setup(); migrator.migrate(); verify(toKvs, never()).dropTable(CHECKPOINT_TABLE); migrator.cleanup(); verify(toKvs, times(1)).dropTable(CHECKPOINT_TABLE); }
@Test public void tablesDelegatedToSourceKvsGetDroppedFromSourceKvsIfMigratable() { fromKvs.createTable(FAKE_ATOMIC_TABLE, AtlasDbConstants.GENERIC_TABLE_METADATA); fromKvs.putUnlessExists(FAKE_ATOMIC_TABLE, ImmutableMap.of(TEST_CELL, TEST_VALUE1)); KeyValueService toTableSplittingKvs = TableSplittingKeyValueService.create( ImmutableList.of(new InMemoryKeyValueService(false), fromKvs), ImmutableMap.of(FAKE_ATOMIC_TABLE, fromKvs)); AtlasDbServices toSplittingServices = createMock(toTableSplittingKvs); ImmutableMigratorSpec spec = ImmutableMigratorSpec.builder() .fromServices(fromServices) .toServices(toSplittingServices) .build(); assertThat(toSplittingServices.getKeyValueService() .get(FAKE_ATOMIC_TABLE, ImmutableMap.of(TEST_CELL, 1L)).get(TEST_CELL).getContents()) .containsExactly(TEST_VALUE1); KeyValueServiceMigrator migrator = KeyValueServiceMigrators.setupMigrator(spec); migrator.setup(); assertThat(fromKvs.getAllTableNames()).doesNotContain(FAKE_ATOMIC_TABLE); }
@Test public void atomicTablesDelegatedToSourceAreNotDropped() { KeyValueService toTableSplittingKvs = TableSplittingKeyValueService.create( ImmutableList.of(new InMemoryKeyValueService(false), fromKvs), Maps.toMap(AtlasDbConstants.HIDDEN_TABLES, ignore -> fromKvs)); AtlasDbServices toSplittingServices = createMock(toTableSplittingKvs); ImmutableMigratorSpec spec = ImmutableMigratorSpec.builder() .fromServices(fromServices) .toServices(toSplittingServices) .build(); fromServices.getTransactionService().putUnlessExists(100_000, 100_001); KeyValueServiceMigrator migrator = KeyValueServiceMigrators.setupMigrator(spec); migrator.setup(); assertThat(toSplittingServices.getTransactionService().get(100_000)).isEqualTo(100_001L); }