/**
* 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.activemq.broker.scheduler.memory;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Date;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Timer;
import java.util.TimerTask;
import java.util.TreeMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import javax.jms.MessageFormatException;
import org.apache.activemq.broker.scheduler.CronParser;
import org.apache.activemq.broker.scheduler.Job;
import org.apache.activemq.broker.scheduler.JobListener;
import org.apache.activemq.broker.scheduler.JobScheduler;
import org.apache.activemq.broker.scheduler.JobSupport;
import org.apache.activemq.util.ByteSequence;
import org.apache.activemq.util.IdGenerator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* Implements an in-memory JobScheduler instance.
*/
public class InMemoryJobScheduler implements JobScheduler {
private static final Logger LOG = LoggerFactory.getLogger(InMemoryJobScheduler.class);
private static final IdGenerator ID_GENERATOR = new IdGenerator();
private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
private final String name;
private final TreeMap<Long, ScheduledTask> jobs = new TreeMap<>();
private final AtomicBoolean started = new AtomicBoolean(false);
private final AtomicBoolean dispatchEnabled = new AtomicBoolean(false);
private final List<JobListener> jobListeners = new CopyOnWriteArrayList<>();
private final Timer timer = new Timer();
public InMemoryJobScheduler(String name) {
this.name = name;
}
@Override
public String getName() throws Exception {
return name;
}
public void start() throws Exception {
if (started.compareAndSet(false, true)) {
startDispatching();
LOG.trace("JobScheduler[{}] started", name);
}
}
public void stop() throws Exception {
if (started.compareAndSet(true, false)) {
stopDispatching();
timer.cancel();
jobs.clear();
LOG.trace("JobScheduler[{}] stopped", name);
}
}
public boolean isStarted() {
return started.get();
}
public boolean isDispatchEnabled() {
return dispatchEnabled.get();
}
@Override
public void startDispatching() throws Exception {
dispatchEnabled.set(true);
}
@Override
public void stopDispatching() throws Exception {
dispatchEnabled.set(false);
}
@Override
public void addListener(JobListener listener) throws Exception {
this.jobListeners.add(listener);
}
@Override
public void removeListener(JobListener listener) throws Exception {
this.jobListeners.remove(listener);
}
@Override
public void schedule(String jobId, ByteSequence payload, long delay) throws Exception {
doSchedule(jobId, payload, "", 0, delay, 0);
}
@Override
public void schedule(String jobId, ByteSequence payload, String cronEntry) throws Exception {
doSchedule(jobId, payload, cronEntry, 0, 0, 0);
}
@Override
public void schedule(String jobId, ByteSequence payload, String cronEntry, long delay, long period, int repeat) throws Exception {
doSchedule(jobId, payload, cronEntry, delay, period, repeat);
}
@Override
public void remove(long time) throws Exception {
doRemoveRange(time, time);
}
@Override
public void remove(String jobId) throws Exception {
doRemoveJob(jobId);
}
@Override
public void removeAllJobs() throws Exception {
doRemoveRange(0, Long.MAX_VALUE);
}
@Override
public void removeAllJobs(long start, long finish) throws Exception {
doRemoveRange(start, finish);
}
@Override
public long getNextScheduleTime() throws Exception {
long nextExecutionTime = -1L;
lock.readLock().lock();
try {
if (!jobs.isEmpty()) {
nextExecutionTime = jobs.entrySet().iterator().next().getKey();
}
} finally {
lock.readLock().unlock();
}
return nextExecutionTime;
}
@Override
public List<Job> getNextScheduleJobs() throws Exception {
List<Job> result = new ArrayList<>();
lock.readLock().lock();
try {
if (!jobs.isEmpty()) {
result.addAll(jobs.entrySet().iterator().next().getValue().getAllJobs());
}
} finally {
lock.readLock().unlock();
}
return result;
}
@Override
public List<Job> getAllJobs() throws Exception {
final List<Job> result = new ArrayList<>();
this.lock.readLock().lock();
try {
for (Map.Entry<Long, ScheduledTask> entry : jobs.entrySet()) {
result.addAll(entry.getValue().getAllJobs());
}
} finally {
this.lock.readLock().unlock();
}
return result;
}
@Override
public List<Job> getAllJobs(long start, long finish) throws Exception {
final List<Job> result = new ArrayList<>();
this.lock.readLock().lock();
try {
for (Map.Entry<Long, ScheduledTask> entry : jobs.entrySet()) {
long jobTime = entry.getKey();
if (start <= jobTime && jobTime <= finish) {
result.addAll(entry.getValue().getAllJobs());
}
}
} finally {
this.lock.readLock().unlock();
}
return result;
}
@Override
public int hashCode() {
return name.hashCode();
}
@Override
public String toString() {
return "JobScheduler: " + name;
}
private void doSchedule(final String jobId, final ByteSequence payload, final String cronEntry, long delay, long period, int repeat) throws IOException {
long startTime = System.currentTimeMillis();
long executionTime = 0;
// round startTime - so we can schedule more jobs at the same time
startTime = ((startTime + 500) / 500) * 500;
if (cronEntry != null && cronEntry.length() > 0) {
try {
executionTime = CronParser.getNextScheduledTime(cronEntry, startTime);
} catch (MessageFormatException e) {
throw new IOException(e.getMessage());
}
}
if (executionTime == 0) {
// start time not set by CRON - so it it to the current time
executionTime = startTime;
}
if (delay > 0) {
executionTime += delay;
} else {
executionTime += period;
}
InMemoryJob newJob = new InMemoryJob(jobId);
newJob.setStart(startTime);
newJob.setCronEntry(cronEntry);
newJob.setDelay(delay);
newJob.setPeriod(period);
newJob.setRepeat(repeat);
newJob.setNextTime(executionTime);
newJob.setPayload(payload.getData());
LOG.trace("JobScheduler adding job[{}] to fire at: {}", jobId, JobSupport.getDateTime(executionTime));
lock.writeLock().lock();
try {
ScheduledTask task = jobs.get(executionTime);
if (task == null) {
task = new ScheduledTask(executionTime);
task.add(newJob);
jobs.put(task.getExecutionTime(), task);
timer.schedule(task, new Date(newJob.getNextTime()));
} else {
task.add(newJob);
}
} finally {
lock.writeLock().unlock();
}
}
private void doReschedule(InMemoryJob job, long nextExecutionTime) {
job.setNextTime(nextExecutionTime);
job.incrementExecutionCount();
if (!job.isCron()) {
job.decrementRepeatCount();
}
LOG.trace("JobScheduler rescheduling job[{}] to fire at: {}", job.getJobId(), JobSupport.getDateTime(nextExecutionTime));
lock.writeLock().lock();
try {
ScheduledTask task = jobs.get(nextExecutionTime);
if (task == null) {
task = new ScheduledTask(nextExecutionTime);
task.add(job);
jobs.put(task.getExecutionTime(), task);
timer.schedule(task, new Date(task.getExecutionTime()));
} else {
task.add(job);
}
} finally {
lock.writeLock().unlock();
}
}
private void doRemoveJob(String jobId) throws IOException {
this.lock.writeLock().lock();
try {
Iterator<Map.Entry<Long, ScheduledTask>> scheduled = jobs.entrySet().iterator();
while (scheduled.hasNext()) {
Map.Entry<Long, ScheduledTask> entry = scheduled.next();
ScheduledTask task = entry.getValue();
if (task.remove(jobId)) {
LOG.trace("JobScheduler removing job[{}]", jobId);
if (task.isEmpty()) {
task.cancel();
scheduled.remove();
}
return;
}
}
} finally {
this.lock.writeLock().unlock();
}
}
private void doRemoveRange(long start, long end) throws IOException {
this.lock.writeLock().lock();
try {
Iterator<Map.Entry<Long, ScheduledTask>> scheduled = jobs.entrySet().iterator();
while (scheduled.hasNext()) {
Map.Entry<Long, ScheduledTask> entry = scheduled.next();
long executionTime = entry.getKey();
if (start <= executionTime && executionTime <= end) {
ScheduledTask task = entry.getValue();
task.cancel();
scheduled.remove();
}
// Don't look beyond the end range.
if (end < executionTime) {
break;
}
}
} finally {
this.lock.writeLock().unlock();
}
}
private boolean canDispatch() {
return isStarted() && isDispatchEnabled();
}
private long calculateNextExecutionTime(InMemoryJob job, long currentTime, int repeat) throws MessageFormatException {
long result = currentTime;
String cron = job.getCronEntry();
if (cron != null && cron.length() > 0) {
result = CronParser.getNextScheduledTime(cron, result);
} else if (job.getRepeat() != 0) {
result += job.getPeriod();
}
return result;
}
private void dispatch(InMemoryJob job) throws IllegalStateException, IOException {
if (canDispatch()) {
LOG.debug("Firing: {}", job);
for (JobListener l : jobListeners) {
l.scheduledJob(job.getJobId(), new ByteSequence(job.getPayload()));
}
}
}
/*
* A TimerTask instance that can aggregate the execution of a number
* scheduled Jobs and handle rescheduling the jobs that require it.
*/
private class ScheduledTask extends TimerTask {
private final Map<String, InMemoryJob> jobs = new TreeMap<>();
private final long executionTime;
public ScheduledTask(long executionTime) {
this.executionTime = executionTime;
}
public long getExecutionTime() {
return executionTime;
}
/**
* @return a Collection containing all the managed jobs for this task.
*/
public Collection<InMemoryJob> getAllJobs() {
return new ArrayList<>(jobs.values());
}
/**
* @return true if the internal list of jobs has become empty.
*/
public boolean isEmpty() {
return jobs.isEmpty();
}
/**
* Adds the job to the internal list of scheduled Jobs managed by this task.
*
* @param newJob
* the new job to add to the list of Jobs.
*/
public void add(InMemoryJob newJob) {
this.jobs.put(newJob.getJobId(), newJob);
}
/**
* Removes the job from the internal list of scheduled Jobs managed by this task.
*
* @param jobId
* the job ID to remove from the list of Jobs.
*
* @return true if the job was removed from the list of managed jobs.
*/
public boolean remove(String jobId) {
return jobs.remove(jobId) != null;
}
@Override
public void run() {
if (!isStarted()) {
return;
}
try {
long currentTime = System.currentTimeMillis();
lock.writeLock().lock();
try {
// Remove this entry as it will now fire any scheduled jobs, if new
// jobs or rescheduled jobs land in the same time slot we want them
// to go into a new ScheduledTask in the Timer instance.
InMemoryJobScheduler.this.jobs.remove(executionTime);
} finally {
lock.writeLock().unlock();
}
long nextExecutionTime = 0;
for (InMemoryJob job : jobs.values()) {
if (!isStarted()) {
break;
}
int repeat = job.getRepeat();
nextExecutionTime = calculateNextExecutionTime(job, currentTime, repeat);
if (!job.isCron()) {
dispatch(job);
if (repeat != 0) {
// Reschedule for the next time, the scheduler will take care of
// updating the repeat counter on the update.
doReschedule(job, nextExecutionTime);
}
} else {
if (repeat == 0) {
// This is a non-repeating Cron entry so we can fire and forget it.
dispatch(job);
}
if (nextExecutionTime > currentTime) {
// Reschedule the cron job as a new event, if the cron entry signals
// a repeat then it will be stored separately and fired as a normal
// event with decrementing repeat.
doReschedule(job, nextExecutionTime);
if (repeat != 0) {
// we have a separate schedule to run at this time
// so the cron job is used to set of a separate schedule
// hence we won't fire the original cron job to the
// listeners but we do need to start a separate schedule
String jobId = ID_GENERATOR.generateId();
ByteSequence payload = new ByteSequence(job.getPayload());
schedule(jobId, payload, "", job.getDelay(), job.getPeriod(), job.getRepeat());
}
}
}
}
} catch (Throwable e) {
LOG.error("Error while processing scheduled job(s).", e);
}
}
}
}