高木浩光＠自宅の日記

package jp.takagi_hiromitsu.winny.crawler;
import java.io.*;
import java.net.*;
import java.util.Set;
import java.util.HashSet;
import java.util.Map;
import java.util.HashMap;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.PriorityBlockingQueue;
import java.util.concurrent.TimeUnit;
import jp.takagi_hiromitsu.winny.net.WinnyProtocolInterpreter;
import jp.takagi_hiromitsu.winny.net.WinnyProtocolException;
import jp.takagi_hiromitsu.winny.net.WinnyCommandVisitor;
import jp.takagi_hiromitsu.winny.net.WinnyProtocolCommander;
import jp.takagi_hiromitsu.winny.net.Command3;
import jp.takagi_hiromitsu.winny.net.Command4;
import jp.takagi_hiromitsu.winny.net.Command13;
import jp.takagi_hiromitsu.winny.net.Command10;
import jp.takagi_hiromitsu.winny.Node;
import jp.takagi_hiromitsu.winny.WinnyKey;
import jp.takagi_hiromitsu.winny.WinnyFileID;
import jp.takagi_hiromitsu.winny.WinnyConfiguration;

public class WinnyWalker {
    Set<Node> foundNodes = new HashSet<Node>();
    final int startTime = currentTime();
    int connectionTriedCount = 0;
    int handshakeSucceededCount = 0;
    int extractionSucceededCount = 0;
    static final int MAX_THREADS = 800;
    ThreadPoolExecutor threadPool;
    final Timer timer = new Timer(true);
    WinnyConfiguration conf = new WinnyConfiguration();
    WinnyWalker(Node initialNode) {
        threadPool = new ThreadPoolExecutor(
            MAX_THREADS, MAX_THREADS, 0L, TimeUnit.SECONDS,
            new PriorityBlockingQueue<Runnable>()
        );
        timer.scheduleAtFixedRate(new MonitorTask(), 0L, 10 * 1000L);
        addNode(initialNode, 0);
    }
    boolean addNode(Node node, int priority) {
        if (!node.isValid()) return false;
        boolean added = foundNodes.add(node);
        if (!added) return false;
        int elapsed = currentTime() - startTime;
        threadPool.execute(new Task(node, priority - elapsed));
        return true;
    }
    class Task implements Runnable, Comparable<Task> {
        Node target;
        int priority;
        int taskSubmittedTime;
        int taskStartTime;
        Task(Node target, int priority) {
            this.target = target;
            this.priority = priority;
            taskSubmittedTime = currentTime();
        }
        public int compareTo(Task another) {
            if (this.priority < another.priority) return -1;
            if (this.priority == another.priority) return 0;
            return 1;
        }
        int handshakeCount = 0;
        int extractedNodeCount = 0;
        int newNodeCount = 0;
        Exception lastException;
        String[] clusterWords;
        public void run() {
            taskStartTime = currentTime();
            int errorCount = 0;
            for (int i = 0; i < 4; i++) {
                if (errorCount >= 1) break;
                if (currentTime() - taskStartTime >= 20) break;
                try {
                    connect(30);
                } catch (IOException e) {
                    errorCount++;
                }
            }
            connectionTriedCount++;
            if (handshakeCount > 0) {
                handshakeSucceededCount++;
            }
            if (extractedNodeCount > 0) {
                extractionSucceededCount++;
            }
            printTasklog();
        }
        private void connect(int timeLimit) throws IOException {
            InputStream is = null;
            OutputStream os = null;
            Socket s = null;
            TimerTask closer = null;
            try {
                s = new Socket(target.addr, target.port);
                closer = new SocketCloser(s);
                timer.schedule(closer, timeLimit * 1000L);
                is = new BufferedInputStream(s.getInputStream());
                os = new BufferedOutputStream(s.getOutputStream());
                new WinnyProtocolCommander(os, conf).sendCommand(new Command10());
                WinnyCommandVisitor visitor = new NodeAddingVisitor();
                new WinnyProtocolInterpreter(is, visitor).interpret();
            } catch (WinnyProtocolException e) {
                lastException = e;
            } catch (IOException e) {
                lastException = e;
                throw e;
            } finally {
                if (closer != null) closer.cancel();
                try {
                    if (os != null) os.close();
                    if (is != null) is.close();
                    if (s != null) s.close();
                } catch (IOException e2) {
                }
            }
        }
        class NodeAddingVisitor extends WinnyCommandVisitor {
            public void visit(Command3 c) {
                handshakeCount++;
                clusterWords = c.clusterWords;
            }
            public void visit(Command4 c) {
                extractedNodeCount++;
                boolean added = addNode(new Node(c.addr, c.port), 0);
                if (added) newNodeCount++;
            }
            public void visit(Command13 c) {
                for (int i = 1; i < c.nodePath.length; i++) {
                    Node n = c.nodePath[i];
                    extractedNodeCount++;
                    boolean added = addNode(n, 0);
                    if (added) newNodeCount++;
                }
                Map<Node,Integer> minimum = new HashMap<Node,Integer>();
                for (WinnyKey k: c.keys) {
                    dumpKey(target, k);
                    Node node;
                    if (k.addr.isSiteLocalAddress() || k.addr.isLinkLocalAddress()) {
                        node = target;
                    } else {
                        node = new Node(k.addr, k.port);
                    }
                    int pri = 1500 - k.keyExpirationTimer;
                    if (pri < 0) pri = 0;
                    Integer old = minimum.get(node);
                    if (old == null || old > pri) {
                        minimum.put(node, pri);
                    }
                }
                for (Node n: minimum.keySet()) {
                    extractedNodeCount++;
                    boolean added = addNode(n, minimum.get(n));
                    if (added) newNodeCount++;
                }
            }
        }
        class SocketCloser extends TimerTask {
            Socket socket;
            int lastExtractedNodeCount = 0;
            SocketCloser(Socket socket) {
                this.socket = socket;
            }
            public void run() {
                try {
                    socket.getInputStream().close();
                } catch (IOException e) {
                }
                try {
                    socket.close();
                } catch (IOException e) {
                }
            }
        }

