477 lines
11 KiB
Go
477 lines
11 KiB
Go
package torrent
|
||
|
||
import (
|
||
"context"
|
||
"sync"
|
||
)
|
||
|
||
type pieceTask struct {
|
||
Index int
|
||
}
|
||
|
||
type assignPieceRequest struct {
|
||
peerID string
|
||
have []bool
|
||
hasInfo bool
|
||
responseCh chan assignPieceResponse
|
||
}
|
||
|
||
type assignPieceResponse struct {
|
||
task pieceTask
|
||
ok bool
|
||
}
|
||
|
||
type reportPieceRequest struct {
|
||
pieceIndex int
|
||
success bool
|
||
responseCh chan bool
|
||
}
|
||
|
||
// cancelPeerRequest отправляет сигнал конкретному пиру отменить кусок в endgame.
|
||
type cancelPeerRequest struct {
|
||
pieceIndex int
|
||
peerID string
|
||
}
|
||
|
||
type pieceScheduler struct {
|
||
assignCh chan assignPieceRequest
|
||
reportCh chan reportPieceRequest
|
||
releasePeerCh chan string
|
||
progressCh chan int
|
||
doneCh chan struct{}
|
||
stopCh chan struct{}
|
||
stopOnce sync.Once
|
||
|
||
setSequentialCh chan bool
|
||
|
||
// cancelCh рассылает сигнал отмены конкретному воркеру в endgame-режиме.
|
||
// Ключ — peerID, значение — канал с индексом куска для отмены.
|
||
cancelMu sync.RWMutex
|
||
cancelSubs map[string]chan int
|
||
}
|
||
|
||
type pieceState uint8
|
||
|
||
const (
|
||
piecePending pieceState = iota
|
||
pieceInProgress // качается одним пиром
|
||
pieceDone
|
||
)
|
||
|
||
// endgameThreshold — количество оставшихся кусков, при котором включается endgame.
|
||
const endgameThreshold = 4
|
||
|
||
func newPieceScheduler(pieceCount int) *pieceScheduler {
|
||
return newPieceSchedulerWithResume(pieceCount, nil)
|
||
}
|
||
|
||
func newPieceSchedulerWithResume(pieceCount int, completedIndices []int) *pieceScheduler {
|
||
ps := &pieceScheduler{
|
||
assignCh: make(chan assignPieceRequest, 128),
|
||
reportCh: make(chan reportPieceRequest, 128),
|
||
releasePeerCh: make(chan string, 128),
|
||
progressCh: make(chan int, 1),
|
||
doneCh: make(chan struct{}),
|
||
stopCh: make(chan struct{}),
|
||
setSequentialCh: make(chan bool),
|
||
cancelSubs: make(map[string]chan int),
|
||
}
|
||
|
||
go ps.run(pieceCount, completedIndices)
|
||
return ps
|
||
}
|
||
|
||
// SubscribeCancel регистрирует канал для получения cancel-сигналов в endgame-режиме.
|
||
// Воркер должен вызвать это перед Acquire, и отписаться после завершения.
|
||
func (ps *pieceScheduler) SubscribeCancel(peerID string) <-chan int {
|
||
ch := make(chan int, 8)
|
||
ps.cancelMu.Lock()
|
||
ps.cancelSubs[peerID] = ch
|
||
ps.cancelMu.Unlock()
|
||
return ch
|
||
}
|
||
|
||
// UnsubscribeCancel отписывает воркера от cancel-сигналов.
|
||
func (ps *pieceScheduler) UnsubscribeCancel(peerID string) {
|
||
ps.cancelMu.Lock()
|
||
delete(ps.cancelSubs, peerID)
|
||
ps.cancelMu.Unlock()
|
||
}
|
||
|
||
// broadcastCancel рассылает сигнал отмены куска всем воркерам кроме победителя.
|
||
func (ps *pieceScheduler) broadcastCancel(pieceIndex int, winnerPeerID string) {
|
||
ps.cancelMu.RLock()
|
||
defer ps.cancelMu.RUnlock()
|
||
for id, ch := range ps.cancelSubs {
|
||
if id == winnerPeerID {
|
||
continue
|
||
}
|
||
select {
|
||
case ch <- pieceIndex:
|
||
default:
|
||
}
|
||
}
|
||
}
|
||
|
||
func (ps *pieceScheduler) Acquire(ctx context.Context, peerID string, have []bool, hasInfo bool) (pieceTask, bool, error) {
|
||
responseCh := make(chan assignPieceResponse, 1)
|
||
req := assignPieceRequest{
|
||
peerID: peerID,
|
||
have: have,
|
||
hasInfo: hasInfo,
|
||
responseCh: responseCh,
|
||
}
|
||
|
||
select {
|
||
case <-ctx.Done():
|
||
return pieceTask{}, false, ctx.Err()
|
||
case <-ps.doneCh:
|
||
return pieceTask{}, false, nil
|
||
case ps.assignCh <- req:
|
||
}
|
||
|
||
select {
|
||
case <-ctx.Done():
|
||
return pieceTask{}, false, ctx.Err()
|
||
case <-ps.doneCh:
|
||
return pieceTask{}, false, nil
|
||
case response := <-responseCh:
|
||
return response.task, response.ok, nil
|
||
}
|
||
}
|
||
|
||
func (ps *pieceScheduler) Report(ctx context.Context, pieceIndex int, success bool) (bool, error) {
|
||
responseCh := make(chan bool, 1)
|
||
req := reportPieceRequest{
|
||
pieceIndex: pieceIndex,
|
||
success: success,
|
||
responseCh: responseCh,
|
||
}
|
||
|
||
select {
|
||
case <-ctx.Done():
|
||
return false, ctx.Err()
|
||
case <-ps.doneCh:
|
||
return false, nil
|
||
case ps.reportCh <- req:
|
||
}
|
||
|
||
select {
|
||
case <-ctx.Done():
|
||
return false, ctx.Err()
|
||
case <-ps.doneCh:
|
||
return false, nil
|
||
case accepted := <-responseCh:
|
||
return accepted, nil
|
||
}
|
||
}
|
||
|
||
func (ps *pieceScheduler) ReleasePeer(peerID string) {
|
||
if peerID == "" {
|
||
return
|
||
}
|
||
select {
|
||
case <-ps.doneCh:
|
||
return
|
||
case ps.releasePeerCh <- peerID:
|
||
default:
|
||
}
|
||
}
|
||
|
||
func (ps *pieceScheduler) Progress() <-chan int {
|
||
return ps.progressCh
|
||
}
|
||
|
||
func (ps *pieceScheduler) Done() <-chan struct{} {
|
||
return ps.doneCh
|
||
}
|
||
|
||
func (ps *pieceScheduler) Stop() {
|
||
ps.stopOnce.Do(func() {
|
||
close(ps.stopCh)
|
||
})
|
||
}
|
||
|
||
func (ps *pieceScheduler) SetSequential(mode bool) {
|
||
select {
|
||
case ps.setSequentialCh <- mode:
|
||
case <-ps.stopCh:
|
||
}
|
||
}
|
||
|
||
func (ps *pieceScheduler) run(pieceCount int, completedIndices []int) {
|
||
states := make([]pieceState, pieceCount)
|
||
availability := make([]int, pieceCount)
|
||
peerAvailability := make(map[string][]bool, 256)
|
||
endgamePeers := make(map[int][]string)
|
||
completed := 0
|
||
sequentialMode := false
|
||
|
||
for _, idx := range completedIndices {
|
||
if idx >= 0 && idx < pieceCount {
|
||
states[idx] = pieceDone
|
||
completed++
|
||
}
|
||
}
|
||
|
||
finish := func() {
|
||
close(ps.doneCh)
|
||
close(ps.progressCh)
|
||
}
|
||
|
||
if pieceCount == 0 {
|
||
finish()
|
||
return
|
||
}
|
||
|
||
for {
|
||
if completed >= pieceCount {
|
||
finish()
|
||
return
|
||
}
|
||
|
||
remaining := pieceCount - completed
|
||
endgame := remaining <= endgameThreshold
|
||
|
||
select {
|
||
case <-ps.stopCh:
|
||
finish()
|
||
return
|
||
case peerID := <-ps.releasePeerCh:
|
||
releasePeerAvailability(peerAvailability, availability, peerID)
|
||
case req := <-ps.assignCh:
|
||
if req.peerID != "" && req.hasInfo {
|
||
updatePeerAvailability(peerAvailability, availability, req.peerID, req.have, pieceCount)
|
||
}
|
||
|
||
pieceIndex := -1
|
||
if endgame {
|
||
pieceIndex = selectPieceEndgame(states, availability, req.have, req.hasInfo, endgamePeers, req.peerID)
|
||
} else {
|
||
if sequentialMode {
|
||
pieceIndex = selectPendingPieceSequential(states, req.have, req.hasInfo)
|
||
} else {
|
||
pieceIndex = selectPendingPieceRarest(states, availability, req.have, req.hasInfo)
|
||
}
|
||
}
|
||
|
||
if pieceIndex >= 0 {
|
||
if !endgame {
|
||
states[pieceIndex] = pieceInProgress
|
||
}
|
||
endgamePeers[pieceIndex] = append(endgamePeers[pieceIndex], req.peerID)
|
||
req.responseCh <- assignPieceResponse{
|
||
task: pieceTask{Index: pieceIndex},
|
||
ok: true,
|
||
}
|
||
continue
|
||
}
|
||
req.responseCh <- assignPieceResponse{ok: false}
|
||
|
||
case req := <-ps.reportCh:
|
||
if req.pieceIndex < 0 || req.pieceIndex >= pieceCount {
|
||
req.responseCh <- false
|
||
continue
|
||
}
|
||
|
||
if req.success {
|
||
if states[req.pieceIndex] == pieceDone {
|
||
req.responseCh <- false
|
||
continue
|
||
}
|
||
states[req.pieceIndex] = pieceDone
|
||
completed++
|
||
|
||
// Endgame: уведомить других пиров отменить этот кусок
|
||
if peers, ok := endgamePeers[req.pieceIndex]; ok && len(peers) > 1 {
|
||
go ps.broadcastCancel(req.pieceIndex, "")
|
||
}
|
||
delete(endgamePeers, req.pieceIndex)
|
||
|
||
select {
|
||
case ps.progressCh <- req.pieceIndex:
|
||
default:
|
||
}
|
||
} else {
|
||
if states[req.pieceIndex] != pieceDone {
|
||
states[req.pieceIndex] = piecePending
|
||
}
|
||
}
|
||
req.responseCh <- true
|
||
|
||
case mode := <-ps.setSequentialCh:
|
||
sequentialMode = mode
|
||
}
|
||
}
|
||
}
|
||
|
||
// selectPieceEndgame выбирает кусок в endgame режиме:
|
||
// разрешает назначать куски в состоянии InProgress другим пирам.
|
||
func selectPieceEndgame(states []pieceState, availability []int, have []bool, hasInfo bool, endgamePeers map[int][]string, peerID string) int {
|
||
// Сначала пробуем найти pending кусок (обычный путь)
|
||
if idx := selectPendingPieceRarest(states, availability, have, hasInfo); idx >= 0 {
|
||
return idx
|
||
}
|
||
|
||
// Endgame: ищем InProgress кусок, который у нас ещё не назначен этому пиру
|
||
bestPiece := -1
|
||
bestAvailability := 0
|
||
|
||
for pieceIndex, state := range states {
|
||
if state != pieceInProgress {
|
||
continue
|
||
}
|
||
// Проверяем, не назначен ли уже этому пиру
|
||
alreadyAssigned := false
|
||
for _, pid := range endgamePeers[pieceIndex] {
|
||
if pid == peerID {
|
||
alreadyAssigned = true
|
||
break
|
||
}
|
||
}
|
||
if alreadyAssigned {
|
||
continue
|
||
}
|
||
// Проверяем наличие у пира
|
||
if hasInfo && (pieceIndex >= len(have) || !have[pieceIndex]) {
|
||
continue
|
||
}
|
||
|
||
count := availability[pieceIndex]
|
||
if bestPiece == -1 || count < bestAvailability {
|
||
bestPiece = pieceIndex
|
||
bestAvailability = count
|
||
}
|
||
}
|
||
|
||
return bestPiece
|
||
}
|
||
|
||
func updatePeerAvailability(peerAvailability map[string][]bool, availability []int, peerID string, have []bool, pieceCount int) {
|
||
normalized := make([]bool, pieceCount)
|
||
copy(normalized, have)
|
||
|
||
if prev, ok := peerAvailability[peerID]; ok {
|
||
same := true
|
||
for i := 0; i < pieceCount; i++ {
|
||
if prev[i] != normalized[i] {
|
||
same = false
|
||
break
|
||
}
|
||
}
|
||
if same {
|
||
return
|
||
}
|
||
|
||
for i := 0; i < pieceCount; i++ {
|
||
if prev[i] && availability[i] > 0 {
|
||
availability[i]--
|
||
}
|
||
if normalized[i] {
|
||
availability[i]++
|
||
}
|
||
}
|
||
peerAvailability[peerID] = normalized
|
||
return
|
||
}
|
||
|
||
for i := 0; i < pieceCount; i++ {
|
||
if normalized[i] {
|
||
availability[i]++
|
||
}
|
||
}
|
||
peerAvailability[peerID] = normalized
|
||
}
|
||
|
||
func releasePeerAvailability(peerAvailability map[string][]bool, availability []int, peerID string) {
|
||
prev, ok := peerAvailability[peerID]
|
||
if !ok {
|
||
return
|
||
}
|
||
|
||
for i := range prev {
|
||
if prev[i] && availability[i] > 0 {
|
||
availability[i]--
|
||
}
|
||
}
|
||
delete(peerAvailability, peerID)
|
||
}
|
||
|
||
func selectPendingPieceRarest(states []pieceState, availability []int, have []bool, hasInfo bool) int {
|
||
firstPending := firstPendingPiece(states)
|
||
if firstPending < 0 {
|
||
return -1
|
||
}
|
||
|
||
// If peer has not sent bitfield/have yet, optimistically probe.
|
||
if !hasInfo {
|
||
return firstPending
|
||
}
|
||
|
||
bestPiece := -1
|
||
bestAvailability := 0
|
||
|
||
for pieceIndex, state := range states {
|
||
if state != piecePending {
|
||
continue
|
||
}
|
||
if pieceIndex >= len(have) || !have[pieceIndex] {
|
||
continue
|
||
}
|
||
|
||
count := availability[pieceIndex]
|
||
if bestPiece == -1 || count < bestAvailability || (count == bestAvailability && pieceIndex < bestPiece) {
|
||
bestPiece = pieceIndex
|
||
bestAvailability = count
|
||
}
|
||
}
|
||
|
||
if bestPiece >= 0 {
|
||
return bestPiece
|
||
}
|
||
|
||
// Some peers send truncated availability info; allow fallback probing.
|
||
if len(have) == 0 || len(have) < len(states) {
|
||
return firstPending
|
||
}
|
||
|
||
return -1
|
||
}
|
||
|
||
func selectPendingPieceSequential(states []pieceState, have []bool, hasInfo bool) int {
|
||
firstPending := firstPendingPiece(states)
|
||
if firstPending < 0 {
|
||
return -1
|
||
}
|
||
|
||
// If peer has not sent bitfield/have yet, optimistically probe.
|
||
if !hasInfo {
|
||
return firstPending
|
||
}
|
||
|
||
for pieceIndex, state := range states {
|
||
if state != piecePending {
|
||
continue
|
||
}
|
||
if pieceIndex >= len(have) || !have[pieceIndex] {
|
||
continue
|
||
}
|
||
return pieceIndex // В последовательном режиме сразу берём первый доступный
|
||
}
|
||
|
||
// Some peers send truncated availability info; allow fallback probing.
|
||
if len(have) == 0 || len(have) < len(states) {
|
||
return firstPending
|
||
}
|
||
|
||
return -1
|
||
}
|
||
|
||
func firstPendingPiece(states []pieceState) int {
|
||
for pieceIndex, state := range states {
|
||
if state == piecePending {
|
||
return pieceIndex
|
||
}
|
||
}
|
||
return -1
|
||
}
|