Pinger.java

package com.intrbiz.bergamot.pinger;

import java.net.InetAddress;
import java.security.SecureRandom;
import java.util.Arrays;
import java.util.Objects;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.TimeUnit;

import com.intrbiz.bergamot.net.raw.RawEngine;
import com.intrbiz.bergamot.net.raw.jna.SockAddrIn;
import com.intrbiz.bergamot.net.raw.model.IPPacket;
import com.intrbiz.bergamot.net.raw.model.IPPayload;
import com.intrbiz.bergamot.net.raw.model.payload.ICMPPacket;

/**
 * Ping lots of targets concurrently measuring the RTT
 */
public class Pinger
{
    private ConcurrentMap<String, PingTarget> targetsByName = new ConcurrentHashMap<String, PingTarget>();
    
    private ConcurrentMap<PingKey, SentMessage> messages = new ConcurrentHashMap<PingKey, SentMessage>();
    
    private RawEngine<SentMessage> engine;
    
    private Timer timer = new Timer();
    
    private SecureRandom random = new SecureRandom();
    
    private final long lookupInterval = TimeUnit.MINUTES.toMillis(5);
    
    public Pinger()
    {
        super();
        this.engine = new RawEngine<SentMessage>(this::recv);
    }
    
    public void start()
    {
        this.engine.start();
    }
    
    private void recv(IPPacket packet)
    {
        if (packet.getPayload() instanceof ICMPPacket)
        {
            ICMPPacket icmp = (ICMPPacket) packet.getPayload();
            // is the ICMP packet
            if (icmp.getType() == 0)
            {
                // lookup the message
                // ignore any messages we cannot correlate
                SentMessage sent = this.messages.remove(new PingKey(packet.getSource().getAddress(), icmp.getId(), icmp.getSequence()));
                if (sent != null)
                {
                    // cancel the timeout
                    sent.cancel();
                    // compute the RTT (in microseconds)
                    long rtt = (System.nanoTime() - sent.sentAt) / 1000L;
                    // update the stats
                    sent.target.reply(rtt);
                }
            }
        }
    }
    
    private void send(PingTarget target)
    {
        // get the target address
        InetAddress address = target.getAddress();
        if (address != null)
        {
            // assemble the ICMP request
            ICMPPacket packet = new ICMPPacket(ICMPPacket.ICMP_TYPE_ECHO, target.getId(), target.nextSequence());
            // per message state
            PingKey key = new PingKey(address.getAddress(), packet.getId(), packet.getSequence());
            SentMessage sent = new SentMessage(target, key);
            this.messages.put(key, sent);
            // send
            this.engine.send(sent, packet, address, 8, this::sent);
        }
        else
        {
            // the address is unknown for some reason
            target.unknownHost();
        }
    }
    
    void sent(SentMessage context, IPPayload payload, InetAddress to, int port, SockAddrIn.ByReference addr)
    {
        // record when we sent the packet
        context.sentAt = System.nanoTime();
        // schedule timeout task
        this.timer.schedule(context, (long) (context.target.getInterval() * 0.75));
    }
    
    // targets
    
    public PingTarget getTarget(String host)
    {
        Objects.requireNonNull(host);
        return this.targetsByName.get(host);
    }
    
    public PingTarget addTarget(String host, long interval, TimeUnit intervalTimeUnit)
    {
        Objects.requireNonNull(intervalTimeUnit);
        return this.addTarget(host, interval, intervalTimeUnit, null);
    }
    
    public PingTarget addTarget(String host, long interval, TimeUnit intervalTimeUnit, long timeout, TimeUnit timeoutTimeUnit)
    {
        Objects.requireNonNull(intervalTimeUnit);
        Objects.requireNonNull(timeoutTimeUnit);
        return this.addTarget(host, interval, intervalTimeUnit, timeout, timeoutTimeUnit, null);
    }
    
    public PingTarget addTarget(String host, long interval)
    {
        return this.addTarget(host, interval, (OnPingUpdate) null);
    }
    
