PowerMsg3 Local Server SDK Scheduler: Efficient Task Scheduling with Netty HashedWheelTimer

{ "title": "PowerMsg3 Local Server SDK Scheduler: Efficient Task Scheduling with Netty HashedWheelTimer", "description": "This code demonstrates a robust task scheduler implementation within the PowerMsg3 Local Server SDK, leveraging Netty's HashedWheelTimer for efficient and reliable scheduling. It allows scheduling tasks with various delay configurations, including one-time execution, fixed-delay repetition, and cancellation capabilities.", "keywords": "PowerMsg3, Local Server SDK, Scheduler, Task Scheduling, Netty, HashedWheelTimer, Timeout, Runnable, Context, Vertx, Spring, PostConstruct, PreDestroy", "content": "package com.taobao.powermsg3.sdk.localserver.scheduler;

import com.taobao.hsf.NamedThreadFactory; import com.taobao.powermsg3.sdk.localserver.PmLocalServerSDKConfig; import io.netty.util.HashedWheelTimer; import io.netty.util.Timeout; import io.netty.util.TimerTask; import io.vertx.core.Context; import lombok.extern.slf4j.Slf4j; import org.apache.commons.collections4.CollectionUtils; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct; import javax.annotation.PreDestroy; import java.time.Duration; import java.util.Objects; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong;

@Component @Slf4j public final class Scheduler { private final HashedWheelTimer hashedWheelTimer;

private final AtomicLong timerIdGen = new AtomicLong();

private final ConcurrentHashMap<Object, Timeout> timerKey2Timeouts = new ConcurrentHashMap<>();

@Autowired
public Scheduler(final PmLocalServerSDKConfig pmLocalServerSDKConfig) {
    final PmLocalServerSDKConfig.SchedulerConfig schedulerConfig = pmLocalServerSDKConfig.getScheduler();
    final int tickInMillis = Math.max(1, schedulerConfig.getTickDurationInMillis());
    final int ticksPerWheel = Math.max(512, schedulerConfig.getTicksPerWheel());
    this.hashedWheelTimer = new HashedWheelTimer(
        new NamedThreadFactory('localserver-scheduler'),
        tickInMillis, TimeUnit.MILLISECONDS,
        ticksPerWheel
    );
}

public boolean cancel(final Object timerKey) {
    final Timeout timeout = timerKey2Timeouts.get(timerKey);
    tryCancel(timeout);
    return true;
}

public Long scheduleOnce(final Duration delay,
                         final Runnable runnable,
                         final Context context) {
    final Long id = timerIdGen.incrementAndGet();
    scheduleOnce(id, delay, runnable, context);
    return id;
}

public void scheduleOnce(final Object timerKey,
                         final Duration delay,
                         final Runnable runnable,
                         final Context context) {
    Objects.requireNonNull(timerKey, 'timerKey should not be null.');
    Objects.requireNonNull(delay, 'delay should not be null.');
    Objects.requireNonNull(runnable, 'runnable should not be null.');
    Objects.requireNonNull(context, 'context should not be null.');
    timerKey2Timeouts.compute(
        timerKey,
        (key, timeout) -> {
            tryCancel(timeout);
            return hashedWheelTimer.newTimeout(t -> {
                context.runOnContext(event -> {
                    timerKey2Timeouts.remove(timerKey);
                    runnable.run();
                });
            }, delay.toNanos(), TimeUnit.NANOSECONDS);
        }
    );
}

public Long scheduleWithFixDelay(final Duration initialDelay,
                                 final Duration delay,
                                 final Runnable runnable,
                                 final Context context) {
    final Long id = timerIdGen.incrementAndGet();
    scheduleWithFixDelay(id, initialDelay, delay, runnable, context);
    return id;
}

public void scheduleWithFixDelay(final Object timerKey,
                                 final Duration initialDelay,
                                 final Duration delay,
                                 final Runnable runnable,
                                 final Context context) {
    Objects.requireNonNull(timerKey, 'timerKey should not be null.');
    Objects.requireNonNull(initialDelay, 'initialDelay should not be null.');
    Objects.requireNonNull(delay, 'delay should not be null.');
    Objects.requireNonNull(runnable, 'runnable should not be null.');
    Objects.requireNonNull(context, 'context should not be null.');
    final Runnable scheduledRunnable = new Runnable() {
        final Runnable self = this;

        @Override
        public void run() {
            try {
                runnable.run();
            } catch (Exception e) {
                log.error('Error when trigger task of runnable:[{}]', runnable, e);
            } finally {
                timerKey2Timeouts.compute(
                    timerKey,
                    (key, timeout) -> {
                        tryCancel(timeout);
                        return scheduleWithTimerKey(delay, self, context);
                    });
            }
        }
    };

    timerKey2Timeouts.compute(timerKey, (key, timeout) -> {
        tryCancel(timeout);
        return scheduleWithTimerKey(initialDelay, scheduledRunnable, context);
    });
}

private boolean tryCancel(final Timeout timeout) {
    if (timeout == null || timeout.isCancelled() || timeout.isExpired()) {
        return true;
    }
    return timeout.cancel();
}

private Timeout scheduleWithTimerKey(final Duration delay,
                                     final Runnable runnable,
                                     final Context context) {
    return hashedWheelTimer.newTimeout(
        t -> context.runOnContext(event -> runnable.run()),
        delay.toNanos(),
        TimeUnit.NANOSECONDS);
}


@PostConstruct
protected void init() {
    hashedWheelTimer.start();
}

@PreDestroy
protected void stop() {
    final Set<Timeout> timeouts = hashedWheelTimer.stop();
    if (CollectionUtils.isNotEmpty(timeouts)) {
        timeouts.stream()
            .filter(timeout -> !timeout.isCancelled() && !timeout.isExpired())
            .forEach(timeout -> {
                final TimerTask task = timeout.task();
                try {
                    task.run(timeout);
                } catch (Exception e) {
                    //Ignore
                }
            });
    }
}

}

public abstract class TimerTaskRunnable extends TimerTask { private final Runnable runnable;

public TimerTaskRunnable(Runnable runnable) {
    this.runnable = runnable;
}

@Override
public void run(Timeout timeout) {
    runnable.run();
}

}