htrace-kudu/src/main/java/org/apache/htrace/impl/KuduSpanReceiver.java - incubator-retired-htrace - Git at Google

 /*
  * 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.htrace.impl;

 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.htrace.core.HTraceConfiguration;
 import org.apache.htrace.core.Span;
 import org.apache.htrace.core.SpanReceiver;
 import org.apache.htrace.core.TimelineAnnotation;
 import org.apache.htrace.protobuf.generated.SpanProtos;
 import org.kududb.client.Bytes;
 import org.kududb.client.KuduClient;
 import org.kududb.client.KuduSession;
 import org.kududb.client.KuduTable;
 import org.kududb.client.Insert;
 import org.kududb.client.PartialRow;

 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ArrayBlockingQueue;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;

 public class KuduSpanReceiver extends SpanReceiver {

   private static final Log LOG = LogFactory.getLog(KuduSpanReceiver.class);

   private static final int SHUTDOWN_TIMEOUT = 30;
   private static final int MAX_ERRORS = 10;
   private final BlockingQueue<Span> queue;
   private final AtomicBoolean running = new AtomicBoolean(true);
   private final KuduClientConfiguration clientConf;
   private final int maxSpanBatchSize;
   private final ThreadFactory threadFactory = new ThreadFactory() {
     private final AtomicLong receiverIndex = new AtomicLong(0);

     @Override
     public Thread newThread(Runnable runnable) {
       Thread thread = new Thread(runnable);
       thread.setDaemon(true);
       thread.setName(String.format("kuduSpanReceiver-%d",
               receiverIndex.getAndIncrement()));
       return thread;
     }
   };
   private ExecutorService service;
   private String table;
   private String column_span_id;
   private String column_span;
   private String column_root_span;
   private String column_root_span_start_time;

   public KuduSpanReceiver(HTraceConfiguration conf) {
     this.clientConf =
             new KuduClientConfiguration(conf.get(KuduReceiverConstants.KUDU_MASTER_HOST_KEY,
                     KuduReceiverConstants.DEFAULT_KUDU_MASTER_HOST),
                     conf.get(KuduReceiverConstants.KUDU_MASTER_PORT_KEY,
                             KuduReceiverConstants.DEFAULT_KUDU_MASTER_PORT));
     if (conf.get(KuduReceiverConstants.KUDU_CLIENT_BOSS_COUNT_KEY) != null) {
       this.clientConf.setBossCount(Integer.valueOf(conf.get(KuduReceiverConstants.KUDU_CLIENT_BOSS_COUNT_KEY)));
     }
     if (conf.get(KuduReceiverConstants.KUDU_CLIENT_WORKER_COUNT_KEY) != null) {
       this.clientConf.setWorkerCount(Integer.valueOf(conf.get(KuduReceiverConstants.KUDU_CLIENT_WORKER_COUNT_KEY)));
     }
     if (conf.get(KuduReceiverConstants.KUDU_CLIENT_STATISTICS_ENABLED_KEY) != null) {
       this.clientConf.setIsStatisticsEnabled(Boolean.valueOf(conf.get(KuduReceiverConstants.KUDU_CLIENT_STATISTICS_ENABLED_KEY)));
     }
     if (conf.get(KuduReceiverConstants.KUDU_CLIENT_TIMEOUT_ADMIN_OPERATION_KEY) != null) {
       this.clientConf
               .setAdminOperationTimeout(Long.valueOf(conf.get(KuduReceiverConstants.KUDU_CLIENT_TIMEOUT_ADMIN_OPERATION_KEY)));
     }
     if (conf.get(KuduReceiverConstants.KUDU_CLIENT_TIMEOUT_OPERATION_KEY) != null) {
       this.clientConf
               .setOperationTimeout(Long.valueOf(conf.get(KuduReceiverConstants.KUDU_CLIENT_TIMEOUT_OPERATION_KEY)));
     }
     if (conf.get(KuduReceiverConstants.KUDU_CLIENT_TIMEOUT_SOCKET_READ_KEY) != null) {
       this.clientConf
               .setSocketReadTimeout(Long.valueOf(conf.get(KuduReceiverConstants.KUDU_CLIENT_TIMEOUT_SOCKET_READ_KEY)));
     }
     this.queue = new ArrayBlockingQueue<Span>(conf.getInt(KuduReceiverConstants.SPAN_BLOCKING_QUEUE_SIZE_KEY,
             KuduReceiverConstants.DEFAULT_SPAN_BLOCKING_QUEUE_SIZE));
     this.table = conf.get(KuduReceiverConstants.KUDU_TABLE_KEY, KuduReceiverConstants.DEFAULT_KUDU_TABLE);
     this.column_span_id = conf.get(KuduReceiverConstants.KUDU_COLUMN_SPAN_ID_KEY,
             KuduReceiverConstants.DEFAULT_KUDU_COLUMN_SPAN_ID);
     this.column_span = conf.get(KuduReceiverConstants.KUDU_COLUMN_SPAN_KEY,
             KuduReceiverConstants.DEFAULT_KUDU_COLUMN_SPAN);
     this.column_root_span = conf.get(KuduReceiverConstants.KUDU_COLUMN_ROOT_SPAN_KEY,
             KuduReceiverConstants.DEFAULT_KUDU_COLUMN_ROOT_SPAN);
     this.column_root_span_start_time = conf.get(KuduReceiverConstants.KUDU_COLUMN_ROOT_SPAN_START_TIME_KEY,
             KuduReceiverConstants.DEFAULT_KUDU_COLUMN_ROOT_SPAN_START_TIME);
     this.maxSpanBatchSize = conf.getInt(KuduReceiverConstants.MAX_SPAN_BATCH_SIZE_KEY,
             KuduReceiverConstants.DEFAULT_MAX_SPAN_BATCH_SIZE);
     if (this.service != null) {
       this.service.shutdownNow();
       this.service = null;
     }
     int numThreads = conf.getInt(KuduReceiverConstants.NUM_PARALLEL_THREADS_KEY,
             KuduReceiverConstants.DEFAULT_NUM_PARALLEL_THREADS);
     this.service = Executors.newFixedThreadPool(numThreads, threadFactory);
     for (int i = 0; i < numThreads; i++) {
       this.service.submit(new KuduSpanReceiver.WriteSpanRunnable());
     }
   }

   @Override
   public void close() throws IOException {
     running.set(false);
     service.shutdown();
     try {
       if (!service.awaitTermination(SHUTDOWN_TIMEOUT, TimeUnit.SECONDS)) {
         LOG.error("Timeout " + SHUTDOWN_TIMEOUT + " " + TimeUnit.SECONDS +
                 " reached while shutting worker threads which process enqued spans." +
                 " Enqueued spans which are left in blocking queue is dropped.");
       }
     } catch (InterruptedException e) {
       LOG.warn("Interrupted exception occured while terminating thread service executor.", e);
     }
   }

   @Override
   public void receiveSpan(Span span) {
     if (running.get()) {
       try {
         this.queue.add(span);
       } catch (IllegalStateException e) {
         LOG.error("Error trying to enqueue span ("
                 + span.getDescription()
                 + ") to the queue. Blocking Queue is currently reached its capacity.");
       }
     }
   }

   private class WriteSpanRunnable implements Runnable {

     private KuduSession session;
     private KuduClient client;

     @Override
     public void run() {
       SpanProtos.Span.Builder sbuilder = SpanProtos.Span.newBuilder();
       SpanProtos.TimelineAnnotation.Builder tlbuilder =
               SpanProtos.TimelineAnnotation.newBuilder();
       List<Span> dequeuedSpans = new ArrayList<Span>(maxSpanBatchSize);
       long errorCount = 0;
       while (running.get() || queue.size() > 0) {
         Span firstSpan = null;
         try {
           firstSpan = queue.poll(1, TimeUnit.SECONDS);
           if (firstSpan != null) {
             dequeuedSpans.add(firstSpan);
             queue.drainTo(dequeuedSpans, maxSpanBatchSize - 1);
           }
         } catch (InterruptedException ie) {
           LOG.error("Interrupted Exception occurred while polling to " +
                   "retrieve first span from blocking queue");
         }
         startSession();
         if (dequeuedSpans.isEmpty()) {
           try {
             this.session.flush();
           } catch (java.lang.Exception e) {
             LOG.error("Failed to flush writes to Kudu.");
             closeSession();
           }
           continue;
         }
         try {
           for (Span span : dequeuedSpans) {
             sbuilder.clear()
                     .setTraceId(span.getSpanId().getHigh())
                     .setStart(span.getStartTimeMillis())
                     .setStop(span.getStopTimeMillis())
                     .setSpanId(span.getSpanId().getLow())
                     .setProcessId(span.getTracerId())
                     .setDescription(span.getDescription());

             if (span.getParents().length == 0) {
               sbuilder.setParentId(0);
             } else if (span.getParents().length > 0) {
               sbuilder.setParentId(span.getParents()[0].getLow());
             }
             for (TimelineAnnotation ta : span.getTimelineAnnotations()) {
               sbuilder.addTimeline(tlbuilder.clear()
                       .setTime(ta.getTime())
                       .setMessage(ta.getMessage())
                       .build());
             }
             KuduTable tableRef = client.openTable(table);
             Insert insert = tableRef.newInsert();
             PartialRow row = insert.getRow();
             row.addBinary(column_span_id, Bytes.fromLong(span.getSpanId().getHigh()));
             row.addBinary(column_span, sbuilder.build().toByteArray());
             if (span.getParents().length == 0) {
               row.addBinary(column_root_span_start_time, Bytes.fromLong(span.getStartTimeMillis()));
               row.addBinary(column_root_span, sbuilder.build().toByteArray());
             }
             this.session.apply(insert);
           }
           dequeuedSpans.clear();
           errorCount = 0;
         } catch (Exception e) {
           errorCount += 1;
           if (errorCount < MAX_ERRORS) {
             try {
               queue.addAll(dequeuedSpans);
             } catch (IllegalStateException ex) {
               LOG.error("Exception occured while writing spans kudu datastore. " +
                       "Trying to re-enqueue de-queued spans to blocking queue for writing but failed. " +
                       "Dropped " + dequeuedSpans.size() + " dequeued span(s) which were due written" +
                       "into kudu datastore");
             }
           }
           closeSession();
           try {
             Thread.sleep(500);
           } catch (InterruptedException e1) {
             LOG.error("Interrupted Exception occurred while allowing kudu to re-stabilized");
           }
         }
       }
       closeSession();
     }

     private void closeSession() {
       try {
         if (this.session != null) {
           this.session.close();
           this.session = null;
         }
       } catch (java.lang.Exception e) {
         LOG.warn("Failed to close Kudu session. " + e.getMessage());
       }
     }

     private void startSession() {
       if (this.session == null) {
         if (this.client == null) {
           client = clientConf.buildClient();
         }
         session = client.newSession();
       }
     }
   }

 }
	/*
	* 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.htrace.impl;

	import org.apache.commons.logging.Log;
	import org.apache.commons.logging.LogFactory;
	import org.apache.htrace.core.HTraceConfiguration;
	import org.apache.htrace.core.Span;
	import org.apache.htrace.core.SpanReceiver;
	import org.apache.htrace.core.TimelineAnnotation;
	import org.apache.htrace.protobuf.generated.SpanProtos;
	import org.kududb.client.Bytes;
	import org.kududb.client.KuduClient;
	import org.kududb.client.KuduSession;
	import org.kududb.client.KuduTable;
	import org.kududb.client.Insert;
	import org.kududb.client.PartialRow;

	import java.io.IOException;
	import java.util.ArrayList;
	import java.util.List;
	import java.util.concurrent.ExecutorService;
	import java.util.concurrent.BlockingQueue;
	import java.util.concurrent.ThreadFactory;
	import java.util.concurrent.TimeUnit;
	import java.util.concurrent.Executors;
	import java.util.concurrent.ArrayBlockingQueue;
	import java.util.concurrent.atomic.AtomicBoolean;
	import java.util.concurrent.atomic.AtomicLong;

	public class KuduSpanReceiver extends SpanReceiver {

	private static final Log LOG = LogFactory.getLog(KuduSpanReceiver.class);

	private static final int SHUTDOWN_TIMEOUT = 30;
	private static final int MAX_ERRORS = 10;
	private final BlockingQueue<Span> queue;
	private final AtomicBoolean running = new AtomicBoolean(true);
	private final KuduClientConfiguration clientConf;
	private final int maxSpanBatchSize;
	private final ThreadFactory threadFactory = new ThreadFactory() {
	private final AtomicLong receiverIndex = new AtomicLong(0);

	@Override
	public Thread newThread(Runnable runnable) {
	Thread thread = new Thread(runnable);
	thread.setDaemon(true);
	thread.setName(String.format("kuduSpanReceiver-%d",
	receiverIndex.getAndIncrement()));
	return thread;
	}
	};
	private ExecutorService service;
	private String table;
	private String column_span_id;
	private String column_span;
	private String column_root_span;
	private String column_root_span_start_time;

	public KuduSpanReceiver(HTraceConfiguration conf) {
	this.clientConf =
	new KuduClientConfiguration(conf.get(KuduReceiverConstants.KUDU_MASTER_HOST_KEY,
	KuduReceiverConstants.DEFAULT_KUDU_MASTER_HOST),
	conf.get(KuduReceiverConstants.KUDU_MASTER_PORT_KEY,
	KuduReceiverConstants.DEFAULT_KUDU_MASTER_PORT));
	if (conf.get(KuduReceiverConstants.KUDU_CLIENT_BOSS_COUNT_KEY) != null) {
	this.clientConf.setBossCount(Integer.valueOf(conf.get(KuduReceiverConstants.KUDU_CLIENT_BOSS_COUNT_KEY)));
	}
	if (conf.get(KuduReceiverConstants.KUDU_CLIENT_WORKER_COUNT_KEY) != null) {
	this.clientConf.setWorkerCount(Integer.valueOf(conf.get(KuduReceiverConstants.KUDU_CLIENT_WORKER_COUNT_KEY)));
	}
	if (conf.get(KuduReceiverConstants.KUDU_CLIENT_STATISTICS_ENABLED_KEY) != null) {
	this.clientConf.setIsStatisticsEnabled(Boolean.valueOf(conf.get(KuduReceiverConstants.KUDU_CLIENT_STATISTICS_ENABLED_KEY)));
	}
	if (conf.get(KuduReceiverConstants.KUDU_CLIENT_TIMEOUT_ADMIN_OPERATION_KEY) != null) {
	this.clientConf
	.setAdminOperationTimeout(Long.valueOf(conf.get(KuduReceiverConstants.KUDU_CLIENT_TIMEOUT_ADMIN_OPERATION_KEY)));
	}
	if (conf.get(KuduReceiverConstants.KUDU_CLIENT_TIMEOUT_OPERATION_KEY) != null) {
	this.clientConf
	.setOperationTimeout(Long.valueOf(conf.get(KuduReceiverConstants.KUDU_CLIENT_TIMEOUT_OPERATION_KEY)));
	}
	if (conf.get(KuduReceiverConstants.KUDU_CLIENT_TIMEOUT_SOCKET_READ_KEY) != null) {
	this.clientConf
	.setSocketReadTimeout(Long.valueOf(conf.get(KuduReceiverConstants.KUDU_CLIENT_TIMEOUT_SOCKET_READ_KEY)));
	}
	this.queue = new ArrayBlockingQueue<Span>(conf.getInt(KuduReceiverConstants.SPAN_BLOCKING_QUEUE_SIZE_KEY,
	KuduReceiverConstants.DEFAULT_SPAN_BLOCKING_QUEUE_SIZE));
	this.table = conf.get(KuduReceiverConstants.KUDU_TABLE_KEY, KuduReceiverConstants.DEFAULT_KUDU_TABLE);
	this.column_span_id = conf.get(KuduReceiverConstants.KUDU_COLUMN_SPAN_ID_KEY,
	KuduReceiverConstants.DEFAULT_KUDU_COLUMN_SPAN_ID);
	this.column_span = conf.get(KuduReceiverConstants.KUDU_COLUMN_SPAN_KEY,
	KuduReceiverConstants.DEFAULT_KUDU_COLUMN_SPAN);
	this.column_root_span = conf.get(KuduReceiverConstants.KUDU_COLUMN_ROOT_SPAN_KEY,
	KuduReceiverConstants.DEFAULT_KUDU_COLUMN_ROOT_SPAN);
	this.column_root_span_start_time = conf.get(KuduReceiverConstants.KUDU_COLUMN_ROOT_SPAN_START_TIME_KEY,
	KuduReceiverConstants.DEFAULT_KUDU_COLUMN_ROOT_SPAN_START_TIME);
	this.maxSpanBatchSize = conf.getInt(KuduReceiverConstants.MAX_SPAN_BATCH_SIZE_KEY,
	KuduReceiverConstants.DEFAULT_MAX_SPAN_BATCH_SIZE);
	if (this.service != null) {
	this.service.shutdownNow();
	this.service = null;
	}
	int numThreads = conf.getInt(KuduReceiverConstants.NUM_PARALLEL_THREADS_KEY,
	KuduReceiverConstants.DEFAULT_NUM_PARALLEL_THREADS);
	this.service = Executors.newFixedThreadPool(numThreads, threadFactory);
	for (int i = 0; i < numThreads; i++) {
	this.service.submit(new KuduSpanReceiver.WriteSpanRunnable());
	}
	}

	@Override
	public void close() throws IOException {
	running.set(false);
	service.shutdown();
	try {
	if (!service.awaitTermination(SHUTDOWN_TIMEOUT, TimeUnit.SECONDS)) {
	LOG.error("Timeout " + SHUTDOWN_TIMEOUT + " " + TimeUnit.SECONDS +
	" reached while shutting worker threads which process enqued spans." +
	" Enqueued spans which are left in blocking queue is dropped.");
	}
	} catch (InterruptedException e) {
	LOG.warn("Interrupted exception occured while terminating thread service executor.", e);
	}
	}

	@Override
	public void receiveSpan(Span span) {
	if (running.get()) {
	try {
	this.queue.add(span);
	} catch (IllegalStateException e) {
	LOG.error("Error trying to enqueue span ("
	+ span.getDescription()
	+ ") to the queue. Blocking Queue is currently reached its capacity.");
	}
	}
	}

	private class WriteSpanRunnable implements Runnable {

	private KuduSession session;
	private KuduClient client;

	@Override
	public void run() {
	SpanProtos.Span.Builder sbuilder = SpanProtos.Span.newBuilder();
	SpanProtos.TimelineAnnotation.Builder tlbuilder =
	SpanProtos.TimelineAnnotation.newBuilder();
	List<Span> dequeuedSpans = new ArrayList<Span>(maxSpanBatchSize);
	long errorCount = 0;
	while (running.get() \|\| queue.size() > 0) {
	Span firstSpan = null;
	try {
	firstSpan = queue.poll(1, TimeUnit.SECONDS);
	if (firstSpan != null) {
	dequeuedSpans.add(firstSpan);
	queue.drainTo(dequeuedSpans, maxSpanBatchSize - 1);
	}
	} catch (InterruptedException ie) {
	LOG.error("Interrupted Exception occurred while polling to " +
	"retrieve first span from blocking queue");
	}
	startSession();
	if (dequeuedSpans.isEmpty()) {
	try {
	this.session.flush();
	} catch (java.lang.Exception e) {
	LOG.error("Failed to flush writes to Kudu.");
	closeSession();
	}
	continue;
	}
	try {
	for (Span span : dequeuedSpans) {
	sbuilder.clear()
	.setTraceId(span.getSpanId().getHigh())
	.setStart(span.getStartTimeMillis())
	.setStop(span.getStopTimeMillis())
	.setSpanId(span.getSpanId().getLow())
	.setProcessId(span.getTracerId())
	.setDescription(span.getDescription());

	if (span.getParents().length == 0) {
	sbuilder.setParentId(0);
	} else if (span.getParents().length > 0) {
	sbuilder.setParentId(span.getParents()[0].getLow());
	}
	for (TimelineAnnotation ta : span.getTimelineAnnotations()) {
	sbuilder.addTimeline(tlbuilder.clear()
	.setTime(ta.getTime())
	.setMessage(ta.getMessage())
	.build());
	}
	KuduTable tableRef = client.openTable(table);
	Insert insert = tableRef.newInsert();
	PartialRow row = insert.getRow();
	row.addBinary(column_span_id, Bytes.fromLong(span.getSpanId().getHigh()));
	row.addBinary(column_span, sbuilder.build().toByteArray());
	if (span.getParents().length == 0) {
	row.addBinary(column_root_span_start_time, Bytes.fromLong(span.getStartTimeMillis()));
	row.addBinary(column_root_span, sbuilder.build().toByteArray());
	}
	this.session.apply(insert);
	}
	dequeuedSpans.clear();
	errorCount = 0;
	} catch (Exception e) {
	errorCount += 1;
	if (errorCount < MAX_ERRORS) {
	try {
	queue.addAll(dequeuedSpans);
	} catch (IllegalStateException ex) {
	LOG.error("Exception occured while writing spans kudu datastore. " +
	"Trying to re-enqueue de-queued spans to blocking queue for writing but failed. " +
	"Dropped " + dequeuedSpans.size() + " dequeued span(s) which were due written" +
	"into kudu datastore");
	}
	}
	closeSession();
	try {
	Thread.sleep(500);
	} catch (InterruptedException e1) {
	LOG.error("Interrupted Exception occurred while allowing kudu to re-stabilized");
	}
	}
	}
	closeSession();
	}

	private void closeSession() {
	try {
	if (this.session != null) {
	this.session.close();
	this.session = null;
	}
	} catch (java.lang.Exception e) {
	LOG.warn("Failed to close Kudu session. " + e.getMessage());
	}
	}

	private void startSession() {
	if (this.session == null) {
	if (this.client == null) {
	client = clientConf.buildClient();
	}
	session = client.newSession();
	}
	}
	}

	}