    public PingTarget addTarget(String host, long interval, long timeout)
    {
        return this.addTarget(host, interval, timeout, null);
    }
    
    public PingTarget addTarget(String host, long interval, TimeUnit intervalTimeUnit, OnPingUpdate updateCallback)
    {
        Objects.requireNonNull(intervalTimeUnit);
        return this.addTarget(host, intervalTimeUnit.toMillis(interval), updateCallback);
    }
    
    public PingTarget addTarget(String host, long interval, TimeUnit intervalTimeUnit, long timeout, TimeUnit timeoutTimeUnit, OnPingUpdate updateCallback)
    {
        Objects.requireNonNull(intervalTimeUnit);
        Objects.requireNonNull(timeoutTimeUnit);
        return this.addTarget(host, intervalTimeUnit.toMillis(interval), timeoutTimeUnit.toMillis(timeout), updateCallback);
    }
    
    public PingTarget addTarget(String host, long interval, OnPingUpdate updateCallback)
    {
        return this.addTarget(host, interval, (long) (interval * 0.75), updateCallback);
    }
    
    public PingTarget addTarget(String host, long interval, long timeout, OnPingUpdate updateCallback)
    {
        Objects.requireNonNull(host);
        final PingTarget newTarget = new PingTarget(host, interval, timeout);
        PingTarget oldTarget = this.targetsByName.putIfAbsent(host, newTarget);
        if (oldTarget == null)
        {
            // setup the target
            oldTarget = newTarget;
            // callbacks
            if (updateCallback != null) newTarget.setUpdateCallback(updateCallback);
            // scheduled tasks
            TimerTask sendTask = new TimerTask()
            {
                public void run()
                {
                    Pinger.this.send(newTarget);
                }
            };
            newTarget.setSendTask(sendTask);
            TimerTask lookupTask = new TimerTask()
            {
                public void run()
                {
                    newTarget.lookup();
                }
            };
            newTarget.setLookupTask(lookupTask);
            // schedule
            this.timer.scheduleAtFixedRate(sendTask, this.random.nextInt((int) interval), interval);
            this.timer.scheduleAtFixedRate(lookupTask, this.lookupInterval, this.lookupInterval);
        }
        return oldTarget;
    }
    
    public void removeTarget(String host)
    {
        Objects.requireNonNull(host);
        // remove the target
        PingTarget target = this.targetsByName.remove(host);
        if (target != null)
        {
            // cancel send task
            TimerTask sendTask = target.getSendTask();
            if (sendTask != null) sendTask.cancel();
            // cancel lookup task
            TimerTask lookupTask = target.getLookupTask();
            if (lookupTask != null) lookupTask.cancel();
        }
    }
    
    // internal data structures
    
    private static final class PingKey
    {
        public final byte[] address;
        
        public final short id;
        
        public final short seq;
        
        public PingKey(byte[] address, short id, short seq)
        {
            this.address = address;
            this.id = id;
            this.seq = seq;
        }

        @Override
        public int hashCode()
        {
            final int prime = 31;
            int result = 1;
            result = prime * result + id;
            result = prime * result + seq;
            return result;
        }

        @Override
        public boolean equals(Object obj)
        {
            if (this == obj) return true;
            if (obj == null) return false;
            if (getClass() != obj.getClass()) return false;
            PingKey other = (PingKey) obj;
            if (!Arrays.equals(address, other.address)) return false;
            if (id != other.id) return false;
            if (seq != other.seq) return false;
            return true;
        }
    }
    
    private final class SentMessage extends TimerTask
    {
        public final PingTarget target;
        
        public final PingKey key;
        
        public long sentAt;
        
        public SentMessage(PingTarget target, PingKey key)
        {
            this.target = target;
            this.key = key;
        }
        
        public void run()
        {
            // timeout
            this.target.timeout();
            // clean up
            Pinger.this.messages.remove(this.key);
        }
    }
}