/*
* 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.hadoop.hive.llap;
import java.util.Map;
import java.util.HashMap;
import java.io.IOException;
import java.net.InetSocketAddress;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.mapred.RecordWriter;
import org.apache.hadoop.util.StringUtils;
import org.apache.hadoop.hive.conf.HiveConf;
import org.apache.hadoop.hive.conf.HiveConf.ConfVars;
import org.apache.hadoop.hive.llap.daemon.rpc.LlapDaemonProtocolProtos.LlapOutputSocketInitMessage;
import org.apache.hadoop.hive.llap.io.ChunkedOutputStream;
import org.apache.hadoop.hive.llap.security.SecretManager;
import com.google.common.base.Preconditions;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import io.netty.bootstrap.ServerBootstrap;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelFutureListener;
import io.netty.channel.ChannelInitializer;
import io.netty.channel.EventLoopGroup;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.protobuf.ProtobufDecoder;
import io.netty.handler.codec.protobuf.ProtobufVarint32FrameDecoder;
import io.netty.handler.codec.string.StringEncoder;
/**
* Responsible for sending back result set data to the connections
* made by external clients via the LLAP input format.
*/
public class LlapOutputFormatService {
private static final Logger LOG = LoggerFactory.getLogger(LlapOutputFormat.class);
private static final AtomicBoolean started = new AtomicBoolean(false);
private static final AtomicBoolean initing = new AtomicBoolean(false);
private static LlapOutputFormatService INSTANCE;
// TODO: the global lock might be to coarse here.
private final Object lock = new Object();
private final Map<String, RecordWriter<?,?>> writers = new HashMap<String, RecordWriter<?,?>>();
private final Map<String, String> errors = new HashMap<String, String>();
private final Configuration conf;
private static final int WAIT_TIME = 5;
private EventLoopGroup eventLoopGroup;
private ServerBootstrap serverBootstrap;
private ChannelFuture listeningChannelFuture;
private int port;
private final SecretManager sm;
private final long writerTimeoutMs;
private LlapOutputFormatService(Configuration conf, SecretManager sm) throws IOException {
this.sm = sm;
this.conf = conf;
this.writerTimeoutMs = HiveConf.getTimeVar(
conf, ConfVars.LLAP_DAEMON_OUTPUT_STREAM_TIMEOUT, TimeUnit.MILLISECONDS);
}
public static void initializeAndStart(Configuration conf, SecretManager sm) throws Exception {
if (!initing.getAndSet(true)) {
INSTANCE = new LlapOutputFormatService(conf, sm);
INSTANCE.start();
started.set(true);
}
}
public static LlapOutputFormatService get() throws IOException {
Preconditions.checkState(started.get(),
"LlapOutputFormatService must be started before invoking get");
return INSTANCE;
}
public void start() throws IOException {
LOG.info("Starting LlapOutputFormatService");
int portFromConf = HiveConf.getIntVar(conf, HiveConf.ConfVars.LLAP_DAEMON_OUTPUT_SERVICE_PORT);
int sendBufferSize = HiveConf.getIntVar(conf,
HiveConf.ConfVars.LLAP_DAEMON_OUTPUT_SERVICE_SEND_BUFFER_SIZE);
eventLoopGroup = new NioEventLoopGroup(1);
serverBootstrap = new ServerBootstrap();
serverBootstrap.group(eventLoopGroup);
serverBootstrap.channel(NioServerSocketChannel.class);
serverBootstrap.childHandler(new LlapOutputFormatServiceChannelHandler(sendBufferSize));
try {
listeningChannelFuture = serverBootstrap.bind(portFromConf).sync();
this.port = ((InetSocketAddress) listeningChannelFuture.channel().localAddress()).getPort();
LOG.info("LlapOutputFormatService: Binding to port: {} with send buffer size: {} ", this.port,
sendBufferSize);
} catch (InterruptedException err) {
throw new IOException("LlapOutputFormatService: Error binding to port " + portFromConf, err);
}
}
public void stop() throws IOException, InterruptedException {
LOG.info("Stopping LlapOutputFormatService");
if (listeningChannelFuture != null) {
listeningChannelFuture.channel().close().sync();
listeningChannelFuture = null;
} else {
LOG.warn("LlapOutputFormatService does not appear to have a listening port to close.");
}
eventLoopGroup.shutdownGracefully(1, WAIT_TIME, TimeUnit.SECONDS).sync();
}
@SuppressWarnings("unchecked")
public <K,V> RecordWriter<K, V> getWriter(String id) throws IOException, InterruptedException {
RecordWriter<?, ?> writer = null;
synchronized (lock) {
long startTime = System.nanoTime();
boolean isFirst = true;
while ((writer = writers.get(id)) == null) {
String error = errors.remove(id);
if (error != null) {
throw new IOException(error);
}
if (isFirst) {
LOG.info("Waiting for writer for " + id);
isFirst = false;
}
if (((System.nanoTime() - startTime) / 1000000) > writerTimeoutMs) {
throw new IOException("The writer for " + id + " has timed out after "
+ writerTimeoutMs + "ms");
}
lock.wait(writerTimeoutMs);
}
}
LOG.info("Returning writer for: "+id);
return (RecordWriter<K, V>) writer;
}
public int getPort() {
return port;
}
protected class LlapOutputFormatServiceHandler
extends SimpleChannelInboundHandler<LlapOutputSocketInitMessage> {
private final int sendBufferSize;
public LlapOutputFormatServiceHandler(final int sendBufferSize) {
this.sendBufferSize = sendBufferSize;
}
@Override
public void channelRead0(ChannelHandlerContext ctx, LlapOutputSocketInitMessage msg) {
String id = msg.getFragmentId();
byte[] tokenBytes = msg.hasToken() ? msg.getToken().toByteArray() : null;
try {
registerReader(ctx, id, tokenBytes);
} catch (Throwable t) {
// Make sure we fail the channel if something goes wrong.
// We internally handle all the "expected" exceptions, so log a lot of information here.
failChannel(ctx, id, StringUtils.stringifyException(t));
}
}
private void registerReader(ChannelHandlerContext ctx, String id, byte[] tokenBytes) {
if (sm != null) {
try {
sm.verifyToken(tokenBytes);
} catch (SecurityException | IOException ex) {
failChannel(ctx, id, ex.getMessage());
return;
}
}
LOG.debug("registering socket for: " + id);
int maxPendingWrites = HiveConf.getIntVar(conf,
HiveConf.ConfVars.LLAP_DAEMON_OUTPUT_SERVICE_MAX_PENDING_WRITES);
@SuppressWarnings("rawtypes")
LlapRecordWriter writer = new LlapRecordWriter(
new ChunkedOutputStream(
new ChannelOutputStream(ctx, id, sendBufferSize, maxPendingWrites),
sendBufferSize));
boolean isFailed = true;
synchronized (lock) {
if (!writers.containsKey(id)) {
isFailed = false;
writers.put(id, writer);
// Add listener to handle any cleanup for when the connection is closed
ctx.channel().closeFuture().addListener(new LlapOutputFormatChannelCloseListener(id));
lock.notifyAll();
}
}
if (isFailed) {
failChannel(ctx, id, "Writer already registered for " + id);
}
}
/** Do not call under lock. */
private void failChannel(ChannelHandlerContext ctx, String id, String error) {
// TODO: write error to the channel? there's no mechanism for that now.
ctx.close();
synchronized (lock) {
errors.put(id, error);
lock.notifyAll();
}
LOG.error(error);
}
}
protected class LlapOutputFormatChannelCloseListener implements ChannelFutureListener {
private String id;
LlapOutputFormatChannelCloseListener(String id) {
this.id = id;
}
@Override
public void operationComplete(ChannelFuture future) throws Exception {
RecordWriter<?, ?> writer = null;
synchronized (INSTANCE) {
writer = writers.remove(id);
}
if (writer == null) {
LOG.warn("Did not find a writer for ID " + id);
}
}
}
protected class LlapOutputFormatServiceChannelHandler extends ChannelInitializer<SocketChannel> {
private final int sendBufferSize;
public LlapOutputFormatServiceChannelHandler(final int sendBufferSize) {
this.sendBufferSize = sendBufferSize;
}
@Override
public void initChannel(SocketChannel ch) throws Exception {
ch.pipeline().addLast(
new ProtobufVarint32FrameDecoder(),
new ProtobufDecoder(LlapOutputSocketInitMessage.getDefaultInstance()),
new StringEncoder(),
new LlapOutputFormatServiceHandler(sendBufferSize));
}
}
}