TaskSchedulerPoller.java

/*
 * 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);
		}

	}

}