495 lines
15 KiB
Go
495 lines
15 KiB
Go
package torrentutil
|
|
|
|
import (
|
|
"fmt"
|
|
"io"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/anacrolix/torrent"
|
|
"github.com/rs/zerolog"
|
|
)
|
|
|
|
// +-----------------------+
|
|
// + anacrolix/torrent +
|
|
// +-----------------------+
|
|
|
|
const (
|
|
piecesForNow = int64(5)
|
|
piecesForHighBefore = int64(2)
|
|
piecesForNext = int64(30)
|
|
piecesForReadahead = int64(30)
|
|
)
|
|
|
|
// readerInfo tracks information about an active reader
|
|
type readerInfo struct {
|
|
id string
|
|
position int64
|
|
lastAccess time.Time
|
|
}
|
|
|
|
// priorityManager manages piece priorities for multiple readers on the same file
|
|
type priorityManager struct {
|
|
mu sync.RWMutex
|
|
readers map[string]*readerInfo
|
|
torrent *torrent.Torrent
|
|
file *torrent.File
|
|
logger *zerolog.Logger
|
|
}
|
|
|
|
// global map to track priority managers per torrent+file combination
|
|
var (
|
|
priorityManagers = make(map[string]*priorityManager)
|
|
priorityManagersMu sync.RWMutex
|
|
)
|
|
|
|
// getPriorityManager gets or creates a priority manager for a torrent+file combination
|
|
func getPriorityManager(t *torrent.Torrent, file *torrent.File, logger *zerolog.Logger) *priorityManager {
|
|
key := fmt.Sprintf("%s:%s", t.InfoHash().String(), file.Path())
|
|
|
|
priorityManagersMu.Lock()
|
|
defer priorityManagersMu.Unlock()
|
|
|
|
if pm, exists := priorityManagers[key]; exists {
|
|
return pm
|
|
}
|
|
|
|
pm := &priorityManager{
|
|
readers: make(map[string]*readerInfo),
|
|
torrent: t,
|
|
file: file,
|
|
logger: logger,
|
|
}
|
|
priorityManagers[key] = pm
|
|
|
|
// Start cleanup goroutine for the first manager
|
|
if len(priorityManagers) == 1 {
|
|
go pm.cleanupStaleReaders()
|
|
}
|
|
|
|
return pm
|
|
}
|
|
|
|
// registerReader registers a new reader with the priority manager
|
|
func (pm *priorityManager) registerReader(readerID string, position int64) {
|
|
pm.mu.Lock()
|
|
defer pm.mu.Unlock()
|
|
|
|
pm.readers[readerID] = &readerInfo{
|
|
id: readerID,
|
|
position: position,
|
|
lastAccess: time.Now(),
|
|
}
|
|
|
|
pm.updatePriorities()
|
|
}
|
|
|
|
// updateReaderPosition updates a reader's position and recalculates priorities
|
|
func (pm *priorityManager) updateReaderPosition(readerID string, position int64) {
|
|
pm.mu.Lock()
|
|
defer pm.mu.Unlock()
|
|
|
|
if reader, exists := pm.readers[readerID]; exists {
|
|
reader.position = position
|
|
reader.lastAccess = time.Now()
|
|
pm.updatePriorities()
|
|
}
|
|
}
|
|
|
|
// unregisterReader removes a reader from tracking
|
|
func (pm *priorityManager) unregisterReader(readerID string) {
|
|
pm.mu.Lock()
|
|
defer pm.mu.Unlock()
|
|
|
|
delete(pm.readers, readerID)
|
|
|
|
// If no more readers, clean up and recalculate priorities
|
|
if len(pm.readers) == 0 {
|
|
pm.resetAllPriorities()
|
|
} else {
|
|
pm.updatePriorities()
|
|
}
|
|
}
|
|
|
|
// updatePriorities recalculates piece priorities based on all active readers
|
|
func (pm *priorityManager) updatePriorities() {
|
|
if pm.torrent == nil || pm.file == nil || pm.torrent.Info() == nil {
|
|
return
|
|
}
|
|
|
|
t := pm.torrent
|
|
file := pm.file
|
|
pieceLength := t.Info().PieceLength
|
|
|
|
if pieceLength == 0 {
|
|
if pm.logger != nil {
|
|
pm.logger.Warn().Msg("torrentutil: piece length is zero, cannot prioritize")
|
|
}
|
|
return
|
|
}
|
|
|
|
numTorrentPieces := int64(t.NumPieces())
|
|
if numTorrentPieces == 0 {
|
|
if pm.logger != nil {
|
|
pm.logger.Warn().Msg("torrentutil: torrent has zero pieces, cannot prioritize")
|
|
}
|
|
return
|
|
}
|
|
|
|
// Calculate file piece range
|
|
fileFirstPieceIdx := file.Offset() / pieceLength
|
|
fileLastPieceIdx := (file.Offset() + file.Length() - 1) / pieceLength
|
|
|
|
// Collect all needed piece ranges from all active readers
|
|
neededPieces := make(map[int64]torrent.PiecePriority)
|
|
|
|
for _, reader := range pm.readers {
|
|
position := reader.position
|
|
// Remove 1MB from the position (for subtitle cluster)
|
|
position -= 1 * 1024 * 1024
|
|
if position < 0 {
|
|
position = 0
|
|
}
|
|
if position < 0 {
|
|
position = 0
|
|
}
|
|
|
|
currentGlobalSeekPieceIdx := (file.Offset() + position) / pieceLength
|
|
|
|
// Pieces needed NOW (immediate)
|
|
for i := int64(0); i < piecesForNow; i++ {
|
|
idx := currentGlobalSeekPieceIdx + i
|
|
if idx >= fileFirstPieceIdx && idx <= fileLastPieceIdx && idx < numTorrentPieces {
|
|
if current, exists := neededPieces[idx]; !exists || current < torrent.PiecePriorityNow {
|
|
neededPieces[idx] = torrent.PiecePriorityNow
|
|
}
|
|
}
|
|
}
|
|
|
|
// Pieces needed HIGH (before current position for rewinds)
|
|
for i := int64(1); i <= piecesForHighBefore; i++ {
|
|
idx := currentGlobalSeekPieceIdx - i
|
|
if idx >= fileFirstPieceIdx && idx <= fileLastPieceIdx && idx >= 0 {
|
|
if current, exists := neededPieces[idx]; !exists || current < torrent.PiecePriorityHigh {
|
|
neededPieces[idx] = torrent.PiecePriorityHigh
|
|
}
|
|
}
|
|
}
|
|
|
|
// Pieces needed NEXT (immediate readahead)
|
|
nextStartIdx := currentGlobalSeekPieceIdx + piecesForNow
|
|
for i := int64(0); i < piecesForNext; i++ {
|
|
idx := nextStartIdx + i
|
|
if idx >= fileFirstPieceIdx && idx <= fileLastPieceIdx && idx < numTorrentPieces {
|
|
if current, exists := neededPieces[idx]; !exists || current < torrent.PiecePriorityNext {
|
|
neededPieces[idx] = torrent.PiecePriorityNext
|
|
}
|
|
}
|
|
}
|
|
|
|
// Pieces needed for READAHEAD (further readahead)
|
|
readaheadStartIdx := nextStartIdx + piecesForNext
|
|
for i := int64(0); i < piecesForReadahead; i++ {
|
|
idx := readaheadStartIdx + i
|
|
if idx >= fileFirstPieceIdx && idx <= fileLastPieceIdx && idx < numTorrentPieces {
|
|
if current, exists := neededPieces[idx]; !exists || current < torrent.PiecePriorityReadahead {
|
|
neededPieces[idx] = torrent.PiecePriorityReadahead
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Reset pieces that are no longer needed by any reader
|
|
for idx := fileFirstPieceIdx; idx <= fileLastPieceIdx; idx++ {
|
|
if idx < 0 || idx >= numTorrentPieces {
|
|
continue
|
|
}
|
|
|
|
piece := t.Piece(int(idx))
|
|
currentPriority := piece.State().Priority
|
|
|
|
if neededPriority, needed := neededPieces[idx]; needed {
|
|
// Set to the highest priority needed by any reader
|
|
if currentPriority != neededPriority {
|
|
piece.SetPriority(neededPriority)
|
|
}
|
|
} else {
|
|
// Only reset to normal if not completely unwanted and not already at highest priority
|
|
if currentPriority != torrent.PiecePriorityNone && currentPriority != torrent.PiecePriorityNow {
|
|
piece.SetPriority(torrent.PiecePriorityNormal)
|
|
}
|
|
}
|
|
}
|
|
|
|
if pm.logger != nil {
|
|
pm.logger.Debug().Msgf("torrentutil: Updated priorities for %d readers, %d pieces prioritized", len(pm.readers), len(neededPieces))
|
|
}
|
|
}
|
|
|
|
// resetAllPriorities resets all file pieces to normal priority
|
|
func (pm *priorityManager) resetAllPriorities() {
|
|
if pm.torrent == nil || pm.file == nil || pm.torrent.Info() == nil {
|
|
return
|
|
}
|
|
|
|
t := pm.torrent
|
|
file := pm.file
|
|
pieceLength := t.Info().PieceLength
|
|
|
|
if pieceLength == 0 {
|
|
return
|
|
}
|
|
|
|
numTorrentPieces := int64(t.NumPieces())
|
|
fileFirstPieceIdx := file.Offset() / pieceLength
|
|
fileLastPieceIdx := (file.Offset() + file.Length() - 1) / pieceLength
|
|
|
|
for idx := fileFirstPieceIdx; idx <= fileLastPieceIdx; idx++ {
|
|
if idx >= 0 && idx < numTorrentPieces {
|
|
piece := t.Piece(int(idx))
|
|
if piece.State().Priority != torrent.PiecePriorityNone {
|
|
piece.SetPriority(torrent.PiecePriorityNormal)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// cleanupStaleReaders periodically removes readers that haven't been accessed recently
|
|
func (pm *priorityManager) cleanupStaleReaders() {
|
|
ticker := time.NewTicker(30 * time.Second)
|
|
defer ticker.Stop()
|
|
|
|
for range ticker.C {
|
|
pm.mu.Lock()
|
|
cutoff := time.Now().Add(-2 * time.Minute)
|
|
|
|
for id, reader := range pm.readers {
|
|
if reader.lastAccess.Before(cutoff) {
|
|
delete(pm.readers, id)
|
|
if pm.logger != nil {
|
|
pm.logger.Debug().Msgf("torrentutil: Cleaned up stale reader %s", id)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Update priorities after cleanup
|
|
if len(pm.readers) > 0 {
|
|
pm.updatePriorities()
|
|
}
|
|
|
|
pm.mu.Unlock()
|
|
}
|
|
}
|
|
|
|
// ReadSeeker implements io.ReadSeekCloser for a torrent file being streamed.
|
|
// It allows dynamic prioritization of pieces when seeking, optimized for streaming
|
|
// and supports multiple concurrent readers on the same file.
|
|
type ReadSeeker struct {
|
|
id string
|
|
torrent *torrent.Torrent
|
|
file *torrent.File
|
|
reader torrent.Reader
|
|
priorityManager *priorityManager
|
|
logger *zerolog.Logger
|
|
}
|
|
|
|
var _ io.ReadSeekCloser = &ReadSeeker{}
|
|
|
|
func NewReadSeeker(t *torrent.Torrent, file *torrent.File, logger ...*zerolog.Logger) io.ReadSeekCloser {
|
|
tr := file.NewReader()
|
|
tr.SetResponsive()
|
|
// Read ahead 5MB for better streaming performance
|
|
// DEVNOTE: Not sure if dynamic prioritization overwrites this but whatever
|
|
tr.SetReadahead(5 * 1024 * 1024)
|
|
|
|
var loggerPtr *zerolog.Logger
|
|
if len(logger) > 0 {
|
|
loggerPtr = logger[0]
|
|
}
|
|
|
|
pm := getPriorityManager(t, file, loggerPtr)
|
|
|
|
rs := &ReadSeeker{
|
|
id: fmt.Sprintf("reader_%d_%d", time.Now().UnixNano(), len(pm.readers)),
|
|
torrent: t,
|
|
file: file,
|
|
reader: tr,
|
|
priorityManager: pm,
|
|
logger: loggerPtr,
|
|
}
|
|
|
|
// Register this reader with the priority manager
|
|
pm.registerReader(rs.id, 0)
|
|
|
|
return rs
|
|
}
|
|
|
|
func (rs *ReadSeeker) Read(p []byte) (n int, err error) {
|
|
return rs.reader.Read(p)
|
|
}
|
|
|
|
func (rs *ReadSeeker) Seek(offset int64, whence int) (int64, error) {
|
|
newOffset, err := rs.reader.Seek(offset, whence)
|
|
if err != nil {
|
|
if rs.logger != nil {
|
|
rs.logger.Error().Err(err).Int64("offset", offset).Int("whence", whence).Msg("torrentutil: ReadSeeker seek error")
|
|
}
|
|
return newOffset, err
|
|
}
|
|
|
|
// Update this reader's position in the priority manager
|
|
rs.priorityManager.updateReaderPosition(rs.id, newOffset)
|
|
|
|
return newOffset, nil
|
|
}
|
|
|
|
// Close closes the underlying torrent file reader and unregisters from priority manager.
|
|
// This makes ReadSeeker implement io.ReadSeekCloser.
|
|
func (rs *ReadSeeker) Close() error {
|
|
// Unregister from priority manager
|
|
rs.priorityManager.unregisterReader(rs.id)
|
|
|
|
if rs.reader != nil {
|
|
return rs.reader.Close()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// PrioritizeDownloadPieces sets high priority for the first 3% of pieces and the last few pieces to ensure faster loading.
|
|
func PrioritizeDownloadPieces(t *torrent.Torrent, file *torrent.File, logger *zerolog.Logger) {
|
|
// Calculate file's pieces
|
|
firstPieceIdx := file.Offset() * int64(t.NumPieces()) / t.Length()
|
|
endPieceIdx := (file.Offset() + file.Length()) * int64(t.NumPieces()) / t.Length()
|
|
|
|
// Prioritize more pieces at the beginning for faster initial loading (3% for beginning)
|
|
numPiecesForStart := (endPieceIdx - firstPieceIdx + 1) * 3 / 100
|
|
if logger != nil {
|
|
logger.Debug().Msgf("torrentuil: Setting high priority for first 3%% - pieces %d to %d (total %d)",
|
|
firstPieceIdx, firstPieceIdx+numPiecesForStart, numPiecesForStart)
|
|
}
|
|
for idx := firstPieceIdx; idx <= firstPieceIdx+numPiecesForStart; idx++ {
|
|
t.Piece(int(idx)).SetPriority(torrent.PiecePriorityNow)
|
|
}
|
|
|
|
// Also prioritize the last few pieces
|
|
numPiecesForEnd := (endPieceIdx - firstPieceIdx + 1) * 1 / 100
|
|
if logger != nil {
|
|
logger.Debug().Msgf("torrentuil: Setting priority for last pieces %d to %d (total %d)",
|
|
endPieceIdx-numPiecesForEnd, endPieceIdx, numPiecesForEnd)
|
|
}
|
|
for idx := endPieceIdx - numPiecesForEnd; idx <= endPieceIdx; idx++ {
|
|
if idx >= 0 && int(idx) < t.NumPieces() {
|
|
t.Piece(int(idx)).SetPriority(torrent.PiecePriorityNow)
|
|
}
|
|
}
|
|
}
|
|
|
|
// PrioritizeRangeRequestPieces attempts to prioritize pieces needed for the range request.
|
|
func PrioritizeRangeRequestPieces(rangeHeader string, t *torrent.Torrent, file *torrent.File, logger *zerolog.Logger) {
|
|
// Parse the range header (format: bytes=START-END)
|
|
var start int64
|
|
_, _ = fmt.Sscanf(rangeHeader, "bytes=%d-", &start)
|
|
|
|
if start >= 0 {
|
|
// Calculate file's pieces range
|
|
fileOffset := file.Offset()
|
|
fileLength := file.Length()
|
|
|
|
// Calculate the total range of pieces for this file
|
|
firstFilePieceIdx := fileOffset * int64(t.NumPieces()) / t.Length()
|
|
endFilePieceIdx := (fileOffset + fileLength) * int64(t.NumPieces()) / t.Length()
|
|
|
|
// Calculate the piece index for this seek offset with small padding
|
|
// Subtract a small amount to ensure we don't miss the beginning of a needed piece
|
|
seekPosition := start
|
|
if seekPosition >= 1024*1024 { // If we're at least 1MB in, add some padding
|
|
seekPosition -= 1024 * 512 // Subtract 512KB to ensure we get the right piece
|
|
}
|
|
seekPieceIdx := (fileOffset + seekPosition) * int64(t.NumPieces()) / t.Length()
|
|
|
|
// Prioritize the next several pieces from this point
|
|
// This is especially important for seeking
|
|
numPiecesToPrioritize := int64(10) // Prioritize next 10 pieces, adjust as needed
|
|
|
|
if seekPieceIdx+numPiecesToPrioritize > endFilePieceIdx {
|
|
numPiecesToPrioritize = endFilePieceIdx - seekPieceIdx
|
|
}
|
|
|
|
if logger != nil {
|
|
logger.Debug().Msgf("torrentutil: Prioritizing range request pieces %d to %d",
|
|
seekPieceIdx, seekPieceIdx+numPiecesToPrioritize)
|
|
}
|
|
|
|
// Set normal priority for pieces far from our current position
|
|
// This allows background downloading while still prioritizing the seek point
|
|
for idx := firstFilePieceIdx; idx <= endFilePieceIdx; idx++ {
|
|
if idx >= 0 && int(idx) < t.NumPieces() {
|
|
// Don't touch the beginning pieces which should maintain their high priority
|
|
// for the next potential restart, and don't touch pieces near our seek point
|
|
if idx > firstFilePieceIdx+100 && idx < seekPieceIdx-100 ||
|
|
idx > seekPieceIdx+numPiecesToPrioritize+100 {
|
|
// Set to normal priority - allow background downloading
|
|
t.Piece(int(idx)).SetPriority(torrent.PiecePriorityNormal)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Now set the highest priority for the pieces we need right now
|
|
for idx := seekPieceIdx; idx < seekPieceIdx+numPiecesToPrioritize; idx++ {
|
|
if idx >= 0 && int(idx) < t.NumPieces() {
|
|
t.Piece(int(idx)).SetPriority(torrent.PiecePriorityNow)
|
|
}
|
|
}
|
|
|
|
// Also prioritize a small buffer before the seek point to handle small rewinds
|
|
// This is useful for MPV's default rewind behavior
|
|
bufferBeforeCount := int64(5) // 5 pieces buffer before seek point
|
|
if seekPieceIdx > firstFilePieceIdx+bufferBeforeCount {
|
|
for idx := seekPieceIdx - bufferBeforeCount; idx < seekPieceIdx; idx++ {
|
|
if idx >= 0 && int(idx) < t.NumPieces() {
|
|
t.Piece(int(idx)).SetPriority(torrent.PiecePriorityHigh)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Also prioritize the next readahead segment after our immediate needs
|
|
// This helps prepare for continued playback
|
|
nextReadStart := seekPieceIdx + numPiecesToPrioritize
|
|
nextReadCount := int64(100) // 100 additional pieces for nextRead
|
|
if nextReadStart+nextReadCount > endFilePieceIdx {
|
|
nextReadCount = endFilePieceIdx - nextReadStart
|
|
}
|
|
|
|
if nextReadCount > 0 {
|
|
if logger != nil {
|
|
logger.Debug().Msgf("torrentutil: Setting next priority for pieces %d to %d",
|
|
nextReadStart, nextReadStart+nextReadCount)
|
|
}
|
|
for idx := nextReadStart; idx < nextReadStart+nextReadCount; idx++ {
|
|
if idx >= 0 && int(idx) < t.NumPieces() {
|
|
t.Piece(int(idx)).SetPriority(torrent.PiecePriorityNext)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Also prioritize the next readahead segment after our immediate needs
|
|
// This helps prepare for continued playback
|
|
readAheadCount := int64(100)
|
|
if nextReadStart+readAheadCount > endFilePieceIdx {
|
|
readAheadCount = endFilePieceIdx - nextReadStart
|
|
}
|
|
|
|
if readAheadCount > 0 {
|
|
if logger != nil {
|
|
logger.Debug().Msgf("torrentutil: Setting read ahead priority for pieces %d to %d",
|
|
nextReadStart, nextReadStart+readAheadCount)
|
|
}
|
|
for idx := nextReadStart; idx < nextReadStart+readAheadCount; idx++ {
|
|
if idx >= 0 && int(idx) < t.NumPieces() {
|
|
t.Piece(int(idx)).SetPriority(torrent.PiecePriorityReadahead)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|