htrace-flume/src/main/java/org/apache/htrace/impl/FlumeSpanReceiver.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 java.io.IOException;
 import java.nio.charset.Charset;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ArrayBlockingQueue;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;

 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.flume.Event;
 import org.apache.flume.FlumeException;
 import org.apache.flume.api.RpcClient;
 import org.apache.flume.api.RpcClientFactory;
 import org.apache.flume.event.EventBuilder;
 import org.apache.htrace.core.HTraceConfiguration;
 import org.apache.htrace.core.Span;
 import org.apache.htrace.core.SpanReceiver;
 import org.apache.htrace.core.TracerId;

 public class FlumeSpanReceiver extends SpanReceiver {
   private static final Log LOG = LogFactory.getLog(FlumeSpanReceiver.class);

   public static final String NUM_THREADS_KEY = "htrace.flume.num-threads";
   public static final int DEFAULT_NUM_THREADS = 1;
   public static final String FLUME_HOSTNAME_KEY = "htrace.flume.hostname";
   public static final String DEFAULT_FLUME_HOSTNAME = "localhost";
   public static final String FLUME_PORT_KEY = "htrace.flume.port";
   public static final String FLUME_BATCHSIZE_KEY = "htrace.flume.batchsize";
   public static final int DEFAULT_FLUME_BATCHSIZE = 100;

   /**
    * How long this receiver will try and wait for all threads to shutdown.
    */
   private static final int SHUTDOWN_TIMEOUT = 30;

   /**
    * How many errors in a row before we start dropping traces on the floor.
    */
   private static final int MAX_ERRORS = 10;

   /**
    * The queue that will get all HTrace spans that are to be sent.
    */
   private final BlockingQueue<Span> queue;

   /**
    * Boolean used to signal that the threads should end.
    */
   private final AtomicBoolean running = new AtomicBoolean(true);

   /**
    * The thread factory used to create new ExecutorService.
    * <p/>
    * This will be the same factory for the lifetime of this object so that
    * no thread names will ever be duplicated.
    */
   private final ThreadFactory tf;

   ////////////////////
   /// Variables that will change on each call to configure()
   ///////////////////
   private ExecutorService service;
   private int maxSpanBatchSize;
   private String flumeHostName;
   private int flumePort;

   public FlumeSpanReceiver(HTraceConfiguration conf) {
     this.queue = new ArrayBlockingQueue<Span>(1000);
     this.tf = new SimpleThreadFactory();
     configure(conf);
   }

   private class SimpleThreadFactory implements ThreadFactory {
     final AtomicLong count = new AtomicLong(0);
     @Override
     public Thread newThread(Runnable arg0) {
       String name = String.format("flumeSpanReceiver-%d", count.getAndIncrement());
       Thread t = new Thread(arg0, name);
       t.setDaemon(true);
       return t;
     }
   }

   private void configure (HTraceConfiguration conf) {

     // Read configuration
     int numThreads = conf.getInt(NUM_THREADS_KEY, DEFAULT_NUM_THREADS);
     this.flumeHostName = conf.get(FLUME_HOSTNAME_KEY, DEFAULT_FLUME_HOSTNAME);
     this.flumePort = conf.getInt(FLUME_PORT_KEY, 0);
     if (this.flumePort == 0) {
       throw new IllegalArgumentException(FLUME_PORT_KEY + " is required in configuration.");
     }
     this.maxSpanBatchSize = conf.getInt(FLUME_BATCHSIZE_KEY, DEFAULT_FLUME_BATCHSIZE);

     // Initialize executors
     // If there are already threads running tear them down.
     if (this.service != null) {
       this.service.shutdownNow();
       this.service = null;
     }
     this.service = Executors.newFixedThreadPool(numThreads, tf);
     for (int i = 0; i < numThreads; i++) {
       this.service.submit(new WriteSpanRunnable());
     }
   }

   private class WriteSpanRunnable implements Runnable {
     private RpcClient flumeClient = null;

     /**
      * This runnable sends a HTrace span to the Flume.
      */
     @Override
     public void run() {
       List<Span> dequeuedSpans = new ArrayList<Span>(maxSpanBatchSize);
       long errorCount = 0;

       while (running.get() || queue.size() > 0) {
         Span firstSpan = null;
         try {
           // Block for up to a second. to try and get a span.
           // We only block for a little bit in order to notice
           // if the running value has changed
           firstSpan = queue.poll(1, TimeUnit.SECONDS);

           // If the poll was successful then it's possible that there
           // will be other spans to get. Try and get them.
           if (firstSpan != null) {
             // Add the first one that we got
             dequeuedSpans.add(firstSpan);
             // Try and get up to 100 queues
             queue.drainTo(dequeuedSpans, maxSpanBatchSize - 1);
           }
         } catch (InterruptedException ie) {
           // Ignored.
         }

         startClient();
         if (dequeuedSpans.isEmpty()) {
           continue;
         }

         try {
           List<Event> events = new ArrayList<Event>(dequeuedSpans.size());
           for (Span span : dequeuedSpans) {
             // Headers allow Flume to filter
             Map<String, String> headers = new HashMap<String, String>();
             headers.put("SpanId",       span.toString());
             headers.put("TracerId",     span.getTracerId());
             headers.put("Description",  span.getDescription());

             String body = span.toJson();

             Event evt = EventBuilder.withBody(body, Charset.forName("UTF-8"), headers);
             events.add(evt);
           }
           flumeClient.appendBatch(events);

           // clear the list for the next time through.
           dequeuedSpans.clear();
           // reset the error counter.
           errorCount = 0;
         } catch (Exception e) {
           errorCount += 1;
           // If there have been ten errors in a row start dropping things.
           if (errorCount < MAX_ERRORS) {
             try {
               queue.addAll(dequeuedSpans);
             } catch (IllegalStateException ex) {
               LOG.error("Drop " + dequeuedSpans.size() +
                         " span(s) because writing to HBase failed.");
             }
           }
           closeClient();
           try {
             // Since there was an error sleep just a little bit to try and allow the
             // HBase some time to recover.
             Thread.sleep(500);
           } catch (InterruptedException e1) {
             // Ignored
           }
         }
       }
       closeClient();
     }

     /**
      * Close Flume RPC client
      */
     private void closeClient() {
       if (flumeClient != null) {
         try {
           flumeClient.close();
         } catch (FlumeException ex) {
           LOG.warn("Error while trying to close Flume Rpc Client.", ex);
         } finally {
           flumeClient = null;
         }
       }
     }

     /**
      * Create / reconnect Flume RPC client
      */
     private void startClient() {
       // If current client is inactive, close it
       if (flumeClient != null && !flumeClient.isActive()) {
         flumeClient.close();
         flumeClient = null;
       }
       // Create client if needed
       if (flumeClient == null) {
         try {
           flumeClient = RpcClientFactory.getDefaultInstance(flumeHostName, flumePort, maxSpanBatchSize);
         } catch (FlumeException e) {
           LOG.warn("Failed to create Flume RPC Client. " + e.getMessage());
         }
       }
     }
   }

   /**
    * Close the receiver and shut down the thread pool.
    *
    * @throws IOException On error.
    */
   @Override
   public void close() throws IOException {
     running.set(false);
     service.shutdown();
     try {
       if (!service.awaitTermination(SHUTDOWN_TIMEOUT, TimeUnit.SECONDS)) {
         LOG.error("Was not able to process all remaining spans upon closing in: " +
             SHUTDOWN_TIMEOUT + " " + TimeUnit.SECONDS +
             ". Left Spans could be dropped.");
        }
     } catch (InterruptedException e1) {
       LOG.warn("Thread interrupted when terminating executor.", e1);
     }
   }

   @Override
   public void receiveSpan(Span span) {
     if (running.get()) {
       try {
         this.queue.add(span);
       } catch (IllegalStateException e) {
         LOG.error("Error trying to append span (" +
             span.getDescription() +
             ") to the queue. Blocking Queue was full.");
       }
     }
   }
 }
	/*
	* 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 java.io.IOException;
	import java.nio.charset.Charset;
	import java.util.ArrayList;
	import java.util.HashMap;
	import java.util.List;
	import java.util.Map;
	import java.util.concurrent.ArrayBlockingQueue;
	import java.util.concurrent.BlockingQueue;
	import java.util.concurrent.ExecutorService;
	import java.util.concurrent.Executors;
	import java.util.concurrent.ThreadFactory;
	import java.util.concurrent.TimeUnit;
	import java.util.concurrent.atomic.AtomicBoolean;
	import java.util.concurrent.atomic.AtomicLong;

	import org.apache.commons.logging.Log;
	import org.apache.commons.logging.LogFactory;
	import org.apache.flume.Event;
	import org.apache.flume.FlumeException;
	import org.apache.flume.api.RpcClient;
	import org.apache.flume.api.RpcClientFactory;
	import org.apache.flume.event.EventBuilder;
	import org.apache.htrace.core.HTraceConfiguration;
	import org.apache.htrace.core.Span;
	import org.apache.htrace.core.SpanReceiver;
	import org.apache.htrace.core.TracerId;

	public class FlumeSpanReceiver extends SpanReceiver {
	private static final Log LOG = LogFactory.getLog(FlumeSpanReceiver.class);

	public static final String NUM_THREADS_KEY = "htrace.flume.num-threads";
	public static final int DEFAULT_NUM_THREADS = 1;
	public static final String FLUME_HOSTNAME_KEY = "htrace.flume.hostname";
	public static final String DEFAULT_FLUME_HOSTNAME = "localhost";
	public static final String FLUME_PORT_KEY = "htrace.flume.port";
	public static final String FLUME_BATCHSIZE_KEY = "htrace.flume.batchsize";
	public static final int DEFAULT_FLUME_BATCHSIZE = 100;

	/**
	* How long this receiver will try and wait for all threads to shutdown.
	*/
	private static final int SHUTDOWN_TIMEOUT = 30;

	/**
	* How many errors in a row before we start dropping traces on the floor.
	*/
	private static final int MAX_ERRORS = 10;

	/**
	* The queue that will get all HTrace spans that are to be sent.
	*/
	private final BlockingQueue<Span> queue;

	/**
	* Boolean used to signal that the threads should end.
	*/
	private final AtomicBoolean running = new AtomicBoolean(true);

	/**
	* The thread factory used to create new ExecutorService.
	* <p/>
	* This will be the same factory for the lifetime of this object so that
	* no thread names will ever be duplicated.
	*/
	private final ThreadFactory tf;

	////////////////////
	/// Variables that will change on each call to configure()
	///////////////////
	private ExecutorService service;
	private int maxSpanBatchSize;
	private String flumeHostName;
	private int flumePort;

	public FlumeSpanReceiver(HTraceConfiguration conf) {
	this.queue = new ArrayBlockingQueue<Span>(1000);
	this.tf = new SimpleThreadFactory();
	configure(conf);
	}

	private class SimpleThreadFactory implements ThreadFactory {
	final AtomicLong count = new AtomicLong(0);
	@Override
	public Thread newThread(Runnable arg0) {
	String name = String.format("flumeSpanReceiver-%d", count.getAndIncrement());
	Thread t = new Thread(arg0, name);
	t.setDaemon(true);
	return t;
	}
	}

	private void configure (HTraceConfiguration conf) {

	// Read configuration
	int numThreads = conf.getInt(NUM_THREADS_KEY, DEFAULT_NUM_THREADS);
	this.flumeHostName = conf.get(FLUME_HOSTNAME_KEY, DEFAULT_FLUME_HOSTNAME);
	this.flumePort = conf.getInt(FLUME_PORT_KEY, 0);
	if (this.flumePort == 0) {
	throw new IllegalArgumentException(FLUME_PORT_KEY + " is required in configuration.");
	}
	this.maxSpanBatchSize = conf.getInt(FLUME_BATCHSIZE_KEY, DEFAULT_FLUME_BATCHSIZE);

	// Initialize executors
	// If there are already threads running tear them down.
	if (this.service != null) {
	this.service.shutdownNow();
	this.service = null;
	}
	this.service = Executors.newFixedThreadPool(numThreads, tf);
	for (int i = 0; i < numThreads; i++) {
	this.service.submit(new WriteSpanRunnable());
	}
	}

	private class WriteSpanRunnable implements Runnable {
	private RpcClient flumeClient = null;

	/**
	* This runnable sends a HTrace span to the Flume.
	*/
	@Override
	public void run() {
	List<Span> dequeuedSpans = new ArrayList<Span>(maxSpanBatchSize);
	long errorCount = 0;

	while (running.get() \|\| queue.size() > 0) {
	Span firstSpan = null;
	try {
	// Block for up to a second. to try and get a span.
	// We only block for a little bit in order to notice
	// if the running value has changed
	firstSpan = queue.poll(1, TimeUnit.SECONDS);

	// If the poll was successful then it's possible that there
	// will be other spans to get. Try and get them.
	if (firstSpan != null) {
	// Add the first one that we got
	dequeuedSpans.add(firstSpan);
	// Try and get up to 100 queues
	queue.drainTo(dequeuedSpans, maxSpanBatchSize - 1);
	}
	} catch (InterruptedException ie) {
	// Ignored.
	}

	startClient();
	if (dequeuedSpans.isEmpty()) {
	continue;
	}

	try {
	List<Event> events = new ArrayList<Event>(dequeuedSpans.size());
	for (Span span : dequeuedSpans) {
	// Headers allow Flume to filter
	Map<String, String> headers = new HashMap<String, String>();
	headers.put("SpanId", span.toString());
	headers.put("TracerId", span.getTracerId());
	headers.put("Description", span.getDescription());

	String body = span.toJson();

	Event evt = EventBuilder.withBody(body, Charset.forName("UTF-8"), headers);
	events.add(evt);
	}
	flumeClient.appendBatch(events);

	// clear the list for the next time through.
	dequeuedSpans.clear();
	// reset the error counter.
	errorCount = 0;
	} catch (Exception e) {
	errorCount += 1;
	// If there have been ten errors in a row start dropping things.
	if (errorCount < MAX_ERRORS) {
	try {
	queue.addAll(dequeuedSpans);
	} catch (IllegalStateException ex) {
	LOG.error("Drop " + dequeuedSpans.size() +
	" span(s) because writing to HBase failed.");
	}
	}
	closeClient();
	try {
	// Since there was an error sleep just a little bit to try and allow the
	// HBase some time to recover.
	Thread.sleep(500);
	} catch (InterruptedException e1) {
	// Ignored
	}
	}
	}
	closeClient();
	}

	/**
	* Close Flume RPC client
	*/
	private void closeClient() {
	if (flumeClient != null) {
	try {
	flumeClient.close();
	} catch (FlumeException ex) {
	LOG.warn("Error while trying to close Flume Rpc Client.", ex);
	} finally {
	flumeClient = null;
	}
	}
	}

	/**
	* Create / reconnect Flume RPC client
	*/
	private void startClient() {
	// If current client is inactive, close it
	if (flumeClient != null && !flumeClient.isActive()) {
	flumeClient.close();
	flumeClient = null;
	}
	// Create client if needed
	if (flumeClient == null) {
	try {
	flumeClient = RpcClientFactory.getDefaultInstance(flumeHostName, flumePort, maxSpanBatchSize);
	} catch (FlumeException e) {
	LOG.warn("Failed to create Flume RPC Client. " + e.getMessage());
	}
	}
	}
	}

	/**
	* Close the receiver and shut down the thread pool.
	*
	* @throws IOException On error.
	*/
	@Override
	public void close() throws IOException {
	running.set(false);
	service.shutdown();
	try {
	if (!service.awaitTermination(SHUTDOWN_TIMEOUT, TimeUnit.SECONDS)) {
	LOG.error("Was not able to process all remaining spans upon closing in: " +
	SHUTDOWN_TIMEOUT + " " + TimeUnit.SECONDS +
	". Left Spans could be dropped.");
	}
	} catch (InterruptedException e1) {
	LOG.warn("Thread interrupted when terminating executor.", e1);
	}
	}

	@Override
	public void receiveSpan(Span span) {
	if (running.get()) {
	try {
	this.queue.add(span);
	} catch (IllegalStateException e) {
	LOG.error("Error trying to append span (" +
	span.getDescription() +
	") to the queue. Blocking Queue was full.");
	}
	}
	}
	}