hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/util/Consumers.java - incubator-hugegraph - 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.hugegraph.util;

 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Queue;
 import java.util.concurrent.ArrayBlockingQueue;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;

 import org.apache.hugegraph.HugeException;
 import org.apache.hugegraph.config.CoreOptions;
 import org.apache.hugegraph.task.TaskManager.ContextCallable;
 import org.slf4j.Logger;

 public final class Consumers<V> {

     public static final int THREADS = 4 + CoreOptions.CPUS / 4;
     public static final int QUEUE_WORKER_SIZE = 1000;
     public static final long CONSUMER_WAKE_PERIOD = 1;

     private static final Logger LOG = Log.logger(Consumers.class);

     private final V QUEUE_END = (V) new Object();
     private final ExecutorService executor;
     private final Consumer<V> consumer;
     private final Runnable doneHandle;
     private final Consumer<Throwable> exceptionHandle;
     private final int workers;
     private final List<Future> runningFutures;
     private final int queueSize;
     private final CountDownLatch latch;
     private final BlockingQueue<V> queue;
     private volatile Throwable exception = null;

     public Consumers(ExecutorService executor, Consumer<V> consumer) {
         this(executor, consumer, null);
     }

     public Consumers(ExecutorService executor,
                      Consumer<V> consumer, Runnable doneHandle) {
         this(executor, consumer, doneHandle, QUEUE_WORKER_SIZE);
     }

     public Consumers(ExecutorService executor,
                      Consumer<V> consumer,
                      Runnable doneHandle,
                      int queueSizePerWorker) {
         this(executor, consumer, doneHandle, null, queueSizePerWorker);
     }

     public Consumers(ExecutorService executor,
                      Consumer<V> consumer,
                      Runnable doneHandle,
                      Consumer<Throwable> exceptionHandle,
                      int queueSizePerWorker) {
         this.executor = executor;
         this.consumer = consumer;
         this.doneHandle = doneHandle;
         this.exceptionHandle = exceptionHandle;

         int workers = THREADS;
         if (this.executor instanceof ThreadPoolExecutor) {
             workers = ((ThreadPoolExecutor) this.executor).getCorePoolSize();
         }
         this.workers = workers;

         this.runningFutures = new ArrayList<>(workers);
         this.queueSize = queueSizePerWorker * workers + 1;
         this.latch = new CountDownLatch(workers);
         this.queue = new ArrayBlockingQueue<>(this.queueSize);
     }

     public void start(String name) {
         this.exception = null;
         if (this.executor == null) {
             return;
         }
         LOG.info("Starting {} workers[{}] with queue size {}...",
                  this.workers, name, this.queueSize);
         for (int i = 0; i < this.workers; i++) {
             this.runningFutures.add(
                     this.executor.submit(new ContextCallable<>(this::runAndDone)));
         }
     }

     private Void runAndDone() {
         try {
             this.run();
         } catch (Throwable e) {
             if (e instanceof StopExecution) {
                 this.queue.clear();
                 putQueueEnd();
             } else {
                 // Only the first exception to one thread can be stored
                 this.exception = e;
                 LOG.error("Error when running task", e);
             }
             exceptionHandle(e);
         } finally {
             this.done();
             this.latch.countDown();
         }
         return null;
     }

     private void run() {
         LOG.debug("Start to work...");
         while (this.consume()) {
             // ignore
         }

         LOG.debug("Worker finished");
     }

     private boolean consume() {
         V elem = null;
         while (elem == null) {
             try {
                 elem = this.queue.poll(CONSUMER_WAKE_PERIOD, TimeUnit.MILLISECONDS);
             } catch (InterruptedException e) {
                 // ignore
                 return false;
             }
         }

         if (elem == QUEUE_END) {
             putQueueEnd();
             return false;
         }
         // do job
         this.consumer.accept(elem);
         return true;
     }

     private void exceptionHandle(Throwable e) {
         if (this.exceptionHandle == null) {
             return;
         }

         try {
             this.exceptionHandle.accept(e);
         } catch (Throwable ex) {
             if (this.exception == null) {
                 this.exception = ex;
             } else {
                 LOG.warn("Error while calling exceptionHandle()", ex);
             }
         }
     }

     private void done() {
         if (this.doneHandle == null) {
             return;
         }

         try {
             this.doneHandle.run();
         } catch (Throwable e) {
             if (this.exception == null) {
                 this.exception = e;
             } else {
                 LOG.warn("Error while calling done()", e);
             }
         }
     }

     private Throwable throwException() {
         assert this.exception != null;
         Throwable e = this.exception;
         this.exception = null;
         return e;
     }

     public void provide(V v) throws Throwable {
         if (this.executor == null) {
             assert this.exception == null;
             // do job directly if without thread pool
             this.consumer.accept(v);
         } else if (this.exception != null) {
             throw this.throwException();
         } else {
             try {
                 this.queue.put(v);
             } catch (InterruptedException e) {
                 LOG.warn("Interrupt while queuing QUEUE_END", e);
             }
         }
     }

     private void putQueueEnd() {
         if (this.executor != null) {
             try {
                 this.queue.put(QUEUE_END);
             } catch (InterruptedException e) {
                 LOG.warn("Interrupted while enqueue", e);
             }
         }
     }

     public void await() throws Throwable {
         if (this.executor == null) {
             // call done() directly if without thread pool
             this.done();
         } else {
             try {
                 putQueueEnd();
                 this.latch.await();
             } catch (InterruptedException e) {
                 String error = "Interrupted while waiting for consumers";
                 for (Future f : this.runningFutures) {
                     f.cancel(true);
                 }
                 this.exception = new HugeException(error, e);
                 LOG.warn(error, e);
             }
         }

         if (this.exception != null) {
             throw this.throwException();
         }
     }

     public ExecutorService executor() {
         return this.executor;
     }

     public static void executeOncePerThread(ExecutorService executor,
                                             int totalThreads,
                                             Runnable callback,
                                             long invokeTimeout)
                                             throws InterruptedException {
         // Ensure callback execute at least once for every thread
         final Map<Thread, Integer> threadsTimes = new ConcurrentHashMap<>();
         final List<Callable<Void>> tasks = new ArrayList<>();
         final Callable<Void> task = () -> {
             Thread current = Thread.currentThread();
             threadsTimes.putIfAbsent(current, 0);
             int times = threadsTimes.get(current);
             if (times == 0) {
                 callback.run();
                 // Let other threads run
                 Thread.yield();
             } else {
                 assert times < totalThreads;
                 assert threadsTimes.size() < totalThreads;
                 E.checkState(tasks.size() == totalThreads,
                              "Bad tasks size: %s", tasks.size());
                 // Let another thread run and wait for it
                 executor.submit(tasks.get(0)).get();
             }
             threadsTimes.put(current, ++times);
             return null;
         };

         // NOTE: expect each task thread to perform a close operation
         for (int i = 0; i < totalThreads; i++) {
             tasks.add(task);
         }
         executor.invokeAll(tasks, invokeTimeout, TimeUnit.SECONDS);
     }

     public static ExecutorService newThreadPool(String prefix, int workers) {
         if (workers == 0) {
             return null;
         } else {
             if (workers < 0) {
                 assert workers == -1;
                 workers = Consumers.THREADS;
             } else if (workers > CoreOptions.CPUS * 2) {
                 workers = CoreOptions.CPUS * 2;
             }
             String name = prefix + "-worker-%d";
             return ExecutorUtil.newFixedThreadPool(workers, name);
         }
     }

     public static ExecutorPool newExecutorPool(String prefix, int workers) {
         return new ExecutorPool(prefix, workers);
     }

     public static RuntimeException wrapException(Throwable e) {
         if (e instanceof RuntimeException) {
             throw (RuntimeException) e;
         }
         throw new HugeException("Error when running task: %s",
                                 HugeException.rootCause(e).getMessage(), e);
     }

     public static class ExecutorPool {

         private static final int POOL_CAPACITY = 2 * CoreOptions.CPUS;

         private final String threadNamePrefix;
         private final int executorWorkers;
         private final AtomicInteger count;

         private final Queue<ExecutorService> executors;

         public ExecutorPool(String prefix, int workers) {
             this.threadNamePrefix = prefix;
             this.executorWorkers = workers;
             this.count = new AtomicInteger();
             this.executors = new ArrayBlockingQueue<>(POOL_CAPACITY);
         }

         public synchronized ExecutorService getExecutor() {
             ExecutorService executor = this.executors.poll();
             if (executor == null) {
                 int count = this.count.incrementAndGet();
                 String prefix = this.threadNamePrefix + "-" + count;
                 executor = newThreadPool(prefix, this.executorWorkers);
             }
             return executor;
         }

         public synchronized void returnExecutor(ExecutorService executor) {
             E.checkNotNull(executor, "executor");
             if (!this.executors.offer(executor)) {
                 try {
                     executor.shutdown();
                 } catch (Exception e) {
                     LOG.warn("close ExecutorService with error:", e);
                 }
             }
         }

         public synchronized void destroy() {
             for (ExecutorService executor : this.executors) {
                 try {
                     executor.shutdownNow();
                 } catch (Exception e) {
                     LOG.warn("close ExecutorService with error:", e);
                 }
             }
             this.executors.clear();
         }
     }

     public static class StopExecution extends HugeException {

         private static final long serialVersionUID = -371829356182454517L;

         public StopExecution(String message) {
             super(message);
         }

         public StopExecution(String message, Object... args) {
             super(message, args);
         }
     }
 }
	/*
	* 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.hugegraph.util;

	import java.util.ArrayList;
	import java.util.List;
	import java.util.Map;
	import java.util.Queue;
	import java.util.concurrent.ArrayBlockingQueue;
	import java.util.concurrent.BlockingQueue;
	import java.util.concurrent.Callable;
	import java.util.concurrent.ConcurrentHashMap;
	import java.util.concurrent.CountDownLatch;
	import java.util.concurrent.ExecutorService;
	import java.util.concurrent.Future;
	import java.util.concurrent.ThreadPoolExecutor;
	import java.util.concurrent.TimeUnit;
	import java.util.concurrent.atomic.AtomicInteger;
	import java.util.function.Consumer;

	import org.apache.hugegraph.HugeException;
	import org.apache.hugegraph.config.CoreOptions;
	import org.apache.hugegraph.task.TaskManager.ContextCallable;
	import org.slf4j.Logger;

	public final class Consumers<V> {

	public static final int THREADS = 4 + CoreOptions.CPUS / 4;
	public static final int QUEUE_WORKER_SIZE = 1000;
	public static final long CONSUMER_WAKE_PERIOD = 1;

	private static final Logger LOG = Log.logger(Consumers.class);

	private final V QUEUE_END = (V) new Object();
	private final ExecutorService executor;
	private final Consumer<V> consumer;
	private final Runnable doneHandle;
	private final Consumer<Throwable> exceptionHandle;
	private final int workers;
	private final List<Future> runningFutures;
	private final int queueSize;
	private final CountDownLatch latch;
	private final BlockingQueue<V> queue;
	private volatile Throwable exception = null;

	public Consumers(ExecutorService executor, Consumer<V> consumer) {
	this(executor, consumer, null);
	}

	public Consumers(ExecutorService executor,
	Consumer<V> consumer, Runnable doneHandle) {
	this(executor, consumer, doneHandle, QUEUE_WORKER_SIZE);
	}

	public Consumers(ExecutorService executor,
	Consumer<V> consumer,
	Runnable doneHandle,
	int queueSizePerWorker) {
	this(executor, consumer, doneHandle, null, queueSizePerWorker);
	}

	public Consumers(ExecutorService executor,
	Consumer<V> consumer,
	Runnable doneHandle,
	Consumer<Throwable> exceptionHandle,
	int queueSizePerWorker) {
	this.executor = executor;
	this.consumer = consumer;
	this.doneHandle = doneHandle;
	this.exceptionHandle = exceptionHandle;

	int workers = THREADS;
	if (this.executor instanceof ThreadPoolExecutor) {
	workers = ((ThreadPoolExecutor) this.executor).getCorePoolSize();
	}
	this.workers = workers;

	this.runningFutures = new ArrayList<>(workers);
	this.queueSize = queueSizePerWorker * workers + 1;
	this.latch = new CountDownLatch(workers);
	this.queue = new ArrayBlockingQueue<>(this.queueSize);
	}

	public void start(String name) {
	this.exception = null;
	if (this.executor == null) {
	return;
	}
	LOG.info("Starting {} workers[{}] with queue size {}...",
	this.workers, name, this.queueSize);
	for (int i = 0; i < this.workers; i++) {
	this.runningFutures.add(
	this.executor.submit(new ContextCallable<>(this::runAndDone)));
	}
	}

	private Void runAndDone() {
	try {
	this.run();
	} catch (Throwable e) {
	if (e instanceof StopExecution) {
	this.queue.clear();
	putQueueEnd();
	} else {
	// Only the first exception to one thread can be stored
	this.exception = e;
	LOG.error("Error when running task", e);
	}
	exceptionHandle(e);
	} finally {
	this.done();
	this.latch.countDown();
	}
	return null;
	}

	private void run() {
	LOG.debug("Start to work...");
	while (this.consume()) {
	// ignore
	}

	LOG.debug("Worker finished");
	}

	private boolean consume() {
	V elem = null;
	while (elem == null) {
	try {
	elem = this.queue.poll(CONSUMER_WAKE_PERIOD, TimeUnit.MILLISECONDS);
	} catch (InterruptedException e) {
	// ignore
	return false;
	}
	}

	if (elem == QUEUE_END) {
	putQueueEnd();
	return false;
	}
	// do job
	this.consumer.accept(elem);
	return true;
	}

	private void exceptionHandle(Throwable e) {
	if (this.exceptionHandle == null) {
	return;
	}

	try {
	this.exceptionHandle.accept(e);
	} catch (Throwable ex) {
	if (this.exception == null) {
	this.exception = ex;
	} else {
	LOG.warn("Error while calling exceptionHandle()", ex);
	}
	}
	}

	private void done() {
	if (this.doneHandle == null) {
	return;
	}

	try {
	this.doneHandle.run();
	} catch (Throwable e) {
	if (this.exception == null) {
	this.exception = e;
	} else {
	LOG.warn("Error while calling done()", e);
	}
	}
	}

	private Throwable throwException() {
	assert this.exception != null;
	Throwable e = this.exception;
	this.exception = null;
	return e;
	}

	public void provide(V v) throws Throwable {
	if (this.executor == null) {
	assert this.exception == null;
	// do job directly if without thread pool
	this.consumer.accept(v);
	} else if (this.exception != null) {
	throw this.throwException();
	} else {
	try {
	this.queue.put(v);
	} catch (InterruptedException e) {
	LOG.warn("Interrupt while queuing QUEUE_END", e);
	}
	}
	}

	private void putQueueEnd() {
	if (this.executor != null) {
	try {
	this.queue.put(QUEUE_END);
	} catch (InterruptedException e) {
	LOG.warn("Interrupted while enqueue", e);
	}
	}
	}

	public void await() throws Throwable {
	if (this.executor == null) {
	// call done() directly if without thread pool
	this.done();
	} else {
	try {
	putQueueEnd();
	this.latch.await();
	} catch (InterruptedException e) {
	String error = "Interrupted while waiting for consumers";
	for (Future f : this.runningFutures) {
	f.cancel(true);
	}
	this.exception = new HugeException(error, e);
	LOG.warn(error, e);
	}
	}

	if (this.exception != null) {
	throw this.throwException();
	}
	}

	public ExecutorService executor() {
	return this.executor;
	}

	public static void executeOncePerThread(ExecutorService executor,
	int totalThreads,
	Runnable callback,
	long invokeTimeout)
	throws InterruptedException {
	// Ensure callback execute at least once for every thread
	final Map<Thread, Integer> threadsTimes = new ConcurrentHashMap<>();
	final List<Callable<Void>> tasks = new ArrayList<>();
	final Callable<Void> task = () -> {
	Thread current = Thread.currentThread();
	threadsTimes.putIfAbsent(current, 0);
	int times = threadsTimes.get(current);
	if (times == 0) {
	callback.run();
	// Let other threads run
	Thread.yield();
	} else {
	assert times < totalThreads;
	assert threadsTimes.size() < totalThreads;
	E.checkState(tasks.size() == totalThreads,
	"Bad tasks size: %s", tasks.size());
	// Let another thread run and wait for it
	executor.submit(tasks.get(0)).get();
	}
	threadsTimes.put(current, ++times);
	return null;
	};

	// NOTE: expect each task thread to perform a close operation
	for (int i = 0; i < totalThreads; i++) {
	tasks.add(task);
	}
	executor.invokeAll(tasks, invokeTimeout, TimeUnit.SECONDS);
	}

	public static ExecutorService newThreadPool(String prefix, int workers) {
	if (workers == 0) {
	return null;
	} else {
	if (workers < 0) {
	assert workers == -1;
	workers = Consumers.THREADS;
	} else if (workers > CoreOptions.CPUS * 2) {
	workers = CoreOptions.CPUS * 2;
	}
	String name = prefix + "-worker-%d";
	return ExecutorUtil.newFixedThreadPool(workers, name);
	}
	}

	public static ExecutorPool newExecutorPool(String prefix, int workers) {
	return new ExecutorPool(prefix, workers);
	}

	public static RuntimeException wrapException(Throwable e) {
	if (e instanceof RuntimeException) {
	throw (RuntimeException) e;
	}
	throw new HugeException("Error when running task: %s",
	HugeException.rootCause(e).getMessage(), e);
	}

	public static class ExecutorPool {

	private static final int POOL_CAPACITY = 2 * CoreOptions.CPUS;

	private final String threadNamePrefix;
	private final int executorWorkers;
	private final AtomicInteger count;

	private final Queue<ExecutorService> executors;

	public ExecutorPool(String prefix, int workers) {
	this.threadNamePrefix = prefix;
	this.executorWorkers = workers;
	this.count = new AtomicInteger();
	this.executors = new ArrayBlockingQueue<>(POOL_CAPACITY);
	}

	public synchronized ExecutorService getExecutor() {
	ExecutorService executor = this.executors.poll();
	if (executor == null) {
	int count = this.count.incrementAndGet();
	String prefix = this.threadNamePrefix + "-" + count;
	executor = newThreadPool(prefix, this.executorWorkers);
	}
	return executor;
	}

	public synchronized void returnExecutor(ExecutorService executor) {
	E.checkNotNull(executor, "executor");
	if (!this.executors.offer(executor)) {
	try {
	executor.shutdown();
	} catch (Exception e) {
	LOG.warn("close ExecutorService with error:", e);
	}
	}
	}

	public synchronized void destroy() {
	for (ExecutorService executor : this.executors) {
	try {
	executor.shutdownNow();
	} catch (Exception e) {
	LOG.warn("close ExecutorService with error:", e);
	}
	}
	this.executors.clear();
	}
	}

	public static class StopExecution extends HugeException {

	private static final long serialVersionUID = -371829356182454517L;

	public StopExecution(String message) {
	super(message);
	}

	public StopExecution(String message, Object... args) {
	super(message, args);
	}
	}
	}