org.aotorrent.common.Piece.java Source code

Introduction

Here is the source code for org.aotorrent.common.Piece.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 io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import org.apache.commons.codec.binary.Hex;
import org.apache.commons.codec.digest.DigestUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.lang.ref.SoftReference;
import java.util.Arrays;
import java.util.BitSet;

/**
 * Project: AOTorrent
 * User:    dmitry
 * Date:    11/8/13
 */
public class Piece implements Comparable<Piece> {
    public static final int DEFAULT_BLOCK_LENGTH = 16384;
    public static final int PIECE_HASH_LENGTH = 20;
    private static final Logger LOGGER = LoggerFactory.getLogger(Piece.class);
    private final TorrentEngine torrentEngine;

    private final Torrent torrent;
    private final int index;
    private final int pieceLength;
    private final byte[] hash;
    private final BitSet blockComplete;
    @Nullable
    private ByteBuf buffer;
    @Nullable
    private SoftReference<ByteBuf> softBuffer;
    private int blockCount;

    private boolean complete = false;
    private int peerCount = 0;

    public Piece(TorrentEngine torrentEngine, Torrent torrent, int index, byte[] hash) {
        this(torrentEngine, torrent, index, hash, torrent.getPieceLength());
    }

    public Piece(TorrentEngine torrentEngine, Torrent torrent, int index, byte[] hash, int pieceLength) {
        this.torrentEngine = torrentEngine;
        this.torrent = torrent;
        this.index = index;
        this.pieceLength = pieceLength;
        this.hash = hash;
        this.blockComplete = new BitSet();
        this.blockCount = (int) Math.ceil((double) pieceLength / (double) DEFAULT_BLOCK_LENGTH);
    }

    void checkExistingData() {
        try {
            buffer = torrent.getFileStorage().read(index, pieceLength);
            byte[] pieceHash = DigestUtils.sha1(buffer.array());

            if (Arrays.equals(pieceHash, hash)) {
                complete = true;
                softBuffer = new SoftReference<>(buffer);
                torrentEngine.setPieceDone(this);
            } else {
                blockComplete.clear();
            }

        } catch (FileNotFoundException e) {
            //Do nothing
        } catch (IOException e) {
            LOGGER.error("IO Error", e);
        } finally {
            buffer = null;
        }
    }

    public void queueWrite(final ByteBuf data, final int offset) {
        torrentEngine.getWriteThreadPool().execute(new Runnable() {
            @Override
            public void run() {
                writeSync(data, offset);
            }
        });
    }

    private void writeSync(ByteBuf data, int offset) {
        if (isComplete()) {
            return;
        }

        if ((data.readableBytes() + offset) > pieceLength) {
            LOGGER.error("Received block is out of bounds");
            return; //TODO exception
        }

        if (buffer == null) {
            buffer = Unpooled.buffer(pieceLength, pieceLength);
        }

        buffer.writerIndex(offset);
        data.getBytes(0, buffer, data.readableBytes());

        int dataBlocks = (int) Math.ceil((double) data.readableBytes() / (double) DEFAULT_BLOCK_LENGTH);
        int blockOffset = offset / DEFAULT_BLOCK_LENGTH;

        data.release();

        blockComplete.set(blockOffset, blockOffset + dataBlocks);

        if (isAllBlocksComplete()) {
            try {
                byte[] pieceHash = DigestUtils.sha1(buffer.array());

                if (Arrays.equals(pieceHash, hash)) {
                    torrent.getFileStorage().store(index, buffer);

                    complete = true;
                    softBuffer = new SoftReference<>(buffer);
                    torrentEngine.setPieceDone(Piece.this);
                    torrentEngine.increaseDownloadedCounter(getPieceLength());
                } else {
                    LOGGER.debug("Piece index: " + index + " hash invalid. Original:" + Hex.encodeHexString(hash)
                            + " actual: " + Hex.encodeHexString(pieceHash));
                    blockComplete.clear();
                }
            } catch (IOException e) {
                LOGGER.error("Can't save piece", e);
            } finally {
                buffer = null;
            }
        }
    }

    public ByteBuf read(int offset, int length) throws IOException {
        //LOGGER.debug("READ is complete = " + isComplete());
        if (isComplete()) {
            ByteBuf bb = (softBuffer != null) ? softBuffer.get() : null;
            if (bb == null) {
                bb = torrent.getFileStorage().read(index, pieceLength);
                softBuffer = new SoftReference<>(bb);
                torrent.missCache();
            } else {
                torrent.hitCache();
            }

            ByteBuf buf = Unpooled.buffer(length, length);
            bb.readerIndex(offset);
            bb.readBytes(buf);
            return buf;
        } else {
            return Unpooled.EMPTY_BUFFER;
        }
    }

    private boolean isAllBlocksComplete() {
        return blockComplete.cardinality() == getBlockCount();
    }

    public void increasePeerCount() {
        peerCount++;
    }

    public boolean isComplete() {
        return complete;
    }

    public int getBlockCount() {
        return blockCount;
    }

    public int getIndex() {
        return index;
    }

    @Override
    public int compareTo(@Nonnull Piece otherPiece) {
        if (isComplete() == otherPiece.isComplete()) {
            return peerCount - otherPiece.peerCount;
        } else {
            if (isComplete()) {
                return 1;
            } else {
                return -1;
            }
        }
    }

    @Override
    public boolean equals(@Nullable Object o) {
        if (this == o) {
            return true;
        }
        if (o == null || getClass() != o.getClass()) {
            return false;
        }

        Piece piece = (Piece) o;

        return index == piece.getIndex() && pieceLength == piece.getPieceLength()
                && Arrays.equals(hash, piece.hash);
    }

    @Override
    public int hashCode() {
        int result = index;
        result = 31 * result + pieceLength;
        result = 31 * result + Arrays.hashCode(hash);
        return result;
    }

    public int getPieceLength() {
        return pieceLength;
    }
}