org.aotorrent.common.TorrentEngine.java Source code

Java tutorial

Introduction

Here is the source code for org.aotorrent.common.TorrentEngine.java

Source

/*
 * Copyright 2014 Napolov Dmitry
 *
 * 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
 *
 *     http://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 org.aotorrent.common;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import io.netty.bootstrap.Bootstrap;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelFutureListener;
import io.netty.channel.EventLoopGroup;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.traffic.GlobalTrafficShapingHandler;
import org.aotorrent.client.OutboundChannelInitializer;
import org.aotorrent.common.connection.AbstractTrackerConnection;
import org.aotorrent.common.connection.PeerConnection;
import org.aotorrent.common.protocol.peer.HandshakeRequest;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nullable;
import javax.annotation.concurrent.GuardedBy;
import java.io.UnsupportedEncodingException;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

/**
 * Project: AOTorrent
 * User:    dmitry
 * Date:    11/8/13
 */
public class TorrentEngine {
    private static final Logger LOGGER = LoggerFactory.getLogger(TorrentEngine.class);
    private static final byte[] peerId = "-AO0001-000000000000".getBytes(); //TODO need to give right peerID

    private final InetSocketAddress address;
    private final EventLoopGroup workerGroup;

    private final GlobalTrafficShapingHandler counter;
    private AtomicLong downloadedTotal = new AtomicLong(0);
    private AtomicLong uploadedTotal = new AtomicLong(0);

    private final Torrent torrent;
    private final Set<AbstractTrackerConnection> trackerConnections = Sets.newLinkedHashSet();
    private final Map<SocketAddress, PeerConnection> peerConnections = Maps.newHashMap();
    private final Object piecesMux = new Object();
    @GuardedBy("piecesMux")
    private final BitSet bitField = new BitSet();
    @GuardedBy("piecesMux")
    private final List<Piece> piecesToDownload = Lists.newArrayList();
    @GuardedBy("piecesMux")
    private final List<Piece> inProgress = Lists.newArrayList();
    private final ExecutorService writeThreadPool;
    private EngineStatus status = EngineStatus.NONE;
    @GuardedBy("piecesMux")
    private List<Piece> pieces = Collections.emptyList();
    @Nullable
    private ExecutorService trackerConnectionThreads;

    public TorrentEngine(Torrent torrent, InetSocketAddress address, EventLoopGroup workerGroup) {
        this.torrent = torrent;
        this.address = address;
        this.workerGroup = workerGroup;

        counter = new GlobalTrafficShapingHandler(workerGroup, 1000);
        writeThreadPool = Executors.newSingleThreadExecutor();
    }

    private void initTrackers() {

        Collection<String> trackers = torrent.getTrackers();

        trackerConnectionThreads = Executors.newFixedThreadPool(trackers.size());

        for (String trackerUrl : trackers) {
            AbstractTrackerConnection trackerConnection = AbstractTrackerConnection.createConnection(this,
                    trackerUrl, torrent.getInfoHash(), address.getAddress(), address.getPort(), workerGroup);
            if (trackerConnection != null) {
                assert trackerConnectionThreads != null;
                trackerConnectionThreads.submit(trackerConnection);
                trackerConnections.add(trackerConnection);
            }
        }
    }

    /**
     * Adds new peers to this engine.
     * <p>If particular peer is already exist, it will be ignored.</p>
     *
     * @param peers Peers to append
     */
    public void appendPeers(Collection<InetSocketAddress> peers) {
        synchronized (peerConnections) {
            peers.removeAll(peerConnections.keySet());

            for (SocketAddress peer : peers) {
                if (peer.equals(address)) {
                    continue;
                }

                final Bootstrap bootstrap = new Bootstrap();
                bootstrap.group(workerGroup);
                bootstrap.channel(NioSocketChannel.class);
                bootstrap.remoteAddress(peer);
                bootstrap.handler(new OutboundChannelInitializer(this));
                final ChannelFuture connect = bootstrap.connect();
                connect.addListener(new ChannelFutureListener() {

                    @Override
                    public void operationComplete(ChannelFuture future) throws Exception {
                        final Throwable cause = future.cause();
                        if (cause != null) {
                            LOGGER.error("Error {}", cause.getLocalizedMessage());
                            return;
                        }
                        final Channel channel = future.channel();
                        final HandshakeRequest handshakeRequest = new HandshakeRequest(getInfoHash(), peerId);
                        final byte[] msg = handshakeRequest.toTransmit();
                        channel.writeAndFlush(msg);
                    }
                });
            }
        }
    }

