piece.go

Documentation: github.com/concurrency-8/piece

     1  package piece
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  
     7  	"github.com/concurrency-8/parser"
     8  )
     9  
    10  // PieceTracker stores flags for blocks of pieces requested and received
    11  // Requested[i][j] = true => jth block of ith piece has been requested
    12  type PieceTracker struct {
    13  	Torrent   parser.TorrentFile
    14  	Requested [][]bool
    15  	Received  [][]bool
    16  	lock      sync.Mutex
    17  }
    18  
    19  // NewPieceTracker returns a new PieceTracker object for the torrent
    20  func NewPieceTracker(torrent parser.TorrentFile) (tracker *PieceTracker) {
    21  	tracker = new(PieceTracker)
    22  	tracker.Torrent = torrent
    23  	numPieces := uint32(len(torrent.Piece) / 20)
    24  	fmt.Println("numPieces: ", numPieces)
    25  	for i := uint32(0); i < numPieces; i++ {
    26  		blocksPerPiece, _ := parser.BlocksPerPiece(torrent, i)
    27  		tracker.Requested = append(tracker.Requested, make([]bool, blocksPerPiece))
    28  		tracker.Received = append(tracker.Received, make([]bool, blocksPerPiece))
    29  	}
    30  	return
    31  }
    32  
    33  // AddRequested flags the request value of a block in a piece
    34  // Invoked while requesting the block of a piece
    35  func (tracker *PieceTracker) AddRequested(block parser.PieceBlock) {
    36  	index := block.Begin / parser.BLOCK_LEN
    37  	tracker.Requested[block.Index][index] = true
    38  }
    39  
    40  // AddReceived flags the received value of a block in a piece
    41  // Invoked when a block is received
    42  func (tracker *PieceTracker) AddReceived(block parser.PieceBlock) {
    43  	index := block.Begin / parser.BLOCK_LEN
    44  	tracker.Received[block.Index][index] = true
    45  }
    46  
    47  // Needed checks if we want a block. If we have already requested all,
    48  // we reset requested to be equal to received and request the remaining pieces
    49  func (tracker *PieceTracker) Needed(block parser.PieceBlock) bool {
    50  
    51  	// Check if all have been requested...
    52  	allRequested := true
    53  	for _, i := range tracker.Requested {
    54  		if !allRequested {
    55  			break
    56  		}
    57  		for _, j := range i {
    58  			allRequested = allRequested && j
    59  		}
    60  	}
    61  
    62  	// If yes, copy received into request...
    63  	if allRequested {
    64  		tracker.lock.Lock()
    65  		tracker.Requested = clone(tracker.Received)
    66  		tracker.lock.Unlock()
    67  	}
    68  
    69  	return !tracker.Requested[block.Index][block.Begin/parser.BLOCK_LEN]
    70  }
    71  
    72  // Deep clones 2-D bool array
    73  func clone(array [][]bool) (result [][]bool) {
    74  	for _, i := range array {
    75  		temp := make([]bool, len(i))
    76  		for index, j := range i {
    77  			temp[index] = j
    78  		}
    79  		result = append(result, temp)
    80  	}
    81  	return
    82  }
    83  
    84  // PieceIsDone tells if the pieceIndex piece has been downloaded successfully
    85  func (tracker *PieceTracker) PieceIsDone(pieceIndex uint32) (result bool) {
    86  	result = true
    87  	for _, i := range tracker.Received[pieceIndex] {
    88  		result = result && i
    89  	}
    90  	return
    91  }
    92  
    93  // IsDone tells if the torrent file has been successfully received
    94  func (tracker *PieceTracker) IsDone() (result bool) {
    95  	result = true
    96  	for _, i := range tracker.Received {
    97  		for _, j := range i {
    98  			result = result && j
    99  		}
   100  	}
   101  	return
   102  }
   103  
   104  // PrintPercentageDone prints the percentage of download completed on the screen
   105  func (tracker *PieceTracker) PrintPercentageDone() {
   106  	downloaded, total := 0.0, 0
   107  	for _, i := range tracker.Received {
   108  		for _, j := range i {
   109  			total++
   110  			if j {
   111  				downloaded++
   112  			}
   113  		}
   114  	}
   115  	percent := float64(downloaded*100) / float64(total)
   116  	fmt.Print("progress:", percent, "\r")
   117  }
   118  
   119  // Reset the piece - Called when invalid SHA
   120  func (tracker *PieceTracker) Reset(index uint32) {
   121  	tracker.lock.Lock()
   122  	for i := range tracker.Requested[index] {
   123  		tracker.Requested[index][i] = false
   124  		tracker.Received[index][i] = false
   125  	}
   126  	tracker.lock.Unlock()
   127  }
   128  
   129  // Fill is used to revive the piecetracker while resuming the torrent
   130  func (tracker *PieceTracker) Fill(index uint32) {
   131  	for i := range tracker.Requested[index] {
   132  		tracker.Requested[index][i] = true
   133  		tracker.Received[index][i] = true
   134  	}
   135  }
   136
View as plain text