я заебался сука

internal/torrent/engine.go

@@ -3,6 +3,7 @@ package torrent
 import (
 	"context"
 	"crypto/rand"
+	"crypto/sha1"
 	"encoding/hex"
 	"errors"
 	"fmt"
@@ -257,84 +258,85 @@ func (e *Engine) runDownload(ctx context.Context, tf *torrentfile.TorrentFile) {
 	}
 	defer partFile.Close()
 
-	pieceDone := make([]bool, len(tf.PieceHashes))
+	scheduler := newPieceScheduler(len(tf.PieceHashes))
+	workerCtx, workerCancel := context.WithCancel(ctx)
+	var workersWG sync.WaitGroup
+
+	startedWorkers := make(map[string]struct{})
+	workerDoneCh := make(chan string, 512)
+	startWorkers := func(peers []PeerStatus) {
+		for _, peer := range peers {
+			key := peerStatusKey(peer.Address, peer.Port)
+			if _, exists := startedWorkers[key]; exists {
+				continue
+			}
+			startedWorkers[key] = struct{}{}
+			workersWG.Add(1)
+			go func(workerKey string, p PeerStatus) {
+				defer workersWG.Done()
+				e.runPeerWorker(workerCtx, tf, p, partFile, scheduler)
+				select {
+				case workerDoneCh <- workerKey:
+				default:
+				}
+			}(key, peer)
+		}
+	}
+
+	var cleanupOnce sync.Once
+	cleanupWorkers := func() {
+		cleanupOnce.Do(func() {
+			workerCancel()
+			scheduler.Stop()
+			workersWG.Wait()
+		})
+	}
+	defer cleanupWorkers()
+
 	e.setPhase("downloading")
 	lastProgressAt := time.Now()
 	lastReannounceAt := time.Now()
+	startWorkers(e.snapshotPeers())
+	heartbeatTicker := time.NewTicker(idleRetryDelay)
+	defer heartbeatTicker.Stop()
+	downloadCompleted := false
 
-	for !allPiecesDone(pieceDone) {
-		if err := ctx.Err(); err != nil {
+	for !downloadCompleted {
+		select {
+		case <-ctx.Done():
 			e.setPhase("stopped")
 			e.setError(nil)
 			return
-		}
-
-		peers := e.snapshotPeers()
-		if len(peers) == 0 {
-			e.setPhase("querying_trackers")
-			refreshCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
-			_ = e.refreshPeersFromTrackers(refreshCtx, tf)
-			cancel()
-			lastReannounceAt = time.Now()
-			e.setPhase("downloading")
-			peers = e.snapshotPeers()
-		}
-
-		madeProgress := false
-		for _, peer := range peers {
-			if allPiecesDone(pieceDone) {
-				break
-			}
-			if err := ctx.Err(); err != nil {
-				e.setPhase("stopped")
-				e.setError(nil)
-				return
-			}
-
-			e.setPeerState(peer.Address, peer.Port, "connecting", "")
-			pieces, _, err := downloadFromPeer(ctx, tf, peer, e.peerID, partFile, pieceDone, func(pieceIndex int, pieceSize int) {
-				e.recordPieceComplete(peer.Address, peer.Port, pieceSize)
-			})
-
-			if err != nil {
-				e.setPeerState(peer.Address, peer.Port, "error", err.Error())
+		case _, ok := <-scheduler.Progress():
+			if !ok {
+				downloadCompleted = true
 				continue
 			}
+			lastProgressAt = time.Now()
+		case <-scheduler.Done():
+			downloadCompleted = true
+		case workerKey := <-workerDoneCh:
+			delete(startedWorkers, workerKey)
+		case <-heartbeatTicker.C:
+			if time.Since(lastReannounceAt) >= reannounceInterval {
+				e.setPhase("querying_trackers")
+				refreshCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
+				_ = e.refreshPeersFromTrackers(refreshCtx, tf)
+				cancel()
+				lastReannounceAt = time.Now()
+				e.setPhase("downloading")
+			}
+			startWorkers(e.snapshotPeers())
 
-			if pieces > 0 {
-				madeProgress = true
-				e.setPeerState(peer.Address, peer.Port, "active", "")
-			} else {
-				e.setPeerState(peer.Address, peer.Port, "idle", "")
+			if time.Since(lastProgressAt) >= downloadStallTimeout {
+				e.setTerminalError("stalled", fmt.Errorf("download stalled: no piece progress for %s", downloadStallTimeout.Round(time.Second)))
+				return
 			}
 		}
-
-		if madeProgress {
-			lastProgressAt = time.Now()
-			continue
-		}
-
-		if time.Since(lastReannounceAt) >= reannounceInterval {
-			e.setPhase("querying_trackers")
-			refreshCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
-			_ = e.refreshPeersFromTrackers(refreshCtx, tf)
-			cancel()
-			lastReannounceAt = time.Now()
-			e.setPhase("downloading")
-		}
-
-		if time.Since(lastProgressAt) >= downloadStallTimeout {
-			e.setTerminalError("stalled", fmt.Errorf("download stalled: no piece progress for %s", downloadStallTimeout.Round(time.Second)))
-			return
-		}
-
-		if err := sleepWithContext(ctx, idleRetryDelay); err != nil {
-			e.setPhase("stopped")
-			e.setError(nil)
-			return
-		}
 	}
 
+	cleanupWorkers()
+
 	e.setPhase("writing_files")
 	if err := materializeDownloadedFiles(tf, partPath, "downloads"); err != nil {
 		e.setTerminalError("failed", fmt.Errorf("write output files: %w", err))
@@ -351,6 +353,138 @@ func (e *Engine) runDownload(ctx context.Context, tf *torrentfile.TorrentFile) {
 	e.mu.Unlock()
 }
 
+func (e *Engine) runPeerWorker(
+	ctx context.Context,
+	tf *torrentfile.TorrentFile,
+	peer PeerStatus,
+	partFile *os.File,
+	scheduler *pieceScheduler,
+) {
+	e.setPeerState(peer.Address, peer.Port, "connecting", "")
+
+	peerAddr := net.JoinHostPort(peer.Address, strconv.Itoa(int(peer.Port)))
+	client, err := newPeerClient(ctx, peerAddr, tf.InfoHash, e.peerID, len(tf.PieceHashes))
+	if err != nil {
+		if ctx.Err() != nil {
+			e.setPeerState(peer.Address, peer.Port, "stopped", "")
+		} else {
+			e.setPeerState(peer.Address, peer.Port, "error", err.Error())
+		}
+		return
+	}
+	defer client.Close()
+
+	e.setPeerState(peer.Address, peer.Port, "ready", "")
+	consecutiveFailures := 0
+
+	for {
+		if ctx.Err() != nil {
+			e.setPeerState(peer.Address, peer.Port, "stopped", "")
+			return
+		}
+
+		have, hasInfo := client.PieceAvailability()
+		task, ok, err := scheduler.Acquire(ctx, have, hasInfo)
+		if err != nil {
+			if ctx.Err() != nil {
+				e.setPeerState(peer.Address, peer.Port, "stopped", "")
+			} else {
+				e.setPeerState(peer.Address, peer.Port, "error", err.Error())
+			}
+			return
+		}
+
+		if !ok {
+			select {
+			case <-ctx.Done():
+				e.setPeerState(peer.Address, peer.Port, "stopped", "")
+				return
+			case <-scheduler.Done():
+				e.setPeerState(peer.Address, peer.Port, "done", "")
+				return
+			case <-time.After(500 * time.Millisecond):
+			}
+			continue
+		}
+
+		e.setPeerState(peer.Address, peer.Port, "requesting", "")
+		pieceSize := pieceSizeForIndex(tf, task.Index)
+		pieceData, err := client.DownloadPiece(ctx, task.Index, pieceSize)
+		if err != nil {
+			_, _ = scheduler.Report(ctx, task.Index, false)
+			if ctx.Err() != nil {
+				e.setPeerState(peer.Address, peer.Port, "stopped", "")
+				return
+			}
+
+			consecutiveFailures++
+			e.setPeerState(peer.Address, peer.Port, "error", err.Error())
+			if shouldDropPeer(err) || consecutiveFailures >= 3 {
+				return
+			}
+			continue
+		}
+
+		hash := sha1.Sum(pieceData)
+		if hash != tf.PieceHashes[task.Index] {
+			_, _ = scheduler.Report(ctx, task.Index, false)
+			consecutiveFailures++
+			e.setPeerState(peer.Address, peer.Port, "error", fmt.Sprintf("piece %d hash mismatch", task.Index))
+			if consecutiveFailures >= 3 {
+				return
+			}
+			continue
+		}
+
+		offset := int64(task.Index * tf.PieceLength)
+		if _, err := partFile.WriteAt(pieceData, offset); err != nil {
+			_, _ = scheduler.Report(ctx, task.Index, false)
+			e.setPeerState(peer.Address, peer.Port, "error", fmt.Sprintf("write piece %d: %v", task.Index, err))
+			return
+		}
+
+		accepted, reportErr := scheduler.Report(ctx, task.Index, true)
+		if reportErr != nil {
+			if ctx.Err() != nil {
+				e.setPeerState(peer.Address, peer.Port, "stopped", "")
+			} else {
+				e.setPeerState(peer.Address, peer.Port, "error", reportErr.Error())
+			}
+			return
+		}
+		if accepted {
+			e.recordPieceComplete(peer.Address, peer.Port, len(pieceData))
+		}
+
+		consecutiveFailures = 0
+		e.setPeerState(peer.Address, peer.Port, "active", "")
+	}
+}
+
+func shouldDropPeer(err error) bool {
+	if err == nil {
+		return false
+	}
+	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+		return true
+	}
+	if isTimeout(err) {
+		return false
+	}
+
+	msg := strings.ToLower(err.Error())
+	switch {
+	case strings.Contains(msg, "peer choked"):
+		return false
+	case strings.Contains(msg, "peer did not unchoke"):
+		return false
+	case strings.Contains(msg, "hash mismatch"):
+		return false
+	default:
+		return true
+	}
+}
+
 func (e *Engine) snapshotPeers() []PeerStatus {
 	e.mu.RLock()
 	defer e.mu.RUnlock()
@@ -621,27 +755,6 @@ func firstTrackerError(statuses []TrackerStatus) error {
 	return nil
 }
 
-func allPiecesDone(pieceDone []bool) bool {
-	for _, done := range pieceDone {
-		if !done {
-			return false
-		}
-	}
-	return true
-}
-
-func sleepWithContext(ctx context.Context, d time.Duration) error {
-	timer := time.NewTimer(d)
-	defer timer.Stop()
-
-	select {
-	case <-ctx.Done():
-		return ctx.Err()
-	case <-timer.C:
-		return nil
-	}
-}
-
 func createPartFile(tf *torrentfile.TorrentFile) (string, *os.File, error) {
 	partsDir := filepath.Join("downloads", ".parts")
 	if err := os.MkdirAll(partsDir, 0o755); err != nil {

internal/torrent/peerwire.go

@@ -3,14 +3,11 @@ package torrent
 import (
 	"bytes"
 	"context"
-	"crypto/sha1"
 	"encoding/binary"
 	"errors"
 	"fmt"
 	"io"
 	"net"
-	"os"
-	"strconv"
 	"time"
 
 	"github.com/veggiedefender/torrent-client/internal/torrentfile"
@@ -49,70 +46,6 @@ type peerClient struct {
 	peerIsChoked bool
 }
 
-func downloadFromPeer(
-	ctx context.Context,
-	tf *torrentfile.TorrentFile,
-	peer PeerStatus,
-	peerID [20]byte,
-	partFile *os.File,
-	pieceDone []bool,
-	onPiece func(pieceIndex int, pieceSize int),
-) (int, int, error) {
-	peerAddr := net.JoinHostPort(peer.Address, strconv.Itoa(int(peer.Port)))
-	client, err := newPeerClient(ctx, peerAddr, tf.InfoHash, peerID, len(tf.PieceHashes))
-	if err != nil {
-		return 0, 0, err
-	}
-	defer client.Close()
-
-	piecesDownloaded := 0
-	bytesDownloaded := 0
-
-	for pieceIndex := 0; pieceIndex < len(tf.PieceHashes); pieceIndex++ {
-		if ctx.Err() != nil {
-			return piecesDownloaded, bytesDownloaded, ctx.Err()
-		}
-
-		if pieceDone[pieceIndex] {
-			continue
-		}
-		if !client.HasPiece(pieceIndex) {
-			continue
-		}
-
-		size := pieceSizeForIndex(tf, pieceIndex)
-		pieceData, err := client.DownloadPiece(ctx, pieceIndex, size)
-		if err != nil {
-			if piecesDownloaded > 0 {
-				return piecesDownloaded, bytesDownloaded, nil
-			}
-			return 0, 0, err
-		}
-
-		sum := sha1.Sum(pieceData)
-		if sum != tf.PieceHashes[pieceIndex] {
-			if piecesDownloaded > 0 {
-				return piecesDownloaded, bytesDownloaded, nil
-			}
-			return 0, 0, fmt.Errorf("piece %d hash mismatch", pieceIndex)
-		}
-
-		offset := int64(pieceIndex * tf.PieceLength)
-		if _, err := partFile.WriteAt(pieceData, offset); err != nil {
-			return piecesDownloaded, bytesDownloaded, err
-		}
-
-		pieceDone[pieceIndex] = true
-		piecesDownloaded++
-		bytesDownloaded += len(pieceData)
-		if onPiece != nil {
-			onPiece(pieceIndex, len(pieceData))
-		}
-	}
-
-	return piecesDownloaded, bytesDownloaded, nil
-}
-
 func newPeerClient(ctx context.Context, addr string, infoHash [20]byte, peerID [20]byte, pieceCount int) (*peerClient, error) {
 	dialer := net.Dialer{Timeout: peerConnectTimeout}
 	conn, err := dialer.DialContext(ctx, "tcp", addr)
@@ -148,14 +81,10 @@ func (pc *peerClient) Close() error {
 	return pc.conn.Close()
 }
 
-func (pc *peerClient) HasPiece(index int) bool {
-	if index < 0 || index >= len(pc.have) {
-		return false
-	}
-	if !pc.hasPieceInfo {
-		return true
-	}
-	return pc.have[index]
+func (pc *peerClient) PieceAvailability() ([]bool, bool) {
+	have := make([]bool, len(pc.have))
+	copy(have, pc.have)
+	return have, pc.hasPieceInfo
 }
 
 func (pc *peerClient) DownloadPiece(ctx context.Context, pieceIndex int, pieceLength int) ([]byte, error) {
@@ -190,7 +119,7 @@ func (pc *peerClient) DownloadPiece(ctx context.Context, pieceIndex int, pieceLe
 }
 
 func (pc *peerClient) readInitialMessages(ctx context.Context) error {
-	if err := pc.setReadDeadlineFromContext(ctx, 1200*time.Millisecond); err != nil {
+	if err := pc.setReadDeadlineFromContext(ctx, 2*time.Second); err != nil {
 		return err
 	}
 	defer pc.conn.SetReadDeadline(time.Time{})

internal/torrent/piecescheduler.go

@@ -0,0 +1,210 @@
+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
+				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 {
+					req.responseCh <- false
+					continue
+				}
+				states[req.pieceIndex] = pieceDone
+				completed++
+				select {
+				case ps.progressCh <- req.pieceIndex:
+				default:
+				}
+				req.responseCh <- true
+				continue
+			}
+
+			if states[req.pieceIndex] == pieceInProgress {
+				states[req.pieceIndex] = piecePending
+			}
+			req.responseCh <- false
+		}
+	}
+}
+
+func selectPendingPiece(states []pieceState, have []bool, hasInfo bool) int {
+	if hasInfo {
+		for pieceIndex, state := range states {
+			if state != piecePending {
+				continue
+			}
+			if pieceIndex >= len(have) {
+				continue
+			}
+			if have[pieceIndex] {
+				return pieceIndex
+			}
+		}
+		return -1
+	}
+
+	for pieceIndex, state := range states {
+		if state == piecePending {
+			return pieceIndex
+		}
+	}
+	return -1
+}