• Home
• com.github.sonus21
• rqueue
• 1.3.1-RELEASE
• source code
• MessageScheduler.java

×

Project download

Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance.
Project price only 1 $

You can buy this project and download/modify it how often you want.

com.github.sonus21.rqueue.core.MessageScheduler Maven / Gradle / Ivy

/*
 * Copyright 2020 Sonu Kumar
 *
 * 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
 *
 *       https://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 com.github.sonus21.rqueue.core;

import static java.lang.Math.max;
import static java.lang.Math.min;

import com.github.sonus21.rqueue.core.RedisScriptFactory.ScriptType;
import com.github.sonus21.rqueue.listener.ConsumerQueueDetail;
import com.github.sonus21.rqueue.utils.QueueInitializationEvent;
import com.github.sonus21.rqueue.utils.SchedulerFactory;
import java.time.Instant;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.concurrent.CancellationException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import org.slf4j.Logger;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationListener;
import org.springframework.data.redis.RedisSystemException;
import org.springframework.data.redis.connection.Message;
import org.springframework.data.redis.connection.MessageListener;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultScriptExecutor;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.data.redis.listener.ChannelTopic;
import org.springframework.data.redis.listener.RedisMessageListenerContainer;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;

public abstract class MessageScheduler
    implements DisposableBean, ApplicationListener {
  private static final long DEFAULT_SCRIPT_EXECUTION_TIME = 5000L;
  private static final long MIN_DELAY = 10L;
  private static final int MAX_MESSAGE = 100;
  private static final long TASK_ALIVE_TIME = -30 * 1000L;

  private final int poolSize;
  private boolean scheduleTaskAtStartup;
  private RedisScript redisScript;
  private MessageSchedulerListener messageSchedulerListener;
  private RedisTemplate redisTemplate;
  private DefaultScriptExecutor defaultScriptExecutor;
  private Map queueRunningState;
  private Map queueNameToScheduledTask;
  private Map channelNameToQueueName;
  private Map queueNameToZsetName;
  private Map queueNameToLastMessageSeenTime;
  private ThreadPoolTaskScheduler scheduler;
  @Autowired private RedisMessageListenerContainer redisMessageListenerContainer;

  public MessageScheduler(
      RedisTemplate redisTemplate, int poolSize, boolean scheduleTaskAtStartup) {
    this.poolSize = poolSize;
    this.scheduleTaskAtStartup = scheduleTaskAtStartup;
    this.redisTemplate = redisTemplate;
  }

  protected abstract Logger getLogger();

  protected abstract long getNextScheduleTime(long currentTime, Long value);

  protected abstract String getChannelName(String queueName);

  protected abstract String getZsetName(String queueName);

  protected abstract String getThreadNamePrefix();

  protected abstract boolean isQueueValid(ConsumerQueueDetail queueDetail);

  private void doStart() {
    for (String queueName : queueRunningState.keySet()) {
      startQueue(queueName);
    }
  }

  private void startQueue(String queueName) {
    if (Boolean.TRUE.equals(queueRunningState.get(queueName))) {
      return;
    }
    queueRunningState.put(queueName, true);
    if (isScheduleTaskAtStartup()) {
      long scheduleAt = System.currentTimeMillis() + MIN_DELAY;
      schedule(queueName, getZsetName(queueName), scheduleAt, false);
    }
    redisMessageListenerContainer.addMessageListener(
        messageSchedulerListener, new ChannelTopic(getChannelName(queueName)));
    channelNameToQueueName.put(getChannelName(queueName), queueName);
    queueNameToZsetName.put(queueName, getZsetName(queueName));
  }

  private void doStop() {
    if (CollectionUtils.isEmpty(queueRunningState)) {
      return;
    }
    for (Map.Entry runningStateByQueue : queueRunningState.entrySet()) {
      if (Boolean.TRUE.equals(runningStateByQueue.getValue())) {
        stopQueue(runningStateByQueue.getKey());
      }
    }
    waitForRunningQueuesToStop();
    queueNameToScheduledTask.clear();
  }

  private void waitForRunningQueuesToStop() {
    for (Map.Entry runningState : queueRunningState.entrySet()) {
      String queueName = runningState.getKey();
      ScheduledTaskDetail scheduledTaskDetail = queueNameToScheduledTask.get(queueName);
      if (scheduledTaskDetail != null) {
        Future future = scheduledTaskDetail.getFuture();
        boolean completedOrCancelled = future.isCancelled() || future.isDone();
        if (!completedOrCancelled) {
          scheduledTaskDetail.getFuture().cancel(true);
        }
      }
    }
  }

  private void stopQueue(String queueName) {
    Assert.isTrue(
        queueRunningState.containsKey(queueName),
        "Queue with name '" + queueName + "' does not exist");
    queueRunningState.put(queueName, false);
  }

  private boolean isScheduleTaskAtStartup() {
    return scheduleTaskAtStartup;
  }

  @Override
  public void destroy() throws Exception {
    doStop();
    if (scheduler != null) {
      scheduler.destroy();
    }
  }

  private void createScheduler(int queueCount) {
    if (queueCount == 0) {
      return;
    }
    scheduler =
        SchedulerFactory.createThreadPoolTaskScheduler(
            min(poolSize, queueCount), getThreadNamePrefix(), 60);
  }

  private boolean isQueueActive(String queueName) {
    Boolean val = queueRunningState.get(queueName);
    if (val == null) {
      return false;
    }
    return val;
  }

  protected synchronized void schedule(
      String queueName, String zsetName, Long startTime, boolean forceSchedule) {
    boolean isQueueActive = isQueueActive(queueName);
    if (!isQueueActive || scheduler == null) {
      return;
    }
    long lastSeenTime = queueNameToLastMessageSeenTime.getOrDefault(queueName, 0L);
    long currentTime = System.currentTimeMillis();
    // ignore too frequents events
    if (!forceSchedule && currentTime - lastSeenTime < MIN_DELAY) {
      return;
    }
    queueNameToLastMessageSeenTime.put(queueName, currentTime);

    ScheduledTaskDetail scheduledTaskDetail = queueNameToScheduledTask.get(queueName);
    if (scheduledTaskDetail == null || forceSchedule) {
      long requiredDelay = max(1, startTime - currentTime);
      long taskStartTime = startTime;
      MessageMoverTask timerTask = new MessageMoverTask(queueName, zsetName);
      Future future;
      if (requiredDelay < MIN_DELAY) {
        future = scheduler.submit(timerTask);
        taskStartTime = currentTime;
      } else {
        future = scheduler.schedule(timerTask, Instant.ofEpochMilli(currentTime + requiredDelay));
      }
      scheduledTaskDetail = new ScheduledTaskDetail(taskStartTime, future);
      queueNameToScheduledTask.put(timerTask.getQueueName(), scheduledTaskDetail);
      return;
    }
    // run existing tasks continue
    long existingDelay = scheduledTaskDetail.getStartTime() - currentTime;
    Future submittedTask = scheduledTaskDetail.getFuture();
    boolean completedOrCancelled = submittedTask.isDone() || submittedTask.isCancelled();
    // tasks older than TASK_ALIVE_TIME are considered dead
    if (!completedOrCancelled && existingDelay < MIN_DELAY && existingDelay > TASK_ALIVE_TIME) {
      try {
        submittedTask.get(DEFAULT_SCRIPT_EXECUTION_TIME, TimeUnit.MILLISECONDS);
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
      } catch (ExecutionException | TimeoutException | CancellationException e) {
        getLogger().debug("{} task failed", scheduledTaskDetail, e);
      }
    }
    // Run was succeeded or cancelled submit new one
    MessageMoverTask timerTask = new MessageMoverTask(queueName, zsetName);
    Future future =
        scheduler.schedule(
            timerTask,
            Instant.ofEpochMilli(getNextScheduleTime(System.currentTimeMillis(), startTime)));
    queueNameToScheduledTask.put(
        timerTask.getQueueName(), new ScheduledTaskDetail(startTime, future));
  }

  private class MessageMoverTask implements Runnable {
    private final String queueName;
    private final String zsetName;

    MessageMoverTask(String queueName, String zsetName) {
      this.queueName = queueName;
      this.zsetName = zsetName;
    }

    String getQueueName() {
      return queueName;
    }

    @Override
    public void run() {
      try {
        if (isQueueActive(queueName)) {
          long currentTime = System.currentTimeMillis();
          Long value =
              defaultScriptExecutor.execute(
                  redisScript, Arrays.asList(queueName, zsetName), currentTime, MAX_MESSAGE);
          schedule(
              queueName, zsetName, getNextScheduleTime(System.currentTimeMillis(), value), true);
        }
      } catch (RedisSystemException e) {
        // no op
      } catch (Exception e) {
        getLogger().warn("Task execution failed for queue: {}", queueName, e);
      }
    }
  }

  private class MessageSchedulerListener implements MessageListener {
    @Override
    public void onMessage(Message message, byte[] pattern) {
      if (message.getBody().length == 0 || message.getChannel().length == 0) {
        return;
      }
      String body = new String(message.getBody());
      String channel = new String(message.getChannel());
      getLogger().debug("Body: {}  Channel: {}", body, channel);
      try {
        Long startTime = Long.parseLong(body);
        String queueName = channelNameToQueueName.get(channel);
        if (queueName == null) {
          getLogger().warn("Unknown channel name {}", channel);
          return;
        }
        String zsetName = queueNameToZsetName.get(queueName);
        if (zsetName == null) {
          getLogger().warn("Unknown zset name {}", queueName);
          return;
        }
        schedule(queueName, zsetName, startTime, false);
      } catch (NumberFormatException e) {
        getLogger().error("Invalid data {} on a channel {}", body, channel);
      }
    }
  }

  @SuppressWarnings("unchecked")
  private void initialize(Map queueDetailMap) {
    Set queueNames = new HashSet<>();
    for (Entry entry : queueDetailMap.entrySet()) {
      String queueName = entry.getKey();
      ConsumerQueueDetail queueDetail = entry.getValue();
      if (isQueueValid(queueDetail)) {
        queueNames.add(queueName);
      }
    }
    defaultScriptExecutor = new DefaultScriptExecutor<>(redisTemplate);
    messageSchedulerListener = new MessageSchedulerListener();
    redisScript = (RedisScript) RedisScriptFactory.getScript(ScriptType.PUSH_MESSAGE);
    queueRunningState = new ConcurrentHashMap<>(queueNames.size());
    queueNameToScheduledTask = new ConcurrentHashMap<>(queueNames.size());
    channelNameToQueueName = new ConcurrentHashMap<>(queueNames.size());
    queueNameToZsetName = new ConcurrentHashMap<>(queueNames.size());
    queueNameToLastMessageSeenTime = new ConcurrentHashMap<>(queueNames.size());
    createScheduler(queueNames.size());
    for (String queueName : queueNames) {
      queueRunningState.put(queueName, false);
    }
  }

  @Override
  public void onApplicationEvent(QueueInitializationEvent event) {
    doStop();
    if (event.isStart()) {
      if (CollectionUtils.isEmpty(event.getQueueDetailMap())) {
        return;
      }
      initialize(event.getQueueDetailMap());
      doStart();
    }
  }
}