/* * 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.flink.runtime.state; import org.apache.commons.io.IOUtils; import org.apache.flink.api.common.typeutils.TypeSerializer; import org.apache.flink.api.common.typeutils.UnloadableDummyTypeSerializer; import org.apache.flink.core.fs.CloseableRegistry; import org.apache.flink.core.fs.FSDataInputStream; import org.apache.flink.core.memory.DataInputView; import org.apache.flink.core.memory.DataInputViewStreamWrapper; import org.apache.flink.runtime.state.metainfo.StateMetaInfoSnapshot; import org.apache.flink.util.Preconditions; import javax.annotation.Nonnull; import java.io.IOException; import java.util.Collection; import java.util.List; import java.util.Map; /** * Implementation of operator state restore operation. */ public class OperatorStateRestoreOperation implements RestoreOperation<Void> { private final CloseableRegistry closeStreamOnCancelRegistry; private final ClassLoader userClassloader; private final Map<String, PartitionableListState<?>> registeredOperatorStates; private final Map<String, BackendWritableBroadcastState<?, ?>> registeredBroadcastStates; private final Collection<OperatorStateHandle> stateHandles; public OperatorStateRestoreOperation( CloseableRegistry closeStreamOnCancelRegistry, ClassLoader userClassloader, Map<String, PartitionableListState<?>> registeredOperatorStates, Map<String, BackendWritableBroadcastState<?, ?>> registeredBroadcastStates, @Nonnull Collection<OperatorStateHandle> stateHandles) { this.closeStreamOnCancelRegistry = closeStreamOnCancelRegistry; this.userClassloader = userClassloader; this.registeredOperatorStates = registeredOperatorStates; this.registeredBroadcastStates = registeredBroadcastStates; this.stateHandles = stateHandles; } @Override public Void restore() throws Exception { if (stateHandles.isEmpty()) { return null; } for (OperatorStateHandle stateHandle : stateHandles) { if (stateHandle == null) { continue; } FSDataInputStream in = stateHandle.openInputStream(); closeStreamOnCancelRegistry.registerCloseable(in); ClassLoader restoreClassLoader = Thread.currentThread().getContextClassLoader(); try { Thread.currentThread().setContextClassLoader(userClassloader); OperatorBackendSerializationProxy backendSerializationProxy = new OperatorBackendSerializationProxy(userClassloader); backendSerializationProxy.read(new DataInputViewStreamWrapper(in)); List<StateMetaInfoSnapshot> restoredOperatorMetaInfoSnapshots = backendSerializationProxy.getOperatorStateMetaInfoSnapshots(); // Recreate all PartitionableListStates from the meta info for (StateMetaInfoSnapshot restoredSnapshot : restoredOperatorMetaInfoSnapshots) { final RegisteredOperatorStateBackendMetaInfo<?> restoredMetaInfo = new RegisteredOperatorStateBackendMetaInfo<>(restoredSnapshot); if (restoredMetaInfo.getPartitionStateSerializer() instanceof UnloadableDummyTypeSerializer) { // must fail now if the previous typeSerializer cannot be restored because there is no typeSerializer // capable of reading previous state // TODO when eager state registration is in place, we can try to get a convert deserializer // TODO from the newly registered typeSerializer instead of simply failing here throw new IOException("Unable to restore operator state [" + restoredSnapshot.getName() + "]." + " The previous typeSerializer of the operator state must be present; the typeSerializer could" + " have been removed from the classpath, or its implementation have changed and could" + " not be loaded. This is a temporary restriction that will be fixed in future versions."); } PartitionableListState<?> listState = registeredOperatorStates.get(restoredSnapshot.getName()); if (null == listState) { listState = new PartitionableListState<>(restoredMetaInfo); registeredOperatorStates.put(listState.getStateMetaInfo().getName(), listState); } else { // TODO with eager state registration in place, check here for typeSerializer migration strategies } } // ... and then get back the broadcast state. List<StateMetaInfoSnapshot> restoredBroadcastMetaInfoSnapshots = backendSerializationProxy.getBroadcastStateMetaInfoSnapshots(); for (StateMetaInfoSnapshot restoredSnapshot : restoredBroadcastMetaInfoSnapshots) { final RegisteredBroadcastStateBackendMetaInfo<?, ?> restoredMetaInfo = new RegisteredBroadcastStateBackendMetaInfo<>(restoredSnapshot); if (restoredMetaInfo.getKeySerializer() instanceof UnloadableDummyTypeSerializer || restoredMetaInfo.getValueSerializer() instanceof UnloadableDummyTypeSerializer) { // must fail now if the previous typeSerializer cannot be restored because there is no typeSerializer // capable of reading previous state // TODO when eager state registration is in place, we can try to get a convert deserializer // TODO from the newly registered typeSerializer instead of simply failing here throw new IOException("Unable to restore broadcast state [" + restoredSnapshot.getName() + "]." + " The previous key and value serializers of the state must be present; the serializers could" + " have been removed from the classpath, or their implementations have changed and could" + " not be loaded. This is a temporary restriction that will be fixed in future versions."); } BackendWritableBroadcastState<?, ?> broadcastState = registeredBroadcastStates.get(restoredSnapshot.getName()); if (broadcastState == null) { broadcastState = new HeapBroadcastState<>(restoredMetaInfo); registeredBroadcastStates.put(broadcastState.getStateMetaInfo().getName(), broadcastState); } else { // TODO with eager state registration in place, check here for typeSerializer migration strategies } } // Restore all the states for (Map.Entry<String, OperatorStateHandle.StateMetaInfo> nameToOffsets : stateHandle.getStateNameToPartitionOffsets().entrySet()) { final String stateName = nameToOffsets.getKey(); PartitionableListState<?> listStateForName = registeredOperatorStates.get(stateName); if (listStateForName == null) { BackendWritableBroadcastState<?, ?> broadcastStateForName = registeredBroadcastStates.get(stateName); Preconditions.checkState(broadcastStateForName != null, "Found state without " + "corresponding meta info: " + stateName); deserializeBroadcastStateValues(broadcastStateForName, in, nameToOffsets.getValue()); } else { deserializeOperatorStateValues(listStateForName, in, nameToOffsets.getValue()); } } } finally { Thread.currentThread().setContextClassLoader(restoreClassLoader); if (closeStreamOnCancelRegistry.unregisterCloseable(in)) { IOUtils.closeQuietly(in); } } } return null; } private <S> void deserializeOperatorStateValues( PartitionableListState<S> stateListForName, FSDataInputStream in, OperatorStateHandle.StateMetaInfo metaInfo) throws IOException { if (null != metaInfo) { long[] offsets = metaInfo.getOffsets(); if (null != offsets) { DataInputView div = new DataInputViewStreamWrapper(in); TypeSerializer<S> serializer = stateListForName.getStateMetaInfo().getPartitionStateSerializer(); for (long offset : offsets) { in.seek(offset); stateListForName.add(serializer.deserialize(div)); } } } } private <K, V> void deserializeBroadcastStateValues( final BackendWritableBroadcastState<K, V> broadcastStateForName, final FSDataInputStream in, final OperatorStateHandle.StateMetaInfo metaInfo) throws Exception { if (metaInfo != null) { long[] offsets = metaInfo.getOffsets(); if (offsets != null) { TypeSerializer<K> keySerializer = broadcastStateForName.getStateMetaInfo().getKeySerializer(); TypeSerializer<V> valueSerializer = broadcastStateForName.getStateMetaInfo().getValueSerializer(); in.seek(offsets[0]); DataInputView div = new DataInputViewStreamWrapper(in); int size = div.readInt(); for (int i = 0; i < size; i++) { broadcastStateForName.put(keySerializer.deserialize(div), valueSerializer.deserialize(div)); } } } } }