    public void removePeer(SocketAddress address) {
        peerConnections.remove(address);
    }

    private List<Piece> createPieces() throws UnsupportedEncodingException {
        List<Piece> pieceList = Lists.newArrayList();

        int pieceCount = (int) Math.ceil((double) torrent.getSize() / torrent.getPieceLength());
        for (int i = 0; i < pieceCount - 1; i++) {
            byte[] hash = Arrays.copyOfRange(torrent.getPieces().getBytes(Torrent.DEFAULT_TORRENT_ENCODING),
                    i * Torrent.INFO_HASH_LENGTH, (i + 1) * Torrent.INFO_HASH_LENGTH);
            final Piece piece = new Piece(this, torrent, i, hash);
            piece.checkExistingData();
            pieceList.add(piece);
        }

        byte[] hash = Arrays.copyOfRange(torrent.getPieces().getBytes(Torrent.DEFAULT_TORRENT_ENCODING),
                (pieceCount - 1) * Torrent.INFO_HASH_LENGTH, (pieceCount) * Torrent.INFO_HASH_LENGTH);
        int lastPieceLength = (int) (torrent.getSize() % torrent.getPieceLength());
        if (lastPieceLength == 0) {
            lastPieceLength = torrent.getPieceLength();
        }
        final Piece piece = new Piece(this, torrent, (pieceCount - 1), hash, lastPieceLength);
        piece.checkExistingData();
        pieceList.add(piece);

        return Collections.unmodifiableList(pieceList);
    }

    /**
     * Starts engine.
     *
     * @throws UnsupportedEncodingException
     */
    public void init() throws UnsupportedEncodingException {
        status = EngineStatus.INIT_PIECES;
        this.pieces = createPieces();

        for (Piece piece : pieces) {
            if (!piece.isComplete()) {
                piecesToDownload.add(piece);
            }
        }

        boolean done = isTorrentDone();

        status = EngineStatus.STARTING;

        initTrackers();

        status = done ? EngineStatus.SEEDING : EngineStatus.DOWNLOADING;
    }

    public byte[] getPeerId() {
        return peerId;
    }

    public BitSet getBitField() {
        return this.bitField;
    }

    public int getPieceCount() {
        return pieces.size();
    }

    /**
     * Gives "nearest" piece to download based on given peer's bitfield
     *
     * @param bitField Peer's bitfield
     * @return Piece or null if peer don't have peers we need.
     */
    @Nullable
    public Piece getNextPiece(BitSet bitField) {

        if (status != EngineStatus.DOWNLOADING) {
            return null;
        }

        synchronized (piecesMux) {
            final boolean endGame = isEndGame();
            if (endGame) {
                synchronized (piecesMux) {
                    if (isEndGame()) {
                        final int lastIndex = inProgress.size() - 1;
                        int rnd = (lastIndex == 0) ? 0 : new Random().nextInt(lastIndex);
                        return inProgress.get(rnd);
                    }
                }
            }

            Collections.sort(piecesToDownload);

            for (Piece piece : piecesToDownload) {
                if (bitField.get(piece.getIndex()) && !piece.isComplete() && !inProgress.contains(piece)) {
                    inProgress.add(piece);
                    return piece;
                }

            }
        }

        return null;
    }

    /**
     * Shows if peer has something that we need.
     *
     * @param bitField Peer's bitfield
     * @return true if peer has pieces we need
     */
    public boolean isUsefulPeer(BitSet bitField) {
        BitSet currentBitField = this.getBitField();

        boolean useful = false;

        synchronized (piecesMux) {
            for (int i = bitField.nextSetBit(0); i >= 0; i = bitField.nextSetBit(i + 1)) {
                pieces.get(i).increasePeerCount();
                if (!currentBitField.get(i)) {
                    useful = true;
                }
            }
        }

        return useful;
    }

    private boolean isTorrentDone() {
        return piecesToDownload.isEmpty();
    }

    private boolean isEndGame() {
        return inProgress.size() == piecesToDownload.size() && !piecesToDownload.isEmpty();
    }

    @Nullable
    public Piece getPiece(int index) {
        return pieces.get(index);
    }

