private void restoreWithRescaling ( Collection < KeyedStateHandle > restoreStateHandles ) throws Exception { // Prepare for restore with rescaling KeyedStateHandle initialHandle = RocksDBIncrementalCheckpointUtils . chooseTheBestStateHandleForInitial ( restoreStateHandles , keyGroupRange ) ; // Init base DB instance if ( initialHandle != null ) { restoreStateHandles . remove ( initialHandle ) ; initDBWithRescaling ( initialHandle ) ; } else { openDB ( ) ; } // Transfer remaining key-groups from temporary instance into base DB byte [ ] startKeyGroupPrefixBytes = new byte [ keyGroupPrefixBytes ] ; RocksDBKeySerializationUtils . serializeKeyGroup ( keyGroupRange . getStartKeyGroup ( ) , startKeyGroupPrefixBytes ) ; byte [ ] stopKeyGroupPrefixBytes = new byte [ keyGroupPrefixBytes ] ; RocksDBKeySerializationUtils . serializeKeyGroup ( keyGroupRange . getEndKeyGroup ( ) + 1 , stopKeyGroupPrefixBytes ) ; for ( KeyedStateHandle rawStateHandle : restoreStateHandles ) { if ( ! ( rawStateHandle instanceof IncrementalRemoteKeyedStateHandle ) ) { throw new IllegalStateException ( "Unexpected state handle type, " + "expected " + IncrementalRemoteKeyedStateHandle . class + ", but found " + rawStateHandle . getClass ( ) ) ; } Path temporaryRestoreInstancePath = new Path ( instanceBasePath . getAbsolutePath ( ) + UUID . randomUUID ( ) . toString ( ) ) ; try ( RestoredDBInstance tmpRestoreDBInfo = restoreDBInstanceFromStateHandle ( ( IncrementalRemoteKeyedStateHandle ) rawStateHandle , temporaryRestoreInstancePath ) ; RocksDBWriteBatchWrapper writeBatchWrapper = new RocksDBWriteBatchWrapper ( this . db ) ) { List < ColumnFamilyDescriptor > tmpColumnFamilyDescriptors = tmpRestoreDBInfo . columnFamilyDescriptors ; List < ColumnFamilyHandle > tmpColumnFamilyHandles = tmpRestoreDBInfo . columnFamilyHandles ; // iterating only the requested descriptors automatically skips the default column family handle for ( int i = 0 ; i < tmpColumnFamilyDescriptors . size ( ) ; ++ i ) { ColumnFamilyHandle tmpColumnFamilyHandle = tmpColumnFamilyHandles . get ( i ) ; ColumnFamilyHandle targetColumnFamilyHandle = getOrRegisterStateColumnFamilyHandle ( null , tmpRestoreDBInfo . stateMetaInfoSnapshots . get ( i ) ) . columnFamilyHandle ; try ( RocksIteratorWrapper iterator = RocksDBOperationUtils . getRocksIterator ( tmpRestoreDBInfo . db , tmpColumnFamilyHandle ) ) { iterator . seek ( startKeyGroupPrefixBytes ) ; while ( iterator . isValid ( ) ) { if ( RocksDBIncrementalCheckpointUtils . beforeThePrefixBytes ( iterator . key ( ) , stopKeyGroupPrefixBytes ) ) { writeBatchWrapper . put ( targetColumnFamilyHandle , iterator . key ( ) , iterator . value ( ) ) ; } else { // Since the iterator will visit the record according to the sorted order, // we can just break here. break ; } iterator . next ( ) ; } } // releases native iterator resources } } finally { cleanUpPathQuietly ( temporaryRestoreInstancePath ) ; } } }
Recovery from multi incremental states with rescaling . For rescaling this method creates a temporary RocksDB instance for a key - groups shard . All contents from the temporary instance are copied into the real restore instance and then the temporary instance is discarded .