/**
* 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.avro.mapred.tether;
import java.io.IOException;
import java.io.ByteArrayOutputStream;
import java.nio.ByteBuffer;
import java.net.InetSocketAddress;
import java.net.URL;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.avro.Schema;
import org.apache.avro.ipc.HttpTransceiver;
import org.apache.avro.ipc.Transceiver;
import org.apache.avro.ipc.SaslSocketTransceiver;
import org.apache.avro.ipc.specific.SpecificRequestor;
import org.apache.avro.io.DecoderFactory;
import org.apache.avro.io.BinaryDecoder;
import org.apache.avro.io.BinaryEncoder;
import org.apache.avro.io.EncoderFactory;
import org.apache.avro.specific.SpecificDatumReader;
import org.apache.avro.specific.SpecificDatumWriter;
/** Base class for Java tether mapreduce programs. Useless except for testing,
* since it's already possible to write Java MapReduce programs without
* tethering. Also serves as an example of how a framework may be
* implemented. */
public abstract class TetherTask<IN,MID,OUT> {
static final Logger LOG = LoggerFactory.getLogger(TetherTask.class);
private Transceiver clientTransceiver;
private OutputProtocol outputClient;
private TaskType taskType;
private int partitions;
private DecoderFactory decoderFactory = DecoderFactory.get();
private BinaryDecoder decoder;
private SpecificDatumReader<IN> inReader;
private SpecificDatumReader<MID> midReader;
private IN inRecord;
private MID midRecord;
private MID midRecordSpare;
private Collector<MID> midCollector;
private Collector<OUT> outCollector;
private TetheredProcess.Protocol proto;
private static class Buffer extends ByteArrayOutputStream {
public ByteBuffer data() {
return ByteBuffer.wrap(buf, 0, count);
}
}
/** Collector for map and reduce output values. */
public class Collector<T> {
private SpecificDatumWriter<T> writer;
private Buffer buffer = new Buffer();
private BinaryEncoder encoder = new EncoderFactory()
.configureBlockSize(512).binaryEncoder(buffer, null);
private Collector(Schema schema) {
this.writer = new SpecificDatumWriter<T>(schema);
}
/** Collect a map or reduce output value. */
public void collect(T record) throws IOException {
buffer.reset();
writer.write(record, encoder);
encoder.flush();
outputClient.output(buffer.data());
}
/** Collect a pre-partitioned map output value. */
public void collect(T record, int partition) throws IOException {
buffer.reset();
writer.write(record, encoder);
encoder.flush();
outputClient.outputPartitioned(partition, buffer.data());
}
}
void open(int inputPort) throws IOException {
// open output client, connecting to parent
String clientPortString = System.getenv("AVRO_TETHER_OUTPUT_PORT");
String protocol = System.getenv("AVRO_TETHER_PROTOCOL");
if (clientPortString == null)
throw new RuntimeException("AVRO_TETHER_OUTPUT_PORT env var is null");
int clientPort = Integer.parseInt(clientPortString);
if (protocol == null) {
throw new RuntimeException("AVRO_TETHER_PROTOCOL env var is null");
}
protocol=protocol.trim().toLowerCase();
if (protocol.equals("http")) {
proto=TetheredProcess.Protocol.HTTP;
} else if (protocol.equals("sasl")) {
proto=TetheredProcess.Protocol.SASL;
} else {
throw new RuntimeException("AVROT_TETHER_PROTOCOL="+protocol+" but this protocol is unsupported");
}
switch (proto) {
case SASL:
this.clientTransceiver =
new SaslSocketTransceiver(new InetSocketAddress(clientPort));
this.outputClient = SpecificRequestor.getClient(OutputProtocol.class, clientTransceiver);
break;
case HTTP:
this.clientTransceiver =new HttpTransceiver(new URL("http://127.0.0.1:"+clientPort));
this.outputClient = SpecificRequestor.getClient(OutputProtocol.class, clientTransceiver);
break;
}
// send inputPort to parent
outputClient.configure(inputPort);
}
void configure(TaskType taskType, CharSequence inSchemaText, CharSequence outSchemaText) {
this.taskType = taskType;
try {
Schema inSchema = Schema.parse(inSchemaText.toString());
Schema outSchema = Schema.parse(outSchemaText.toString());
switch (taskType) {
case MAP:
this.inReader = new SpecificDatumReader<IN>(inSchema);
this.midCollector = new Collector<MID>(outSchema);
break;
case REDUCE:
this.midReader = new SpecificDatumReader<MID>(inSchema);
this.outCollector = new Collector<OUT>(outSchema);
break;
}
} catch (Throwable e) {
fail(e.toString());
}
}
void partitions(int partitions) { this.partitions = partitions; }
/** Return the number of map output partitions of this job. */
public int partitions() { return partitions; }
void input(ByteBuffer data, long count) {
try {
decoder = decoderFactory.binaryDecoder(data.array(), decoder);
for (long i = 0; i < count; i++) {
switch (taskType) {
case MAP:
inRecord = inReader.read(inRecord, decoder);
map(inRecord, midCollector);
break;
case REDUCE:
MID prev = midRecord;
midRecord = midReader.read(midRecordSpare, decoder);
if (prev != null && !midRecord.equals(prev))
reduceFlush(prev, outCollector);
reduce(midRecord, outCollector);
midRecordSpare = prev;
break;
}
}
} catch (Throwable e) {
LOG.warn("failing: "+e, e);
fail(e.toString());
}
}
void complete() {
if (taskType == TaskType.REDUCE && midRecord != null)
try {
reduceFlush(midRecord, outCollector);
} catch (Throwable e) {
LOG.warn("failing: "+e, e);
fail(e.toString());
}
LOG.info("TetherTask: Sending complete to parent process.");
outputClient.complete();
LOG.info("TetherTask: Done sending complete to parent process.");
}
/** Called with input values to generate intermediate values. */
public abstract void map(IN record, Collector<MID> collector)
throws IOException;
/** Called with sorted intermediate values. */
public abstract void reduce(MID record, Collector<OUT> collector)
throws IOException;
/** Called with the last intermediate value in each equivalence run. */
public abstract void reduceFlush(MID record, Collector<OUT> collector)
throws IOException;
/** Call to update task status. */
public void status(String message) {
outputClient.status(message);
}
/** Call to increment a counter. */
public void count(String group, String name, long amount) {
outputClient.count(group, name, amount);
}
/** Call to fail the task. */
public void fail(String message) {
outputClient.fail(message);
close();
}
void close() {
LOG.info("Closing the transciever");
if (clientTransceiver != null)
try {
clientTransceiver.close();
} catch (IOException e) {} // ignore
}
}