Java tutorial
/* * 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; } } }