        private void printTasklog() {
            taskOut.printf(
                "%21s pri: %4d life: %3d con: %d new: %5.1f%% (%4d/%4d) %d %d %s %s\n",
                target.toString(),
                priority + (taskStartTime - startTime),
                currentTime() - taskStartTime,
                handshakeCount,
                (float)newNodeCount / extractedNodeCount * 100,
                newNodeCount,
                extractedNodeCount,
                currentTime() - startTime,
                currentTime(),
                clusterWords == null ?
                    "" :
                    clusterWords[0] + "|" + clusterWords[1] + "|" + clusterWords[2],
                lastException == null ? "" : lastException
            );
            taskOut.flush();
        }
    }
    class MonitorTask extends TimerTask {
        public void run() {
            int elapsed = currentTime() - startTime;
            monitorOut.printf(
//                "Elapsed time: %7d; Nodes found: %6d, tried: %6d, " +
//                    "connected: %6d, succeeded: %6d, live: %5.1f%%, " +
//                    "extracted: %5.1f; %Queue lingth: %6d; Threads: %4d;" +
//                    " /sec tried: %5.3f, connected: %5.3f\n",
                "%7d; %6d, %6d, %6d, %6d, %5.1f%%, %5.1f%%; %6d; %4d; %5.3f, %5.3f\n",
                elapsed,
                foundNodes.size(),
                connectionTriedCount,
                handshakeSucceededCount,
                extractionSucceededCount,
                (float)handshakeSucceededCount / connectionTriedCount * 100,
                (float)extractionSucceededCount / connectionTriedCount * 100,
                threadPool.getQueue().size(),
                Thread.activeCount(),
                (float)connectionTriedCount / elapsed,
                (float)handshakeSucceededCount / elapsed
            );
            monitorOut.flush();
        }
    }
    private void dumpKey(Node node, WinnyKey k) {
        keydumpOut.printf(
            "%s %s:%d %s %d %d %d %d \"%s\" %d %s\n",
            node.toString(),
            k.addr.getHostAddress().toString(),
            k.port,
            k.fileid.toString(),
            k.keyExpirationTimer,
            currentTime(),
            k.keyUpdateTime,
            k.referenceCount,
            k.trip,
            k.filesize,
            k.filename
        );
    }
    private static int currentTime() {
        return (int)(System.currentTimeMillis() / 1000);
    }

    static final String usage = "Usage: java WinnyWalker host port [outputdir]";
    static PrintStream keydumpOut, taskOut, monitorOut;
    public static void main(String[] args) throws Exception {
        if (args.length < 2 || args.length > 3) {
            System.err.println(usage);
            System.exit(1);
        }
        String host = args[0];
        int port = Integer.parseInt(args[1]);
        Node initialNode = new Node(InetAddress.getByName(host), port);
        File dir = new File(".");
        if (args.length > 2) {
            dir = new File(args[2]);
        }
        keydumpOut = newPrintStream(dir, "keydump.log");
        taskOut = newPrintStream(dir, "task.log");
        monitorOut = newPrintStream(dir, "monitor.log");
        new WinnyWalker(initialNode);
    }
    static PrintStream newPrintStream(File dir, String file) throws Exception {
        OutputStream os = new FileOutputStream(new File(dir, file));
        return new PrintStream(new BufferedOutputStream(os));
    }
}