From 20d05be69f6be6277bfb40d52bc0e6c6b7243f79 Mon Sep 17 00:00:00 2001 From: itexpert228 <67105314+fdaser1337@users.noreply.github.com> Date: Fri, 6 Mar 2026 14:09:24 +0300 Subject: [PATCH] =?UTF-8?q?=D1=8F=20=D0=B7=D0=B0=D0=B5=D0=B1=D0=B0=D0=BB?= =?UTF-8?q?=D1=81=D1=8F=20=D1=81=D1=83=D0=BA=D0=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/torrent/engine.go | 283 ++++++++++++++++++++--------- internal/torrent/peerwire.go | 81 +-------- internal/torrent/piecescheduler.go | 210 +++++++++++++++++++++ 3 files changed, 413 insertions(+), 161 deletions(-) create mode 100644 internal/torrent/piecescheduler.go diff --git a/internal/torrent/engine.go b/internal/torrent/engine.go index 200ee14..2407453 100644 --- a/internal/torrent/engine.go +++ b/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 { diff --git a/internal/torrent/peerwire.go b/internal/torrent/peerwire.go index ae2ab7b..8fb674f 100644 --- a/internal/torrent/peerwire.go +++ b/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{}) diff --git a/internal/torrent/piecescheduler.go b/internal/torrent/piecescheduler.go new file mode 100644 index 0000000..4808468 --- /dev/null +++ b/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 +}