@SuppressWarnings("unchecked")
private <X extends StateHandle<?>> X serializeDeserialize(X handle)
throws IOException, ClassNotFoundException {
byte[] serialized = InstantiationUtil.serializeObject(handle);
return (X)
InstantiationUtil.deserializeObject(
serialized, Thread.currentThread().getContextClassLoader());
}
private <T> T copyObject(T object) {
try {
return InstantiationUtil.deserializeObject(
InstantiationUtil.serializeObject(object), getClass().getClassLoader());
} catch (IOException | ClassNotFoundException e) {
throw new RuntimeException("Failed to copy object.");
}
}
private org.apache.hadoop.fs.FileSystem instantiateFileSystem(
Class<? extends org.apache.hadoop.fs.FileSystem> fsClass) throws IOException {
try {
return fsClass.newInstance();
} catch (ExceptionInInitializerError e) {
throw new IOException(
"The filesystem class '"
+ fsClass.getName()
+ "' throw an exception upon initialization.",
e.getException());
} catch (Throwable t) {
String errorMessage = InstantiationUtil.checkForInstantiationError(fsClass);
if (errorMessage != null) {
throw new IOException(
"The filesystem class '"
+ fsClass.getName()
+ "' cannot be instantiated: "
+ errorMessage);
} else {
throw new IOException(
"An error occurred while instantiating the filesystem class '"
+ fsClass.getName()
+ "'.",
t);
}
}
}
/**
* Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current
* state handle
*
* @throws IOException
* @throws ClassNotFoundException
* @throws RocksDBException
*/
private void restoreKVStateMetaData()
throws IOException, ClassNotFoundException, RocksDBException {
// read number of k/v states
int numColumns = currentStateHandleInView.readInt();
// those two lists are aligned and should later have the same size!
currentStateHandleKVStateColumnFamilies = new ArrayList<>(numColumns);
// restore the empty columns for the k/v states through the metadata
for (int i = 0; i < numColumns; i++) {
StateDescriptor<?, ?> stateDescriptor =
InstantiationUtil.deserializeObject(
currentStateHandleInStream, rocksDBKeyedStateBackend.userCodeClassLoader);
Tuple2<ColumnFamilyHandle, StateDescriptor<?, ?>> columnFamily =
rocksDBKeyedStateBackend.kvStateInformation.get(stateDescriptor.getName());
if (null == columnFamily) {
ColumnFamilyDescriptor columnFamilyDescriptor =
new ColumnFamilyDescriptor(
stateDescriptor.getName().getBytes(), rocksDBKeyedStateBackend.columnOptions);
columnFamily =
new Tuple2<ColumnFamilyHandle, StateDescriptor<?, ?>>(
rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor),
stateDescriptor);
rocksDBKeyedStateBackend.kvStateInformation.put(stateDescriptor.getName(), columnFamily);
}
currentStateHandleKVStateColumnFamilies.add(columnFamily.f0);
}
}
/** Writes the {@code Context} to the given state checkpoint output. */
protected void writeToState(StateBackend.CheckpointStateOutputView out) throws IOException {
keySerializer.serialize(key, out);
windowSerializer.serialize(window, out);
out.writeLong(watermarkTimer);
out.writeLong(processingTimeTimer);
byte[] serializedState = InstantiationUtil.serializeObject(state);
out.writeInt(serializedState.length);
out.write(serializedState, 0, serializedState.length);
MultiplexingStreamRecordSerializer<IN> recordSerializer =
new MultiplexingStreamRecordSerializer<>(inputSerializer);
out.writeInt(windowBuffer.size());
for (StreamRecord<IN> element : windowBuffer.getElements()) {
recordSerializer.serialize(element, out);
}
}
/**
* Constructs a new {@code Context} by reading from a {@link DataInputView} that contains a
* serialized context that we wrote in {@link
* #writeToState(StateBackend.CheckpointStateOutputView)}
*/
@SuppressWarnings("unchecked")
protected Context(DataInputView in, ClassLoader userClassloader) throws Exception {
this.key = keySerializer.deserialize(in);
this.window = windowSerializer.deserialize(in);
this.watermarkTimer = in.readLong();
this.processingTimeTimer = in.readLong();
int stateSize = in.readInt();
byte[] stateData = new byte[stateSize];
in.read(stateData);
state = InstantiationUtil.deserializeObject(stateData, userClassloader);
this.windowBuffer = windowBufferFactory.create();
int numElements = in.readInt();
MultiplexingStreamRecordSerializer<IN> recordSerializer =
new MultiplexingStreamRecordSerializer<>(inputSerializer);
for (int i = 0; i < numElements; i++) {
windowBuffer.storeElement(recordSerializer.deserialize(in).<IN>asRecord());
}
}
private void writeKVStateMetaData() throws IOException, InterruptedException {
// write number of k/v states
outputView.writeInt(stateBackend.kvStateInformation.size());
int kvStateId = 0;
// iterate all column families, where each column family holds one k/v state, to write the
// metadata
for (Map.Entry<String, Tuple2<ColumnFamilyHandle, StateDescriptor<?, ?>>> column :
stateBackend.kvStateInformation.entrySet()) {
// be cooperative and check for interruption from time to time in the hot loop
checkInterrupted();
// write StateDescriptor for this k/v state
InstantiationUtil.serializeObject(outStream, column.getValue().f1);
// retrieve iterator for this k/v states
readOptions = new ReadOptions();
readOptions.setSnapshot(snapshot);
RocksIterator iterator = stateBackend.db.newIterator(column.getValue().f0, readOptions);
kvStateIterators.add(new Tuple2<>(iterator, kvStateId));
++kvStateId;
}
}
@Override
public void invoke() throws Exception {
if (LOG.isDebugEnabled()) {
LOG.debug(getLogString("Starting data sink operator"));
}
ExecutionConfig executionConfig;
try {
ExecutionConfig c =
(ExecutionConfig)
InstantiationUtil.readObjectFromConfig(
getJobConfiguration(), ExecutionConfig.CONFIG_KEY, getUserCodeClassLoader());
if (c != null) {
executionConfig = c;
} else {
LOG.warn("The execution config returned by the configuration was null");
executionConfig = new ExecutionConfig();
}
} catch (IOException e) {
throw new RuntimeException("Could not load ExecutionConfig from Job Configuration: " + e);
} catch (ClassNotFoundException e) {
throw new RuntimeException("Could not load ExecutionConfig from Job Configuration: " + e);
}
boolean objectReuseEnabled = executionConfig.isObjectReuseEnabled();
try {
// initialize local strategies
MutableObjectIterator<IT> input1;
switch (this.config.getInputLocalStrategy(0)) {
case NONE:
// nothing to do
localStrategy = null;
input1 = reader;
break;
case SORT:
// initialize sort local strategy
try {
// get type comparator
TypeComparatorFactory<IT> compFact =
this.config.getInputComparator(0, getUserCodeClassLoader());
if (compFact == null) {
throw new Exception("Missing comparator factory for local strategy on input " + 0);
}
// initialize sorter
UnilateralSortMerger<IT> sorter =
new UnilateralSortMerger<IT>(
getEnvironment().getMemoryManager(),
getEnvironment().getIOManager(),
this.reader,
this,
this.inputTypeSerializerFactory,
compFact.createComparator(),
this.config.getRelativeMemoryInput(0),
this.config.getFilehandlesInput(0),
this.config.getSpillingThresholdInput(0));
this.localStrategy = sorter;
input1 = sorter.getIterator();
} catch (Exception e) {
throw new RuntimeException(
"Initializing the input processing failed"
+ (e.getMessage() == null ? "." : ": " + e.getMessage()),
e);
}
break;
default:
throw new RuntimeException("Invalid local strategy for DataSinkTask");
}
// read the reader and write it to the output
final TypeSerializer<IT> serializer = this.inputTypeSerializerFactory.getSerializer();
final MutableObjectIterator<IT> input = input1;
final OutputFormat<IT> format = this.format;
// check if task has been canceled
if (this.taskCanceled) {
return;
}
if (LOG.isDebugEnabled()) {
LOG.debug(getLogString("Starting to produce output"));
}
// open
format.open(
this.getEnvironment().getIndexInSubtaskGroup(),
this.getEnvironment().getNumberOfSubtasks());
if (objectReuseEnabled) {
IT record = serializer.createInstance();
// work!
while (!this.taskCanceled && ((record = input.next(record)) != null)) {
format.writeRecord(record);
}
} else {
IT record;
// work!
while (!this.taskCanceled && ((record = input.next()) != null)) {
format.writeRecord(record);
}
}
// close. We close here such that a regular close throwing an exception marks a task as
// failed.
if (!this.taskCanceled) {
this.format.close();
this.format = null;
}
} catch (Exception ex) {
// make a best effort to clean up
try {
if (!cleanupCalled && format instanceof CleanupWhenUnsuccessful) {
cleanupCalled = true;
((CleanupWhenUnsuccessful) format).tryCleanupOnError();
}
} catch (Throwable t) {
LOG.error("Cleanup on error failed.", t);
}
ex = ExceptionInChainedStubException.exceptionUnwrap(ex);
if (ex instanceof CancelTaskException) {
// forward canceling exception
throw ex;
}
// drop, if the task was canceled
else if (!this.taskCanceled) {
if (LOG.isErrorEnabled()) {
LOG.error(getLogString("Error in user code: " + ex.getMessage()), ex);
}
throw ex;
}
} finally {
if (this.format != null) {
// close format, if it has not been closed, yet.
// This should only be the case if we had a previous error, or were canceled.
try {
this.format.close();
} catch (Throwable t) {
if (LOG.isWarnEnabled()) {
LOG.warn(getLogString("Error closing the output format"), t);
}
}
}
// close local strategy if necessary
if (localStrategy != null) {
try {
this.localStrategy.close();
} catch (Throwable t) {
LOG.error("Error closing local strategy", t);
}
}
RegularPactTask.clearReaders(new MutableReader[] {inputReader});
}
if (!this.taskCanceled) {
if (LOG.isDebugEnabled()) {
LOG.debug(getLogString("Finished data sink operator"));
}
} else {
if (LOG.isDebugEnabled()) {
LOG.debug(getLogString("Data sink operator cancelled"));
}
}
}
/** For backwards compatibility, remove again later! */
@Deprecated
private void restoreOldSavepointKeyedState(Collection<KeyGroupsStateHandle> restoreState)
throws Exception {
if (restoreState.isEmpty()) {
return;
}
Preconditions.checkState(1 == restoreState.size(), "Only one element expected here.");
HashMap<String, RocksDBStateBackend.FinalFullyAsyncSnapshot> namedStates =
InstantiationUtil.deserializeObject(
restoreState.iterator().next().openInputStream(), userCodeClassLoader);
Preconditions.checkState(1 == namedStates.size(), "Only one element expected here.");
DataInputView inputView =
namedStates.values().iterator().next().stateHandle.getState(userCodeClassLoader);
// clear k/v state information before filling it
kvStateInformation.clear();
// first get the column family mapping
int numColumns = inputView.readInt();
Map<Byte, StateDescriptor> columnFamilyMapping = new HashMap<>(numColumns);
for (int i = 0; i < numColumns; i++) {
byte mappingByte = inputView.readByte();
ObjectInputStream ooIn =
new InstantiationUtil.ClassLoaderObjectInputStream(
new DataInputViewStream(inputView), userCodeClassLoader);
StateDescriptor stateDescriptor = (StateDescriptor) ooIn.readObject();
columnFamilyMapping.put(mappingByte, stateDescriptor);
// this will fill in the k/v state information
getColumnFamily(stateDescriptor);
}
// try and read until EOF
try {
// the EOFException will get us out of this...
while (true) {
byte mappingByte = inputView.readByte();
ColumnFamilyHandle handle = getColumnFamily(columnFamilyMapping.get(mappingByte));
byte[] keyAndNamespace = BytePrimitiveArraySerializer.INSTANCE.deserialize(inputView);
ByteArrayInputStreamWithPos bis = new ByteArrayInputStreamWithPos(keyAndNamespace);
K reconstructedKey = keySerializer.deserialize(new DataInputViewStreamWrapper(bis));
int len = bis.getPosition();
int keyGroup =
(byte) KeyGroupRangeAssignment.assignToKeyGroup(reconstructedKey, numberOfKeyGroups);
if (keyGroupPrefixBytes == 1) {
// copy and override one byte (42) between key and namespace
System.arraycopy(keyAndNamespace, 0, keyAndNamespace, 1, len);
keyAndNamespace[0] = (byte) keyGroup;
} else {
byte[] largerKey = new byte[1 + keyAndNamespace.length];
// write key-group
largerKey[0] = (byte) ((keyGroup >> 8) & 0xFF);
largerKey[1] = (byte) (keyGroup & 0xFF);
// write key
System.arraycopy(keyAndNamespace, 0, largerKey, 2, len);
// skip one byte (42), write namespace
System.arraycopy(
keyAndNamespace, 1 + len, largerKey, 2 + len, keyAndNamespace.length - len - 1);
keyAndNamespace = largerKey;
}
byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(inputView);
db.put(handle, keyAndNamespace, value);
}
} catch (EOFException e) {
// expected
}
}
/**
* Writes to the output stream the error return code, and the given exception in serialized form.
*
* @param out Thr output stream to write to.
* @param t The exception to send.
* @throws IOException Thrown, if the output stream could not be written to.
*/
private static void writeErrorToStream(OutputStream out, Throwable t) throws IOException {
byte[] bytes = InstantiationUtil.serializeObject(t);
out.write(RETURN_ERROR);
writeLength(bytes.length, out);
out.write(bytes);
}