Ztorrent/internal/torrent/piecescheduler.go

package torrent

import (
	"context"
	"sync"
)

type pieceTask struct {
	Index int
}

type assignPieceRequest struct {
	have       []bool
	hasInfo    bool
	responseCh chan assignPieceResponse
}

type assignPieceResponse struct {
	task pieceTask
	ok   bool
}

type reportPieceRequest struct {
	pieceIndex int
	success    bool
	responseCh chan bool
}

type pieceScheduler struct {
	assignCh   chan assignPieceRequest
	reportCh   chan reportPieceRequest
	progressCh chan int
	doneCh     chan struct{}
	stopCh     chan struct{}
	stopOnce   sync.Once
}

type pieceState uint8

const (
	piecePending pieceState = iota
	pieceInProgress
	pieceDone
)

func newPieceScheduler(pieceCount int) *pieceScheduler {
	ps := &pieceScheduler{
		assignCh:   make(chan assignPieceRequest, 128),
		reportCh:   make(chan reportPieceRequest, 128),
		progressCh: make(chan int, 128),
		doneCh:     make(chan struct{}),
		stopCh:     make(chan struct{}),
	}

	go ps.run(pieceCount)
	return ps
}

func (ps *pieceScheduler) Acquire(ctx context.Context, have []bool, hasInfo bool) (pieceTask, bool, error) {
	responseCh := make(chan assignPieceResponse, 1)
	req := assignPieceRequest{
		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) 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) run(pieceCount int) {
	states := make([]pieceState, pieceCount)
	completed := 0
	finish := func() {
		close(ps.doneCh)
		close(ps.progressCh)
	}

	if pieceCount == 0 {
		finish()
		return
	}

	for {
		if completed >= pieceCount {
			finish()
			return
		}

		select {
		case <-ps.stopCh:
			finish()
			return
		case req := <-ps.assignCh:
			pieceIndex := selectPendingPiece(states, req.have, req.hasInfo)
			if pieceIndex >= 0 {
				states[pieceIndex] = pieceInProgress
				debugf("scheduler assigned piece %d", pieceIndex)
				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 >= len(states) {
				req.responseCh <- false
				continue
			}

			if req.success {
				if states[req.pieceIndex] == pieceDone {
					debugf("scheduler report piece %d ignored (already done)", req.pieceIndex)
					req.responseCh <- false
					continue
				}
				states[req.pieceIndex] = pieceDone
				completed++
				debugf("scheduler report piece %d success (%d/%d)", req.pieceIndex, completed, pieceCount)
				select {
				case ps.progressCh <- req.pieceIndex:
				default:
				}
				req.responseCh <- true
				continue
			}

			if states[req.pieceIndex] == pieceInProgress {
				states[req.pieceIndex] = piecePending
				debugf("scheduler report piece %d failed, re-queued", req.pieceIndex)
			}
			req.responseCh <- false
		}
	}
}

func selectPendingPiece(states []pieceState, have []bool, hasInfo bool) int {
	firstPending := firstPendingPiece(states)
	if firstPending < 0 {
		return -1
	}

	// If peer has not sent bitfield/have yet, treat it as potentially having any piece.
	if !hasInfo {
		return firstPending
	}

	// Prefer pieces explicitly advertised by peer.
	for pieceIndex, state := range states {
		if state != piecePending {
			continue
		}
		if pieceIndex < len(have) && have[pieceIndex] {
			return pieceIndex
		}
	}

	// Some peers send missing/truncated availability info; allow optimistic 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
}