/*
* Copyright 2006-2010 the original author or authors.
*
* 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 org.springframework.batch.poller.scheduling;
import java.util.Queue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicReference;
import org.springframework.batch.poller.Poller;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.BeanFactory;
import org.springframework.beans.factory.BeanFactoryAware;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.scheduling.TaskScheduler;
import org.springframework.scheduling.Trigger;
import org.springframework.scheduling.concurrent.ConcurrentTaskScheduler;
import org.springframework.scheduling.support.PeriodicTrigger;
import org.springframework.util.ErrorHandler;
import org.springframework.util.ReflectionUtils;
/**
* A {@link Poller} implementation that uses a {@link TaskScheduler} to poll in
* a background thread.
*
* @author Dave Syer
*
*/
public class TaskSchedulerPoller<T> implements Poller<T>, BeanFactoryAware, InitializingBean {
private static final String TASK_SCHEDULER_BEAN_NAME = "taskScheduler";
private volatile Trigger trigger;
private volatile boolean initialized;
private final Object initializationMonitor = new Object();
private TaskScheduler taskScheduler;
private BeanFactory beanFactory;
public void setTrigger(Trigger trigger) {
this.trigger = trigger;
}
public void setTaskScheduler(TaskScheduler taskScheduler) {
this.taskScheduler = taskScheduler;
}
public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
this.beanFactory = beanFactory;
}
public void afterPropertiesSet() throws Exception {
initialize();
}
private void initialize() {
synchronized (this.initializationMonitor) {
if (this.initialized) {
return;
}
if (this.trigger == null) {
this.trigger = new PeriodicTrigger(100L);
}
if (taskScheduler == null && beanFactory != null) {
taskScheduler = beanFactory.getBean(TASK_SCHEDULER_BEAN_NAME, TaskScheduler.class);
}
}
}
private ScheduledFuture<?> getSchedule(final Callable<T> callable, final Queue<T> queue, final AtomicReference<Throwable> throwable) {
TaskScheduler scheduler = taskScheduler;
if (scheduler == null) {
ConcurrentTaskScheduler concurrentTaskScheduler = new ConcurrentTaskScheduler();
concurrentTaskScheduler.setErrorHandler(new PropagatingErrorHandler());
scheduler = concurrentTaskScheduler;
}
Runnable task = new Runnable() {
public void run() {
if (!queue.isEmpty() || throwable.get() != null) {
return;
}
T result;
try {
result = callable.call();
}
catch (RuntimeException e) {
throwable.set(e);
throw e;
}
catch (Exception e) {
throwable.set(e);
throw new IllegalStateException("Could not obtain result", e);
}
if (result != null) {
queue.add(result);
}
}
};
ScheduledFuture<?> schedule = scheduler.schedule(task, trigger);
return schedule;
}
/**
* @param callback a {@link Callable} to use to retrieve a result
* @return the result, or null if the operation times out
*
* @see Poller#poll(Callable)
*/
public Future<T> poll(Callable<T> callback) throws Exception {
if (!initialized) {
initialize();
}
final BlockingQueue<T> queue = new LinkedBlockingQueue<T>(1);
final AtomicReference<Throwable> throwable = new AtomicReference<Throwable>();
final ScheduledFuture<?> schedule = getSchedule(callback, queue, throwable);
return new Future<T>() {
public boolean cancel(boolean mayInterruptIfRunning) {
return schedule.cancel(mayInterruptIfRunning);
}
public T get() throws InterruptedException, ExecutionException {
try {
T result = queue.take();
if (throwable.get()!=null) {
throw new ExecutionException(throwable.get());
}
return result;
}
finally {
cancelAndMaybeRethrow(schedule);
}
}
public T get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
try {
T result = queue.poll(timeout, unit);
if (throwable.get()!=null) {
throw new ExecutionException(throwable.get());
}
return result;
}
finally {
cancelAndMaybeRethrow(schedule);
}
}
public boolean isCancelled() {
return schedule.isCancelled();
}
public boolean isDone() {
return schedule.isDone() || !queue.isEmpty();
}
private void cancelAndMaybeRethrow(final ScheduledFuture<?> schedule) throws InterruptedException, ExecutionException {
try {
// Just returns null if the task was successful.
schedule.get();
}
catch (ExecutionException e) {
throw e;
}
catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw e;
}
schedule.cancel(true);
}
};
}
/**
* An {@link ErrorHandler} implementation that propagates the throwable as a
* runtime exception.
*/
static class PropagatingErrorHandler implements ErrorHandler {
public void handleError(Throwable t) {
ReflectionUtils.rethrowRuntimeException(t);
}
}
}