bt.torrent.messaging.Assignments Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of bt-core Show documentation
Show all versions of bt-core Show documentation
BitTorrent Client Library (Core)
package bt.torrent.messaging;
import bt.net.Peer;
import bt.runtime.Config;
import bt.data.Bitfield;
import bt.torrent.BitfieldBasedStatistics;
import bt.torrent.selector.PieceSelector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.BitSet;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.Map;
import java.util.Optional;
import java.util.Random;
import java.util.Set;
public class Assignments {
private static final Logger LOGGER = LoggerFactory.getLogger(Assignments.class);
private static final int MAX_ASSIGNED_PIECES_PER_PEER = 50;
private Config config;
private Bitfield bitfield;
private PieceSelector selector;
private BitfieldBasedStatistics pieceStatistics;
private Set assignedPieces;
private Map assignments;
private Map> peers;
private Random random;
public Assignments(Bitfield bitfield, PieceSelector selector, BitfieldBasedStatistics pieceStatistics, Config config) {
this.bitfield = bitfield;
this.selector = selector;
this.pieceStatistics = pieceStatistics;
this.config = config;
this.assignedPieces = new HashSet<>();
this.assignments = new HashMap<>();
this.peers = new HashMap<>();
this.random = new Random(System.currentTimeMillis());
}
public Assignment get(Peer peer) {
return assignments.get(peer);
}
public void remove(Assignment assignment) {
assignment.abort();
assignments.remove(assignment.getPeer());
assignedPieces.remove(assignment.getPiece());
}
public int count() {
return assignments.size();
}
public int workersCount() {
return peers.size();
}
public Optional assign(Peer peer) {
LinkedList pieces = peers.get(peer);
if (pieces == null || pieces.isEmpty()) {
return Optional.empty();
}
boolean endgame = isEndgame();
StringBuilder buf = LOGGER.isTraceEnabled() ? new StringBuilder() : null;
if (LOGGER.isTraceEnabled()) {
buf.append("Trying to claim next assignment for peer ");
buf.append(peer);
buf.append(". Number of remaining pieces: ");
buf.append(bitfield.getPiecesRemaining());
buf.append(", number of pieces in progress: ");
buf.append(assignedPieces.size());
buf.append(", endgame: " + endgame);
buf.append(". ");
}
Optional selectedPiece;
if (endgame) {
// take random piece to minimize number of pieces
// requested from different peers at the same time
Integer pieceIndex = pieces.remove(random.nextInt(pieces.size()));
selectedPiece = Optional.of(pieceIndex);
} else {
Integer piece;
boolean assigned = true;
Iterator iter = pieces.iterator();
do {
piece = iter.next();
if (bitfield.isComplete(piece)) {
iter.remove();
if (LOGGER.isTraceEnabled()) {
buf.append("Checking next piece in queue: {" + piece + "}; piece is completed. ");
}
continue;
}
assigned = assignedPieces.contains(piece);
if (assigned && LOGGER.isTraceEnabled()) {
buf.append("Checking next piece in queue: {" + piece + "}; piece is assigned. ");
}
} while (assigned && iter.hasNext());
if (!assigned) {
iter.remove();
}
selectedPiece = assigned ? Optional.empty() : Optional.of(piece);
}
if (LOGGER.isTraceEnabled()) {
if (selectedPiece.isPresent()) {
buf.append(" => Assigning piece #");
buf.append(selectedPiece.get());
buf.append(" to current peer");
} else {
buf.append(" => No pieces to assign.");
}
LOGGER.trace(buf.toString());
}
return selectedPiece.isPresent() ? Optional.of(assign(peer, selectedPiece.get())) : Optional.empty();
}
private boolean isEndgame() {
// if all remaining pieces are requested,
// that would mean that we have entered the "endgame" mode
return bitfield.getPiecesRemaining() <= assignedPieces.size();
}
private Assignment assign(Peer peer, Integer piece) {
Assignment assignment = new Assignment(peer, piece, config.getMaxPieceReceivingTime());
assignments.put(peer, assignment);
assignedPieces.add(piece);
if (LOGGER.isTraceEnabled()) {
LOGGER.trace("Assigning piece #{} to peer: {}", piece, peer);
}
return assignment;
}
/**
* Updates the lists of interesting pieces for the provided peers.
*
* @return Collection of peers that have interesting pieces and can be given an assignment
*/
// TODO: select from seeders first
public Set update(Set ready, Set choking) {
Iterator suggested = selector.getNextPieces(pieceStatistics).iterator();
if (LOGGER.isTraceEnabled()) {
LOGGER.trace("Updating assignments. Piece selector has more pieces: {}, number of ready peers: {}, number of assigned peers: {}",
suggested.hasNext(), ready.size(), assignments.size());
}
while (suggested.hasNext() && ready.size() > 0) {
Integer piece = suggested.next();
Iterator iter = new HashSet<>(ready).iterator();
while (iter.hasNext()) {
Peer peer = iter.next();
Optional peerBitfield = pieceStatistics.getPeerBitfield(peer);
if (!peerBitfield.isPresent()) {
iter.remove();
continue;
}
LinkedList queue = peers.get(peer);
if (queue != null && queue.size() > MAX_ASSIGNED_PIECES_PER_PEER) {
iter.remove();
continue;
}
boolean hasPiece = peerBitfield.get().isVerified(piece);
if (hasPiece) {
if (queue == null) {
queue = new LinkedList<>();
peers.put(peer, queue);
}
if (!queue.contains(piece)) {
queue.add(piece);
if (LOGGER.isTraceEnabled()) {
LOGGER.trace("Adding piece #{} to peer's queue: {}. Number of pieces in peer's queue: {}",
piece, peer, queue.size());
}
}
}
}
}
Iterator>> iter = peers.entrySet().iterator();
while (iter.hasNext()) {
Map.Entry> e = iter.next();
Peer peer = e.getKey();
LinkedList pieces = e.getValue();
if (!ready.contains(peer) || pieces.isEmpty()) {
iter.remove();
}
}
Set result = new HashSet<>(peers.keySet());
for (Peer peer : choking) {
if (hasInterestingPieces(peer)) {
result.add(peer);
}
}
return result;
}
private boolean hasInterestingPieces(Peer peer) {
Optional peerBitfieldOptional = pieceStatistics.getPeerBitfield(peer);
if (!peerBitfieldOptional.isPresent()) {
return false;
}
BitSet peerBitfield = BitSet.valueOf(peerBitfieldOptional.get().getBitmask());
BitSet localBitfield = BitSet.valueOf(bitfield.getBitmask());
peerBitfield.andNot(localBitfield);
return peerBitfield.cardinality() > 0;
}
}