apoc.Pools Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of apoc-common Show documentation
Show all versions of apoc-common Show documentation
Data types package for Neo4j Procedures
/*
* Copyright (c) "Neo4j"
* Neo4j Sweden AB [http://neo4j.com]
*
* This file is part of Neo4j.
*
* Licensed 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 apoc;
import apoc.periodic.PeriodicUtils;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.FutureTask;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.function.BiConsumer;
import java.util.stream.Stream;
import org.neo4j.graphdb.GraphDatabaseService;
import org.neo4j.graphdb.Transaction;
import org.neo4j.kernel.api.procedure.GlobalProcedures;
import org.neo4j.kernel.lifecycle.LifecycleAdapter;
import org.neo4j.logging.Log;
import org.neo4j.logging.internal.LogService;
public class Pools extends LifecycleAdapter {
public static final int DEFAULT_SCHEDULED_THREADS = Runtime.getRuntime().availableProcessors() / 4;
public static final int DEFAULT_POOL_THREADS = Runtime.getRuntime().availableProcessors() * 2;
private final Log log;
private final ApocConfig apocConfig;
private ExecutorService singleExecutorService;
private ScheduledExecutorService scheduledExecutorService;
private ExecutorService defaultExecutorService;
private final Map jobList = new ConcurrentHashMap<>();
public Pools(LogService log, GlobalProcedures globalProceduresRegistry, ApocConfig apocConfig) {
this.log = log.getInternalLog(Pools.class);
this.apocConfig = apocConfig;
// expose this config instance via `@Context ApocConfig config`
globalProceduresRegistry.registerComponent((Class) getClass(), ctx -> this, true);
this.log.info("successfully registered Pools for @Context");
}
@Override
public void init() {
int threads =
Math.max(1, apocConfig.getInt(ApocConfig.APOC_CONFIG_JOBS_POOL_NUM_THREADS, DEFAULT_POOL_THREADS));
int queueSize = Math.max(1, apocConfig.getInt(ApocConfig.APOC_CONFIG_JOBS_QUEUE_SIZE, threads * 5));
// ensure we use daemon threads everywhere
ThreadFactory threadFactory = r -> {
Thread t = Executors.defaultThreadFactory().newThread(r);
t.setDaemon(true);
return t;
};
this.singleExecutorService = new ThreadPoolExecutor(
1,
1,
0L,
TimeUnit.SECONDS,
new ArrayBlockingQueue<>(queueSize),
threadFactory,
new CallerBlocksPolicy());
this.defaultExecutorService = new ThreadPoolExecutor(
threads / 2,
threads,
30L,
TimeUnit.SECONDS,
new ArrayBlockingQueue<>(queueSize),
threadFactory,
new CallerBlocksPolicy());
this.scheduledExecutorService = Executors.newScheduledThreadPool(
Math.max(
1,
apocConfig.getInt(
ApocConfig.APOC_CONFIG_JOBS_SCHEDULED_NUM_THREADS, DEFAULT_SCHEDULED_THREADS)),
threadFactory);
scheduledExecutorService.scheduleAtFixedRate(
() -> {
for (Iterator> it =
jobList.entrySet().iterator();
it.hasNext(); ) {
Map.Entry entry = it.next();
if (entry.getValue().isDone() || entry.getValue().isCancelled()) it.remove();
}
},
10,
10,
TimeUnit.SECONDS);
}
@Override
public void shutdown() {
Stream.of(singleExecutorService, defaultExecutorService, scheduledExecutorService)
.forEach(service -> {
try {
service.shutdown();
service.awaitTermination(10, TimeUnit.SECONDS);
} catch (InterruptedException e) {
throw new RuntimeException("Shutdown failed to complete with error: " + e.getMessage());
}
});
}
public ExecutorService getSingleExecutorService() {
return singleExecutorService;
}
public ScheduledExecutorService getScheduledExecutorService() {
return scheduledExecutorService;
}
public ExecutorService getDefaultExecutorService() {
return defaultExecutorService;
}
public Map getJobList() {
return jobList;
}
static class CallerBlocksPolicy implements RejectedExecutionHandler {
@Override
public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
// Submit again by directly injecting the task into the work queue, waiting if necessary, but also
// periodically checking if the pool has been
// shut down.
FutureTask task = new FutureTask<>(r, null);
BlockingQueue queue = executor.getQueue();
while (!executor.isShutdown()) {
try {
if (queue.offer(task, 250, TimeUnit.MILLISECONDS)) {
while (!executor.isShutdown()) {
try {
task.get(250, TimeUnit.MILLISECONDS);
return; // Success!
} catch (TimeoutException ignore) {
// This is fine an expected. We just want to check that the executor hasn't been shut
// down.
}
}
}
} catch (InterruptedException | ExecutionException e) {
throw new RuntimeException(e);
}
}
}
}
public Future processBatch(List batch, GraphDatabaseService db, BiConsumer action) {
return defaultExecutorService.submit(() -> {
try (Transaction tx = db.beginTx()) {
batch.forEach(t -> action.accept(tx, t));
tx.commit();
}
return null;
});
}
public static T force(Future future) throws ExecutionException {
while (true) {
try {
return future.get();
} catch (InterruptedException e) {
Thread.interrupted();
}
}
}
}