Java tutorial
/* * 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 * <p> * http://www.apache.org/licenses/LICENSE-2.0 * <p> * 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.nifi.util.hive; import java.io.IOException; import java.lang.reflect.UndeclaredThrowableException; import java.security.PrivilegedExceptionAction; import java.util.List; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import org.apache.hadoop.hive.conf.HiveConf; import org.apache.hadoop.security.UserGroupInformation; import org.apache.hive.hcatalog.streaming.HiveEndPoint; import org.apache.hive.hcatalog.streaming.RecordWriter; import org.apache.hive.hcatalog.streaming.SerializationError; import org.apache.hive.hcatalog.streaming.StreamingConnection; import org.apache.hive.hcatalog.streaming.StreamingException; import org.apache.hive.hcatalog.streaming.StreamingIOFailure; import org.apache.hive.hcatalog.streaming.StrictJsonWriter; import org.apache.hive.hcatalog.streaming.TransactionBatch; import org.slf4j.Logger; import org.slf4j.LoggerFactory; public class HiveWriter { private static final Logger LOG = LoggerFactory.getLogger(HiveWriter.class); private final HiveEndPoint endPoint; private final StreamingConnection connection; private final int txnsPerBatch; private final RecordWriter recordWriter; private final ExecutorService callTimeoutPool; private final long callTimeout; private final Object txnBatchLock = new Object(); private final UserGroupInformation ugi; private TransactionBatch txnBatch; private long lastUsed; // time of last flush on this writer protected boolean closed; // flag indicating HiveWriter was closed private int totalRecords = 0; public HiveWriter(HiveEndPoint endPoint, int txnsPerBatch, boolean autoCreatePartitions, long callTimeout, ExecutorService callTimeoutPool, UserGroupInformation ugi, HiveConf hiveConf) throws InterruptedException, ConnectFailure { try { this.ugi = ugi; this.callTimeout = callTimeout; this.callTimeoutPool = callTimeoutPool; this.endPoint = endPoint; this.connection = newConnection(endPoint, autoCreatePartitions, hiveConf, ugi); this.txnsPerBatch = txnsPerBatch; this.recordWriter = getRecordWriter(endPoint, ugi, hiveConf); this.txnBatch = nextTxnBatch(recordWriter); this.closed = false; this.lastUsed = System.currentTimeMillis(); } catch (InterruptedException | RuntimeException | ConnectFailure e) { throw e; } catch (Exception e) { throw new ConnectFailure(endPoint, e); } } protected RecordWriter getRecordWriter(HiveEndPoint endPoint, UserGroupInformation ugi, HiveConf hiveConf) throws StreamingException, IOException, InterruptedException { if (ugi == null) { return new StrictJsonWriter(endPoint, hiveConf); } else { try { return ugi.doAs((PrivilegedExceptionAction<StrictJsonWriter>) () -> new StrictJsonWriter(endPoint, hiveConf)); } catch (UndeclaredThrowableException e) { Throwable cause = e.getCause(); if (cause instanceof StreamingException) { throw (StreamingException) cause; } else { throw e; } } } } @Override public String toString() { return "{ endPoint = " + endPoint + ", TransactionBatch = " + txnBatch + " }"; } /** * Write the record data to Hive * * @throws IOException if an error occurs during the write * @throws InterruptedException if the write operation is interrupted */ public synchronized void write(final byte[] record) throws WriteFailure, SerializationError, InterruptedException { if (closed) { throw new IllegalStateException( "This hive streaming writer was closed " + "and thus no longer able to write : " + endPoint); } // write the tuple try { LOG.debug("Writing event to {}", endPoint); callWithTimeout(new CallRunner<Void>() { @Override public Void call() throws StreamingException, InterruptedException { txnBatch.write(record); totalRecords++; return null; } }); } catch (SerializationError se) { throw new SerializationError(endPoint.toString() + " SerializationError", se); } catch (StreamingException | TimeoutException e) { throw new WriteFailure(endPoint, txnBatch.getCurrentTxnId(), e); } } /** * Commits the current Txn if totalRecordsPerTransaction > 0 . * If 'rollToNext' is true, will switch to next Txn in batch or to a * new TxnBatch if current Txn batch is exhausted */ public void flush(boolean rollToNext) throws CommitFailure, TxnBatchFailure, TxnFailure, InterruptedException { // if there are no records do not call flush if (totalRecords <= 0) return; try { synchronized (txnBatchLock) { commitTxn(); nextTxn(rollToNext); totalRecords = 0; lastUsed = System.currentTimeMillis(); } } catch (StreamingException e) { throw new TxnFailure(txnBatch, e); } } /** Queues up a heartbeat request on the current and remaining txns using the * heartbeatThdPool and returns immediately */ public void heartBeat() throws InterruptedException { // 1) schedule the heartbeat on one thread in pool synchronized (txnBatchLock) { try { callWithTimeout(new CallRunner<Void>() { @Override public Void call() throws Exception { try { LOG.info("Sending heartbeat on batch " + txnBatch); txnBatch.heartbeat(); } catch (StreamingException e) { LOG.warn("Heartbeat error on batch " + txnBatch, e); } return null; } }); } catch (InterruptedException e) { throw e; } catch (Exception e) { LOG.warn("Unable to send heartbeat on Txn Batch " + txnBatch, e); // Suppressing exceptions as we don't care for errors on heartbeats } } } /** * Returns totalRecords written so far in a transaction * @returns totalRecords */ public int getTotalRecords() { return totalRecords; } /** * Flush and Close current transactionBatch. */ public void flushAndClose() throws TxnBatchFailure, TxnFailure, CommitFailure, IOException, InterruptedException { flush(false); close(); } /** * Close the Transaction Batch and connection * @throws IOException if an error occurs during close * @throws InterruptedException if the close operation is interrupted */ public void close() throws IOException, InterruptedException { closeTxnBatch(); closeConnection(); closed = true; } protected void closeConnection() throws InterruptedException { LOG.info("Closing connection to end point : {}", endPoint); try { callWithTimeout(new CallRunner<Void>() { @Override public Void call() throws Exception { connection.close(); // could block return null; } }); } catch (Exception e) { LOG.warn("Error closing connection to EndPoint : " + endPoint, e); // Suppressing exceptions as we don't care for errors on connection close } } protected void commitTxn() throws CommitFailure, InterruptedException { LOG.debug("Committing Txn id {} to {}", txnBatch.getCurrentTxnId(), endPoint); try { callWithTimeout(new CallRunner<Void>() { @Override public Void call() throws Exception { txnBatch.commit(); // could block return null; } }); } catch (StreamingException | TimeoutException e) { throw new CommitFailure(endPoint, txnBatch.getCurrentTxnId(), e); } } protected StreamingConnection newConnection(HiveEndPoint endPoint, boolean autoCreatePartitions, HiveConf conf, UserGroupInformation ugi) throws InterruptedException, ConnectFailure { try { return callWithTimeout(() -> { return endPoint.newConnection(autoCreatePartitions, conf, ugi); // could block }); } catch (StreamingException | TimeoutException e) { throw new ConnectFailure(endPoint, e); } } protected TransactionBatch nextTxnBatch(final RecordWriter recordWriter) throws InterruptedException, TxnBatchFailure { LOG.debug("Fetching new Txn Batch for {}", endPoint); TransactionBatch batch = null; try { batch = callWithTimeout(() -> { return connection.fetchTransactionBatch(txnsPerBatch, recordWriter); // could block }); batch.beginNextTransaction(); LOG.debug("Acquired {}. Switching to first txn", batch); } catch (TimeoutException | StreamingException e) { throw new TxnBatchFailure(endPoint, e); } return batch; } protected void closeTxnBatch() throws InterruptedException { try { LOG.debug("Closing Txn Batch {}", txnBatch); callWithTimeout(new CallRunner<Void>() { @Override public Void call() throws Exception { if (txnBatch != null) { txnBatch.close(); // could block } return null; } }); } catch (InterruptedException e) { throw e; } catch (Exception e) { LOG.warn("Error closing txn batch " + txnBatch, e); } } /** * Aborts the current Txn and switches to next Txn. * @throws StreamingException if could not get new Transaction Batch, or switch to next Txn */ public void abort() throws StreamingException, TxnBatchFailure, InterruptedException { synchronized (txnBatchLock) { abortTxn(); nextTxn(true); // roll to next } } /** * Aborts current Txn in the txnBatch. */ protected void abortTxn() throws InterruptedException { LOG.info("Aborting Txn id {} on End Point {}", txnBatch.getCurrentTxnId(), endPoint); try { callWithTimeout(new CallRunner<Void>() { @Override public Void call() throws StreamingException, InterruptedException { txnBatch.abort(); // could block return null; } }); } catch (InterruptedException e) { throw e; } catch (TimeoutException e) { LOG.warn("Timeout while aborting Txn " + txnBatch.getCurrentTxnId() + " on EndPoint: " + endPoint, e); } catch (Exception e) { LOG.warn("Error aborting Txn " + txnBatch.getCurrentTxnId() + " on EndPoint: " + endPoint, e); // Suppressing exceptions as we don't care for errors on abort } } /** * if there are remainingTransactions in current txnBatch, begins nextTransactions * otherwise creates new txnBatch. * @param rollToNext Whether to roll to the next transaction batch */ protected void nextTxn(boolean rollToNext) throws StreamingException, InterruptedException, TxnBatchFailure { if (txnBatch.remainingTransactions() == 0) { closeTxnBatch(); txnBatch = null; if (rollToNext) { txnBatch = nextTxnBatch(recordWriter); } } else if (rollToNext) { LOG.debug("Switching to next Txn for {}", endPoint); txnBatch.beginNextTransaction(); // does not block } } /** * If the current thread has been interrupted, then throws an * exception. * @throws InterruptedException uf the current thread has been interrupted */ protected static void checkAndThrowInterruptedException() throws InterruptedException { if (Thread.currentThread().interrupted()) { throw new InterruptedException("Timed out before Hive call was made. " + "Your callTimeout might be set too low or Hive calls are " + "taking too long."); } } /** * Execute the callable on a separate thread and wait for the completion * for the specified amount of time in milliseconds. In case of timeout * cancel the callable and throw an IOException */ private <T> T callWithTimeout(final CallRunner<T> callRunner) throws TimeoutException, StreamingException, InterruptedException { Future<T> future = callTimeoutPool.submit(() -> { if (ugi == null) { return callRunner.call(); } try { return ugi.doAs((PrivilegedExceptionAction<T>) () -> callRunner.call()); } catch (UndeclaredThrowableException e) { Throwable cause = e.getCause(); // Unwrap exception so it is thrown the same way as without ugi if (!(cause instanceof Exception)) { throw e; } throw (Exception) cause; } }); try { if (callTimeout > 0) { return future.get(callTimeout, TimeUnit.MILLISECONDS); } else { return future.get(); } } catch (TimeoutException eT) { future.cancel(true); throw eT; } catch (ExecutionException e1) { Throwable cause = e1.getCause(); if (cause instanceof IOException) { throw new StreamingIOFailure("I/O Failure", (IOException) cause); } else if (cause instanceof StreamingException) { throw (StreamingException) cause; } else if (cause instanceof InterruptedException) { throw (InterruptedException) cause; } else if (cause instanceof RuntimeException) { throw (RuntimeException) cause; } else if (cause instanceof TimeoutException) { throw new StreamingException("Operation Timed Out.", (TimeoutException) cause); } else { throw new RuntimeException(e1); } } } public long getLastUsed() { return lastUsed; } private byte[] generateRecord(List<String> tuple) { StringBuilder buf = new StringBuilder(); for (String o : tuple) { buf.append(o); buf.append(","); } return buf.toString().getBytes(); } /** * Simple interface whose <tt>call</tt> method is called by * {#callWithTimeout} in a new thread inside a * {@linkplain java.security.PrivilegedExceptionAction#run()} call. * @param <T> the type of object returned from the call */ private interface CallRunner<T> { T call() throws Exception; } public static class Failure extends Exception { public Failure(String message, Throwable cause) { super(message, cause); } } public static class WriteFailure extends Failure { public WriteFailure(HiveEndPoint endPoint, Long currentTxnId, Throwable cause) { super("Failed writing to : " + endPoint + ". TxnID : " + currentTxnId, cause); } } public static class CommitFailure extends Failure { public CommitFailure(HiveEndPoint endPoint, Long txnID, Throwable cause) { super("Commit of Txn " + txnID + " failed on EndPoint: " + endPoint, cause); } } public static class ConnectFailure extends Failure { public ConnectFailure(HiveEndPoint ep, Throwable cause) { super("Failed connecting to EndPoint " + ep, cause); } } public static class TxnBatchFailure extends Failure { public TxnBatchFailure(HiveEndPoint ep, Throwable cause) { super("Failed acquiring Transaction Batch from EndPoint: " + ep, cause); } } public static class TxnFailure extends Failure { public TxnFailure(TransactionBatch txnBatch, Throwable cause) { super("Failed switching to next Txn in TxnBatch " + txnBatch, cause); } } }