    /**
     * Remove piece inProgress state.
     * <p>Send HAVE messages to other peers if this peace is downloaded correctly.</p>
     *
     * @param piece Piece
     */
    public void setPieceDone(@Nullable Piece piece) {
        if (piece != null) {
            if (piece.isComplete()) {
                synchronized (piecesMux) {
                    inProgress.remove(piece);
                    piecesToDownload.remove(piece);
                    bitField.set(piece.getIndex());
                }
                synchronized (peerConnections) {
                    for (PeerConnection peerConnection : peerConnections.values()) {
                        peerConnection.haveNewPiece(piece.getIndex());
                    }
                }
            } else {
                synchronized (piecesMux) {
                    inProgress.remove(piece);
                }
            }
        }
    }

    public byte[] getInfoHash() {
        return torrent.getInfoHash();
    }

    public void shutdown() {
        if (trackerConnectionThreads != null) {
            trackerConnectionThreads.shutdown();
            trackerConnectionThreads.shutdownNow();
        }
        writeThreadPool.shutdown();
        try {
            writeThreadPool.awaitTermination(10, TimeUnit.SECONDS);
        } catch (InterruptedException e) {
            writeThreadPool.shutdownNow();
        }
    }

    public void increaseUploadedCounter(long amount) {
        uploadedTotal.addAndGet(amount);
    }

    public void increaseDownloadedCounter(long amount) {
        downloadedTotal.addAndGet(amount);
    }

    public String getName() {
        return torrent.getName();
    }

    public long getSize() {
        return torrent.getSize();
    }

    public long getDownloaded() {
        long downloaded = 0;
        for (Piece piece : pieces) {
            if (piece.isComplete()) {
                downloaded += piece.getPieceLength();
            }
        }
        return downloaded;
    }

    public int getConnectedPeers() {
        return peerConnections.size();
    }

    public long getDownloadSpeed() {
        return counter.trafficCounter().lastReadThroughput() / 1000;
    }

    public long getUploadSpeed() {
        return counter.trafficCounter().lastWriteThroughput() / 1000;
    }

    public int getDownloadingFromCount() {
        int count = 0;
        for (PeerConnection peerConnection : peerConnections.values()) {
            if (peerConnection.isDownloadingFrom()) {
                count++;
            }
        }
        return count;
    }

    public int getUploadingToCount() {
        int count = 0;
        for (PeerConnection peerConnection : peerConnections.values()) {
            if (peerConnection.isUploadingTo()) {
                count++;
            }
        }
        return count;
    }

    public long getDownloadedTotal() {
        return downloadedTotal.get();
    }

    public long getUploadedTotal() {
        return uploadedTotal.get();
    }

    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder();
        sb.append("Status: ").append(status).append(" ");
        if (status == EngineStatus.DOWNLOADING || status == EngineStatus.SEEDING) {
            sb.append("Known/Connected peers: ").append("?/").append(peerConnections.size()).append(" ");
            sb.append("Pieces(Downloaded/Total(InProgress)): ").append(getBitField().cardinality()).append('/')
                    .append(pieces.size()).append("(").append(inProgress.size()).append(")").append(" ");
            sb.append("Size(Downloaded/Total): ")
                    .append(getBitField().cardinality() * ((long) torrent.getPieceLength()) / 1024).append('/')
                    .append(torrent.getSize() / 1024).append(" ");
            sb.append("(").append(getBitField().cardinality() * 100 / pieces.size()).append("%)").append(" ");
            sb.append("Speed(Downloaded/Uploaded): ").append(counter.trafficCounter().lastReadThroughput() / 1000)
                    .append('/').append(counter.trafficCounter().lastWriteThroughput() / 1000).append(" ");
            sb.append("Cache(hit/miss): ").append(torrent.getCacheHit()).append('/').append(torrent.getCacheMiss());
        }
        return String.valueOf(sb);
    }

    public void registerConnection(SocketAddress socketAddress, PeerConnection peerConnection) {
        synchronized (peerConnections) {
            peerConnections.put(socketAddress, peerConnection);
        }
    }

    public Object getPiecesMux() {
        return piecesMux;
    }

    public GlobalTrafficShapingHandler getCounter() {
        return counter;
    }

    public ExecutorService getWriteThreadPool() {
        return writeThreadPool;
    }

    private enum EngineStatus {
        NONE("None"), INIT_PIECES("Initializing pieces"), STARTING("Starting"), DOWNLOADING("Downloading"), SEEDING(
                "Seeding");

        private final String status;

        EngineStatus(String status) {
            this.status = status;
        }

        @Override
        public String toString() {
            return status;
        }
    }
}