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 java.io.IOException;
import java.util.ArrayList;
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;
/**
*
*/
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) throws MemoryAllocationException
{
this.serializer = serializerFactory.getSerializer();
this.memManager = memManager;
this.memory = new ArrayList(numPages);
memManager.allocatePages(owner, this.memory, numPages);
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 {
synchronized (this.lock) {
if (this.closed) {
return;
}
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");
}
this.memManager.release(this.buffer.close());
this.memManager.release(this.memory);
}
private void setException(Throwable t) {
synchronized (this.lock) {
this.exception = t;
this.lock.notifyAll();
}
try {
close();
} catch (Throwable ex) {}
}
private void writingDone() throws IOException {
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();
}
}
}