Kafka TimingWheel

package com.yit;

import java.util.Date;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.DelayQueue;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.locks.StampedLock;

/**
 * Created by jenkin.z.chen on 2018/6/26.
 */
public class Timer {
    private TimingWheel timingWheel;
    private DelayQueue<TimerTaskList> delayQueue;
    private final StampedLock stampedLock = new StampedLock();
    private BlockingQueue<TimerTaskEntry> blockingQueue = new LinkedBlockingDeque<>();
    public Timer(long tickMs,int wheelSize,long startMs){
        delayQueue = new DelayQueue();
        timingWheel = new TimingWheel(tickMs,wheelSize,startMs,delayQueue);
    }

    public BlockingQueue<TimerTaskEntry> getBlockingQueue() {
        return blockingQueue;
    }

    public void add(TimerTaskEntry timerTaskEntry){
        long stamp = stampedLock.tryOptimisticRead();
        if (!stampedLock.validate(stamp)) {
            stamp = stampedLock.readLock();
            try {
                timingWheel.add(timerTaskEntry);
                return;
            }finally {
                stampedLock.unlock(stamp);
            }
        }
        timingWheel.add(timerTaskEntry);
    }
    boolean reinsert(TimerTaskEntry timerTaskEntry){
        return timingWheel.add(timerTaskEntry);
    }
    private boolean advanceClock() throws InterruptedException {
        while(true){
            TimerTaskList timerTaskList = delayQueue.take();
            long stamp = stampedLock.writeLock();
            try{
                timingWheel.advanceClock(timerTaskList.getExpiration());
                timerTaskList.flush(this);
            }finally {
                stampedLock.unlock(stamp);
            }

        }
    }

    public static void main(String[] args) throws InterruptedException {
        Timer timer = new Timer(1000,10,System.currentTimeMillis());
        Thread thread = new Thread(()->{
            try {
                timer.advanceClock();
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        });
        thread.start();

       // for(int i=1;i<Integer.MAX_VALUE;i++){
            long longTime =  1*1000+System.currentTimeMillis();
            //System.out.println("print time :"+longTime);
            timer.add(new TimerTaskEntry(longTime));
            long longTime1 =  12*1000+System.currentTimeMillis();
            long longTime3 =  17*1000+System.currentTimeMillis();
            //System.out.println("print time :"+longTime);
            timer.add(new TimerTaskEntry(longTime1));
            timer.add(new TimerTaskEntry(longTime3));

            while(true)
            {
                System.out.println("*******************************"+new Date(timer.getBlockingQueue().take().getExpirations()));
            }

    }
}

 

package com.yit;

/**
 * Created by jenkin.z.chen on 2018/6/26.
 */
public class TimerTaskEntry {
    private final long expirations;
    public TimerTaskEntry(long expirations){
        this.expirations = expirations;
    }

    public long getExpirations() {
        return expirations;
    }
}

 

 

package com.yit;

import java.util.Date;
import java.util.LinkedList;
import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.Delayed;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

/**
 * Created by jenkin.z.chen on 2018/6/26.
 */
public class TimerTaskList implements Delayed {
    private CopyOnWriteArrayList<TimerTaskEntry> list = new CopyOnWriteArrayList<>();
    private final AtomicLong expiration = new AtomicLong(-1);
    public TimerTaskList(){

    }

    public void flush(Timer timer){
        list.stream().forEach(timerTaskEntry->{
            if(!timer.reinsert(timerTaskEntry)){
                timer.getBlockingQueue().offer(timerTaskEntry);
            }
            this.list.remove(timerTaskEntry);
        });
        expiration.set(-1);
    }
    public void add(TimerTaskEntry string){
        list.add(string);
    }
    public boolean setExpiration(long expiration){
        return this.expiration.getAndSet(expiration)!=expiration;
    }
    @Override
    public long getDelay(TimeUnit unit) {
        return unit.convert(Math.max(expiration.get() - System.currentTimeMillis(), 0), TimeUnit.MILLISECONDS);
    }

    public long getExpiration() {
        return expiration.get();
    }

    @Override
    public int compareTo(Delayed o) {
        TimerTaskList other = (TimerTaskList) o;

         if(getExpiration() < other.getExpiration())
             return -1;
        else if(getExpiration() > other.getExpiration())
             return 1;
        else return 0;
    }
}

 

 

package com.yit;

import java.util.concurrent.DelayQueue;

/**
 * Created by jenkin.z.chen on 2018/6/26.
 */
public class TimingWheel {
    private final long tickMs;
    private final int wheelSize;
    private final DelayQueue<TimerTaskList> delayQueue;
    private final long interval;
    private final TimerTaskList[] buckets;
    private long currentTime;
    private volatile TimingWheel overflowWheel;

    public TimingWheel(long tickMs, int wheelSize,long startMs, DelayQueue<TimerTaskList> delayQueue) {
        this.tickMs = tickMs;
        this.wheelSize = wheelSize;
        this.delayQueue = delayQueue;
        this.interval = tickMs*wheelSize;
        this.buckets = new TimerTaskList[wheelSize];
        for(int i=0;i<wheelSize;i++){
            buckets[i] = new TimerTaskList();
        }
        this.currentTime = startMs-startMs%tickMs;
        this.overflowWheel = null;
    }
    public boolean add(TimerTaskEntry timerTaskEntry){
        long expiration = timerTaskEntry.getExpirations();
        boolean flag;
        if (expiration < currentTime + tickMs) {
           return  false;
        } else if (expiration < currentTime + interval) {
            // Put in its own bucket
            long virtualId = expiration / tickMs;
            TimerTaskList bucket = buckets[(int) (virtualId % wheelSize)];
            bucket.add(timerTaskEntry);
            if (bucket.setExpiration(virtualId * tickMs)) {
                delayQueue.offer(bucket);
            }
            return true;
        } else {
            // Out of the interval. Put it into the parent timer
            if (overflowWheel == null) addOverflowWheel();
            flag = overflowWheel.add(timerTaskEntry);
        }
        return flag;
    }
    private synchronized void addOverflowWheel(){
        if (overflowWheel == null) {
            overflowWheel = new TimingWheel(
                    interval,
                    wheelSize,
                    currentTime,
                    delayQueue
            );
        }
    }
    public void advanceClock(long timeMs){
        if (timeMs >= currentTime + tickMs) {
            currentTime = timeMs - (timeMs % tickMs);
            // Try to advance the clock of the overflow wheel if present
            if (overflowWheel != null) {
                overflowWheel.advanceClock(currentTime);
            }
        }
    }
}