refactor: split mutex, add download queue, wire context cancellation

- Split single sync.Mutex into per-map mutexes (activeJobs, stepStates, rateLimiter, pendingInput, cancelFuncs) to reduce contention - Add download semaphore (channel-buffered, 3 concurrent) as worker pool - Add cancelFuncs map with context cancellation wired into download goroutines for clean shutdown - Replace playlist busy-poll (500ms sleep loop) with synchronous runDownload + Done channel for completion tracking - Use full UUID for job IDs instead of 8-char prefix - Download thumbnail URL to temp file before sending (InputFileUpload) - Change MAX_FILE_SIZE_MB default from 2000 to 50 - Clean up cancel funcs after normal download completion
2026-06-25 23:23:19 +03:30
parent 3290bb55c0
commit dfe946a41b
8 changed files with 233 additions and 111 deletions
--- a/internal/handler/download.go
+++ b/internal/handler/download.go
@@ -3,8 +3,10 @@ package handler
 import (
 	"context"
 	"fmt"
+	"io"
 	"log/slog"
 	"math"
+	"net/http"
 	"os"
 	"path/filepath"
 	"strings"
@@ -33,20 +35,20 @@ type stepState struct {
 }

 func (h *Handler) getUserStepState(chatID int64) *stepState {
-	h.mu.Lock()
-	defer h.mu.Unlock()
+	h.stepStatesMu.Lock()
+	defer h.stepStatesMu.Unlock()
 	return h.stepStates[chatID]
 }

 func (h *Handler) setUserStepState(chatID int64, ss *stepState) {
-	h.mu.Lock()
-	defer h.mu.Unlock()
+	h.stepStatesMu.Lock()
+	defer h.stepStatesMu.Unlock()
 	h.stepStates[chatID] = ss
 }

 func (h *Handler) clearStepState(chatID int64) {
-	h.mu.Lock()
-	defer h.mu.Unlock()
+	h.stepStatesMu.Lock()
+	defer h.stepStatesMu.Unlock()
 	delete(h.stepStates, chatID)
 }

@@ -355,6 +357,8 @@ func (h *Handler) handleContainerSelection(ctx context.Context, chatID int64, ms
 		return
 	}

+	dlCtx, dlCancel := context.WithCancel(context.Background())
+
 	job := &domain.DownloadJob{
 		ID:             h.generateJobID(),
 		UserID:         userID,
@@ -369,17 +373,33 @@ func (h *Handler) handleContainerSelection(ctx context.Context, chatID int64, ms
 		Status:         domain.StatusDownloading,
 		Title:          ss.MediaInfo.Title,
 		FileSize:       fileSize,
+		Done:           make(chan struct{}),
 	}

 	h.editText(ctx, chatID, msgID, "Processing...", nil)
 	h.setActiveJob(userID, job)
+	h.setCancelFunc(userID, dlCancel)
 	h.clearStepState(chatID)

-	go h.runDownload(context.Background(), job)
+	go h.runDownload(dlCtx, job)
 }

 // runDownload manages the download lifecycle, progress, and completion.
 func (h *Handler) runDownload(ctx context.Context, job *domain.DownloadJob) {
+	defer func() {
+		if job.Done != nil {
+			close(job.Done)
+		}
+	}()
+
+	// Acquire semaphore slot or wait for cancellation
+	select {
+	case h.downloadSem <- struct{}{}:
+	case <-ctx.Done():
+		return
+	}
+	defer func() { <-h.downloadSem }()
+
 	kb := progressKeyboard(0, "", "")
 	msg, err := h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{
 		ChatID:      job.ChatID,
@@ -396,61 +416,77 @@ func (h *Handler) runDownload(ctx context.Context, job *domain.DownloadJob) {
 	if err != nil {
 		h.editText(ctx, job.ChatID, job.MessageID, fmt.Sprintf("Download failed: %s", err.Error()), nil)
 		job.Status = domain.StatusFailed
-		h.mu.Lock()
+		h.activeJobsMu.Lock()
 		delete(h.activeJobs, job.UserID)
-		h.mu.Unlock()
+		h.activeJobsMu.Unlock()
 		return
 	}

 	lastProgress := 0.0
-	for upd := range updates {
-		slog.Info("progress update", "pct", upd.Progress, "speed", upd.Speed, "status", upd.Status)
+loop:
+	for {
+		select {
+		case upd, ok := <-updates:
+			if !ok {
+				break loop
+			}
+			slog.Info("progress update", "pct", upd.Progress, "speed", upd.Speed, "status", upd.Status)

-		h.mu.Lock()
-		currentJob := h.activeJobs[job.UserID]
-		h.mu.Unlock()
-		if currentJob == nil {
-			return
+			h.activeJobsMu.Lock()
+			currentJob := h.activeJobs[job.UserID]
+			h.activeJobsMu.Unlock()
+			if currentJob == nil {
+				return
+			}
+
+			if upd.Status == domain.StatusCancelled {
+				h.editText(ctx, job.ChatID, job.MessageID, "Download cancelled.", nil)
+				h.activeJobsMu.Lock()
+				delete(h.activeJobs, job.UserID)
+				h.activeJobsMu.Unlock()
+				return
+			}
+
+			if upd.Status == domain.StatusFailed {
+				h.editText(ctx, job.ChatID, job.MessageID, "Download failed.", nil)
+				h.activeJobsMu.Lock()
+				delete(h.activeJobs, job.UserID)
+				h.activeJobsMu.Unlock()
+				return
+			}
+
+			if upd.Progress > 0 && lastProgress == 0 {
+				lastProgress = upd.Progress
+			} else if upd.Progress-lastProgress < domain.ProgressThreshold && upd.Status != domain.StatusCompleted {
+				continue
+			} else {
+				lastProgress = upd.Progress
+			}
+
+			if upd.Progress > domain.QuotaProgressThreshold && !job.QuotaApplied && job.FileSize > 0 {
+				job.QuotaApplied = true
+				h.applyQuota(job.UserID, job.FileSize)
+			}
+
+			kb := progressKeyboard(upd.Progress, upd.Speed, upd.ETA)
+			text := fmt.Sprintf("Downloading: %s", job.Title)
+			h.editText(ctx, job.ChatID, job.MessageID, text, &kb)
+
+		case <-ctx.Done():
+			dl.CancelDownload(job)
+			for range updates {
+			}
+			break loop
 		}
-
-		if upd.Status == domain.StatusCancelled {
-			h.editText(ctx, job.ChatID, job.MessageID, "Download cancelled.", nil)
-			h.mu.Lock()
-			delete(h.activeJobs, job.UserID)
-			h.mu.Unlock()
-			return
-		}
-
-		if upd.Status == domain.StatusFailed {
-			h.editText(ctx, job.ChatID, job.MessageID, "Download failed.", nil)
-			h.mu.Lock()
-			delete(h.activeJobs, job.UserID)
-			h.mu.Unlock()
-			return
-		}
-
-		// Always show the first real progress update, then throttle at 5% intervals
-		if upd.Progress > 0 && lastProgress == 0 {
-			lastProgress = upd.Progress
-		} else if upd.Progress-lastProgress < domain.ProgressThreshold && upd.Status != domain.StatusCompleted {
-			continue
-		} else {
-			lastProgress = upd.Progress
-		}
-
-		if upd.Progress > domain.QuotaProgressThreshold && !job.QuotaApplied && job.FileSize > 0 {
-			job.QuotaApplied = true
-			h.applyQuota(job.UserID, job.FileSize)
-		}
-
-		kb := progressKeyboard(upd.Progress, upd.Speed, upd.ETA)
-		text := fmt.Sprintf("Downloading: %s", job.Title)
-		h.editText(ctx, job.ChatID, job.MessageID, text, &kb)
 	}

-	h.mu.Lock()
+	h.activeJobsMu.Lock()
 	delete(h.activeJobs, job.UserID)
-	h.mu.Unlock()
+	h.activeJobsMu.Unlock()
+
+	h.cancelFuncsMu.Lock()
+	delete(h.cancelFuncs, job.UserID)
+	h.cancelFuncsMu.Unlock()

 	if job.Status == domain.StatusCompleted {
 		pattern := filepath.Join(h.downloadDir, job.ID+"_*")
@@ -518,12 +554,19 @@ func (h *Handler) sendDocument(ctx context.Context, chatID int64, filePath, titl
 func (h *Handler) handleCancelDownload(ctx context.Context, chatID int64, msgID int, cbID string, userID int64) {
 	defer h.answerCb(ctx, cbID)

+	h.cancelFuncsMu.Lock()
+	if cancel, ok := h.cancelFuncs[userID]; ok {
+		cancel()
+		delete(h.cancelFuncs, userID)
+	}
+	h.cancelFuncsMu.Unlock()
+
 	job := h.getActiveJob(userID)
 	if job != nil {
 		dl.CancelDownload(job)
-		h.mu.Lock()
+		h.activeJobsMu.Lock()
 		delete(h.activeJobs, userID)
-		h.mu.Unlock()
+		h.activeJobsMu.Unlock()
 		h.editText(ctx, chatID, msgID, "Download cancelled.", nil)
 		return
 	}
@@ -647,21 +690,71 @@ func (h *Handler) sendMediaPreview(ctx context.Context, chatID int64, info *doma
 	}

 	if info.Thumbnail != "" {
-		photo := &models.InputFileString{Data: info.Thumbnail}
-		_, err := h.Bot.SendPhoto(ctx, &tgbot.SendPhotoParams{
-			ChatID:  chatID,
-			Photo:   photo,
-			Caption: caption,
+		thumbPath, err := downloadFile(info.Thumbnail, h.downloadDir)
+		if err != nil {
+			slog.Warn("download thumbnail failed, falling back to text", "url", info.Thumbnail, "error", err)
+			h.sendText(ctx, chatID, caption)
+			return
+		}
+		defer os.Remove(thumbPath)
+
+		f, err := os.Open(thumbPath)
+		if err != nil {
+			slog.Warn("open thumbnail failed, falling back to text", "error", err)
+			h.sendText(ctx, chatID, caption)
+			return
+		}
+		defer f.Close()
+
+		photo := &models.InputFileUpload{Filename: "thumb.jpg", Data: f}
+		_, err = h.Bot.SendPhoto(ctx, &tgbot.SendPhotoParams{
+			ChatID: chatID, Photo: photo, Caption: caption,
 		})
 		if err == nil {
 			return
 		}
-		slog.Warn("send thumbnail failed, falling back to text", "error", err)
+		slog.Warn("send thumbnail photo failed, falling back to text", "error", err)
 	}

 	h.sendText(ctx, chatID, caption)
 }

+// downloadFile downloads a URL to a temp file in the given directory and returns the path.
+func downloadFile(url, dir string) (string, error) {
+	resp, err := http.Get(url)
+	if err != nil {
+		return "", fmt.Errorf("http get: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return "", fmt.Errorf("unexpected status: %s", resp.Status)
+	}
+
+	if err := os.MkdirAll(dir, 0755); err != nil {
+		return "", fmt.Errorf("create download dir: %w", err)
+	}
+
+	f, err := os.CreateTemp(dir, "thumb_*")
+	if err != nil {
+		return "", fmt.Errorf("create temp file: %w", err)
+	}
+
+	_, err = io.Copy(f, resp.Body)
+	if err != nil {
+		f.Close()
+		os.Remove(f.Name())
+		return "", fmt.Errorf("write temp file: %w", err)
+	}
+
+	if err := f.Close(); err != nil {
+		os.Remove(f.Name())
+		return "", fmt.Errorf("close temp file: %w", err)
+	}
+
+	return f.Name(), nil
+}
+
 func formatDuration(seconds float64) string {
 	if seconds <= 0 {
 		return "0:00"