/*
* 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.io.network.api;
import java.io.IOException;
import org.apache.flink.core.io.IOReadableWritable;
import org.apache.flink.runtime.event.task.AbstractEvent;
import org.apache.flink.runtime.io.network.Buffer;
import org.apache.flink.runtime.io.network.bufferprovider.BufferProvider;
import org.apache.flink.runtime.io.network.channels.EndOfSuperstepEvent;
import org.apache.flink.runtime.io.network.serialization.RecordSerializer;
import org.apache.flink.runtime.io.network.serialization.SpanningRecordSerializer;
import org.apache.flink.runtime.jobgraph.tasks.AbstractInvokable;
/**
* A record writer connects the application to an output gate. It allows the application
* of emit (send out) to the output gate. The output gate will then take care of distributing
* the emitted records among the output channels.
*
* @param <T>
* the type of the record that can be emitted with this record writer
*/
public class RecordWriter<T extends IOReadableWritable> extends BufferWriter {
private final BufferProvider bufferPool;
private final ChannelSelector<T> channelSelector;
private int numChannels;
/** RecordSerializer per outgoing channel */
private RecordSerializer<T>[] serializers;
// -----------------------------------------------------------------------------------------------------------------
public RecordWriter(AbstractInvokable invokable) {
this(invokable, new RoundRobinChannelSelector<T>());
}
public RecordWriter(AbstractInvokable invokable, ChannelSelector<T> channelSelector) {
// initialize the gate
super(invokable);
this.bufferPool = invokable.getEnvironment().getOutputBufferProvider();
this.channelSelector = channelSelector;
}
// -----------------------------------------------------------------------------------------------------------------
@SuppressWarnings("unchecked")
public void initializeSerializers() {
this.numChannels = this.outputGate.getNumChannels();
this.serializers = new RecordSerializer[numChannels];
for (int i = 0; i < this.numChannels; i++) {
this.serializers[i] = new SpanningRecordSerializer<T>();
}
}
public void emit(final T record) throws IOException, InterruptedException {
for (int targetChannel : this.channelSelector.selectChannels(record, this.numChannels)) {
// serialize with corresponding serializer and send full buffer
RecordSerializer<T> serializer = this.serializers[targetChannel];
RecordSerializer.SerializationResult result = serializer.addRecord(record);
while (result.isFullBuffer()) {
Buffer buffer = serializer.getCurrentBuffer();
if (buffer != null) {
sendBuffer(buffer, targetChannel);
}
buffer = this.bufferPool.requestBufferBlocking(this.bufferPool.getBufferSize());
result = serializer.setNextBuffer(buffer);
}
}
}
public void flush() throws IOException, InterruptedException {
for (int targetChannel = 0; targetChannel < this.numChannels; targetChannel++) {
RecordSerializer<T> serializer = this.serializers[targetChannel];
Buffer buffer = serializer.getCurrentBuffer();
if (buffer != null) {
sendBuffer(buffer, targetChannel);
}
serializer.clear();
}
}
@Override
public void broadcastEvent(AbstractEvent event) throws IOException, InterruptedException {
for (int targetChannel = 0; targetChannel < this.numChannels; targetChannel++) {
RecordSerializer<T> serializer = this.serializers[targetChannel];
Buffer buffer = serializer.getCurrentBuffer();
if (buffer == null) {
super.sendEvent(event, targetChannel);
} else {
super.sendBufferAndEvent(buffer, event, targetChannel);
buffer = this.bufferPool.requestBufferBlocking(this.bufferPool.getBufferSize());
serializer.setNextBuffer(buffer);
}
}
}
@Override
public void sendEndOfSuperstep() throws IOException, InterruptedException {
for (int targetChannel = 0; targetChannel < this.numChannels; targetChannel++) {
RecordSerializer<T> serializer = this.serializers[targetChannel];
Buffer buffer = serializer.getCurrentBuffer();
if (buffer == null) {
super.sendEvent(EndOfSuperstepEvent.INSTANCE, targetChannel);
} else {
super.sendBufferAndEvent(buffer, EndOfSuperstepEvent.INSTANCE, targetChannel);
buffer = this.bufferPool.requestBufferBlocking(this.bufferPool.getBufferSize());
serializer.setNextBuffer(buffer);
}
}
}
public void clearBuffers() {
if (this.serializers != null) {
for (RecordSerializer<?> s: this.serializers) {
Buffer b = s.getCurrentBuffer();
if (b != null) {
b.recycleBuffer();
}
}
}
}
}