org.apache.flink.runtime.operators.TempBarrier Maven / Gradle / Ivy
/*
* 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.operators;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.common.typeutils.TypeSerializerFactory;
import org.apache.flink.core.memory.DataInputView;
import org.apache.flink.core.memory.MemorySegment;
import org.apache.flink.runtime.io.disk.InputViewIterator;
import org.apache.flink.runtime.io.disk.SpillingBuffer;
import org.apache.flink.runtime.io.disk.iomanager.IOManager;
import org.apache.flink.runtime.jobgraph.tasks.AbstractInvokable;
import org.apache.flink.runtime.memory.ListMemorySegmentSource;
import org.apache.flink.runtime.memory.MemoryAllocationException;
import org.apache.flink.runtime.memory.MemoryManager;
import org.apache.flink.runtime.operators.util.CloseableInputProvider;
import org.apache.flink.util.MutableObjectIterator;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
/** This class facilitates JVM-local exchange between stages of a batch job. */
public class TempBarrier implements CloseableInputProvider {
private final SpillingBuffer buffer;
private final TypeSerializer serializer;
private final TempWritingThread tempWriter;
private final MemoryManager memManager;
private final Object lock = new Object();
private volatile Throwable exception;
private final ArrayList memory;
private volatile boolean writingDone;
private volatile boolean closed;
// --------------------------------------------------------------------------------------------
public TempBarrier(
AbstractInvokable owner,
MutableObjectIterator input,
TypeSerializerFactory serializerFactory,
MemoryManager memManager,
IOManager ioManager,
int numPages,
List preAllocated)
throws MemoryAllocationException {
this.serializer = serializerFactory.getSerializer();
this.memManager = memManager;
this.memory = new ArrayList<>(numPages);
this.memory.addAll(preAllocated);
if (numPages > memory.size()) {
memManager.allocatePages(owner, memory, numPages - preAllocated.size());
}
this.buffer =
new SpillingBuffer(
ioManager,
new ListMemorySegmentSource(this.memory),
memManager.getPageSize());
this.tempWriter =
new TempWritingThread(input, serializerFactory.getSerializer(), this.buffer);
}
// --------------------------------------------------------------------------------------------
public void startReading() {
this.tempWriter.start();
}
/**
* This method resets the input!
*
* @see org.apache.flink.runtime.operators.util.CloseableInputProvider#getIterator()
*/
@Override
public MutableObjectIterator getIterator() throws InterruptedException, IOException {
synchronized (this.lock) {
while (this.exception == null && !this.writingDone) {
this.lock.wait(5000);
}
}
if (this.exception != null) {
throw new RuntimeException(
"An error occurred creating the temp table.", this.exception);
} else if (this.writingDone) {
final DataInputView in = this.buffer.flip();
return new InputViewIterator<>(in, this.serializer);
} else {
return null;
}
}
@Override
public void close() throws IOException {
memManager.release(prepareToClose());
}
List closeAndGetLeftoverMemory() throws IOException {
return prepareToClose();
}
private List prepareToClose() throws IOException {
synchronized (this.lock) {
if (this.closed) {
return Collections.emptyList();
}
if (this.exception == null) {
this.exception = new Exception("The dam has been closed.");
}
this.lock.notifyAll();
}
try {
this.tempWriter.shutdown();
this.tempWriter.join();
} catch (InterruptedException iex) {
throw new IOException("Interrupted");
}
List toRelease = new ArrayList<>();
toRelease.addAll(buffer.close());
toRelease.addAll(memory);
memory.clear();
return toRelease;
}
private void setException(Throwable t) {
synchronized (this.lock) {
this.exception = t;
this.lock.notifyAll();
}
try {
close();
} catch (Throwable ex) {
// ignored
}
}
private void writingDone() {
synchronized (this.lock) {
this.writingDone = true;
this.lock.notifyAll();
}
}
// --------------------------------------------------------------------------------------------
private final class TempWritingThread extends Thread {
private final MutableObjectIterator input;
private final TypeSerializer serializer;
private final SpillingBuffer buffer;
private volatile boolean running = true;
private TempWritingThread(
MutableObjectIterator input,
TypeSerializer serializer,
SpillingBuffer buffer) {
super("Temp writer");
setDaemon(true);
this.input = input;
this.serializer = serializer;
this.buffer = buffer;
}
@Override
public void run() {
final MutableObjectIterator input = this.input;
final TypeSerializer serializer = this.serializer;
final SpillingBuffer buffer = this.buffer;
try {
T record = serializer.createInstance();
while (this.running && ((record = input.next(record)) != null)) {
serializer.serialize(record, buffer);
}
TempBarrier.this.writingDone();
} catch (Throwable t) {
TempBarrier.this.setException(t);
}
}
public void shutdown() {
this.running = false;
this.interrupt();
}
}
}