From dfe946a41bc984a5c85dcf725c2a593a274037a2 Mon Sep 17 00:00:00 2001 From: db123 Date: Thu, 25 Jun 2026 23:23:19 +0330 Subject: [PATCH] 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 --- .env.example | 2 +- README.md | 2 +- cmd/bot/main.go | 2 +- internal/domain/constants.go | 2 +- internal/domain/types.go | 1 + internal/handler/download.go | 213 +++++++++++++++++++++++++---------- internal/handler/handler.go | 103 +++++++++++------ internal/handler/playlist.go | 19 ++-- 8 files changed, 233 insertions(+), 111 deletions(-) diff --git a/.env.example b/.env.example index 96202d1..f3a44b9 100644 --- a/.env.example +++ b/.env.example @@ -3,7 +3,7 @@ HTTP_PROXY= HTTPS_PROXY= DB_PATH=data/uptodown.db DOWNLOAD_DIR=/tmp/uptodown -MAX_FILE_SIZE_MB=2000 +MAX_FILE_SIZE_MB=50 COOKIES_FILE= TZ=UTC BOT_ALLOWED_USERS= diff --git a/README.md b/README.md index 9802520..16220da 100644 --- a/README.md +++ b/README.md @@ -41,7 +41,7 @@ Copy `.env.example` to `.env` and configure: | `HTTPS_PROXY` | empty | Outbound HTTPS proxy | | `DB_PATH` | `data/uptodown.db` | SQLite database path | | `DOWNLOAD_DIR` | `/tmp/uptodown` | Temporary download directory | -| `MAX_FILE_SIZE_MB` | `2000` | Maximum upload file size in MB | +| `MAX_FILE_SIZE_MB` | `50` | Maximum upload file size in MB | | `COOKIES_FILE` | empty | Path to cookies.txt for auth | | `TZ` | `UTC` | Timezone | | `BOT_ALLOWED_USERS` | empty | Comma-separated allowed user IDs | diff --git a/cmd/bot/main.go b/cmd/bot/main.go index d389191..af2e198 100644 --- a/cmd/bot/main.go +++ b/cmd/bot/main.go @@ -29,7 +29,7 @@ func main() { token := os.Getenv("BOT_TOKEN") dbPath := getEnvDefault("DB_PATH", "data/uptodown.db") downloadDir := getEnvDefault("DOWNLOAD_DIR", "/tmp/uptodown") - maxFileSizeMB := getEnvIntDefault("MAX_FILE_SIZE_MB", 2000) + maxFileSizeMB := getEnvIntDefault("MAX_FILE_SIZE_MB", 50) cookiesFile := os.Getenv("COOKIES_FILE") if token == "" { diff --git a/internal/domain/constants.go b/internal/domain/constants.go index 03ba98c..d2da964 100644 --- a/internal/domain/constants.go +++ b/internal/domain/constants.go @@ -6,7 +6,7 @@ const ( RateLimitInterval = 1 * time.Second ProgressThreshold = 5.0 MaxFileSizeEnvKey = "MAX_FILE_SIZE_MB" - DefaultMaxFileSize = 2000 + DefaultMaxFileSize = 50 DefaultDailyQuotaMB = 100 DefaultWeeklyQuotaMB = 500 DefaultMonthlyQuotaMB = 2000 diff --git a/internal/domain/types.go b/internal/domain/types.go index 413a71c..f42fb6c 100644 --- a/internal/domain/types.go +++ b/internal/domain/types.go @@ -129,6 +129,7 @@ type DownloadJob struct { Title string QuotaApplied bool StderrLog string + Done chan struct{} CancelCh chan struct{} Cmd *exec.Cmd } diff --git a/internal/handler/download.go b/internal/handler/download.go index 401b1c8..a9ecf24 100644 --- 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" diff --git a/internal/handler/handler.go b/internal/handler/handler.go index 69304a1..af7f3f6 100644 --- a/internal/handler/handler.go +++ b/internal/handler/handler.go @@ -36,18 +36,29 @@ type PendingInput struct { Data map[string]string } +const maxConcurrentDownloads = 3 + type Handler struct { Bot *tgbot.Bot Repo repo.Repository YtDlp *dl.Client AllowedUsers map[int64]bool + activeJobs map[int64]*domain.DownloadJob stepStates map[int64]*stepState lastMsgTime map[int64]time.Time pendingInput map[int64]*PendingInput - mu sync.Mutex - downloadDir string - maxFileSize int64 + cancelFuncs map[int64]context.CancelFunc + + activeJobsMu sync.Mutex + stepStatesMu sync.Mutex + rateLimiterMu sync.Mutex + pendingInputMu sync.Mutex + cancelFuncsMu sync.Mutex + + downloadDir string + maxFileSize int64 + downloadSem chan struct{} } func NewHandler(bot *tgbot.Bot, store repo.Repository, ytdlp *dl.Client, allowed map[int64]bool, downloadDir string, maxFileSize int64) *Handler { @@ -60,8 +71,10 @@ func NewHandler(bot *tgbot.Bot, store repo.Repository, ytdlp *dl.Client, allowed stepStates: make(map[int64]*stepState), lastMsgTime: make(map[int64]time.Time), pendingInput: make(map[int64]*PendingInput), + cancelFuncs: make(map[int64]context.CancelFunc), downloadDir: downloadDir, maxFileSize: maxFileSize, + downloadSem: make(chan struct{}, maxConcurrentDownloads), } go h.periodicCleanup() return h @@ -70,27 +83,29 @@ func NewHandler(bot *tgbot.Bot, store repo.Repository, ytdlp *dl.Client, allowed func (h *Handler) periodicCleanup() { for { time.Sleep(domain.CleanupInterval) - h.mu.Lock() - // Clean rate limit entries - cutoff := time.Now().Add(-domain.RateLimitInterval * 2) - for uid, t := range h.lastMsgTime { - if t.Before(cutoff) { - delete(h.lastMsgTime, uid) + func() { + h.rateLimiterMu.Lock() + defer h.rateLimiterMu.Unlock() + cutoff := time.Now().Add(-domain.RateLimitInterval * 2) + for uid, t := range h.lastMsgTime { + if t.Before(cutoff) { + delete(h.lastMsgTime, uid) + } } - } + }() - // Clean stale pending inputs - pendingCutoff := time.Now().Add(-5 * time.Minute) - for chatID, pi := range h.pendingInput { - if pi.CreatedAt.Before(pendingCutoff) { - delete(h.pendingInput, chatID) + func() { + h.pendingInputMu.Lock() + defer h.pendingInputMu.Unlock() + pendingCutoff := time.Now().Add(-5 * time.Minute) + for chatID, pi := range h.pendingInput { + if pi.CreatedAt.Before(pendingCutoff) { + delete(h.pendingInput, chatID) + } } - } + }() - h.mu.Unlock() - - // Clean old temp files (>1h) h.cleanOldTempFiles() } } @@ -116,8 +131,8 @@ func (h *Handler) cleanOldTempFiles() { } func (h *Handler) isRateLimited(userID int64) bool { - h.mu.Lock() - defer h.mu.Unlock() + h.rateLimiterMu.Lock() + defer h.rateLimiterMu.Unlock() now := time.Now() if last, ok := h.lastMsgTime[userID]; ok && now.Sub(last) < domain.RateLimitInterval { return true @@ -147,19 +162,19 @@ func (h *Handler) HandleMessage(ctx context.Context, msg *models.Message) { slog.Info("message", "chat_id", msg.Chat.ID, "text", msg.Text) if msg.Text[0] == '/' { - h.mu.Lock() + h.pendingInputMu.Lock() delete(h.pendingInput, msg.Chat.ID) - h.mu.Unlock() + h.pendingInputMu.Unlock() h.routeCommand(ctx, msg, userID) return } - h.mu.Lock() + h.pendingInputMu.Lock() pi, hasPI := h.pendingInput[msg.Chat.ID] if hasPI { delete(h.pendingInput, msg.Chat.ID) } - h.mu.Unlock() + h.pendingInputMu.Unlock() if hasPI { h.processPendingInput(ctx, msg, userID, pi) @@ -212,9 +227,9 @@ func (h *Handler) HandleCallback(ctx context.Context, cb *models.CallbackQuery) } slog.Info("callback", "chat_id", chatID, "data", cb.Data) - h.mu.Lock() + h.pendingInputMu.Lock() delete(h.pendingInput, chatID) - h.mu.Unlock() + h.pendingInputMu.Unlock() msgID := cb.Message.Message.ID data := cb.Data @@ -294,7 +309,7 @@ func (h *Handler) routePrefixedCallback(ctx context.Context, chatID int64, msgID } func (h *Handler) setPending(ctx context.Context, chatID int64, msgID int, userID int64, kind PendingKind, data map[string]string) { - h.mu.Lock() + h.pendingInputMu.Lock() h.pendingInput[chatID] = &PendingInput{ Kind: kind, UserID: userID, @@ -302,7 +317,7 @@ func (h *Handler) setPending(ctx context.Context, chatID int64, msgID int, userI MsgID: msgID, Data: data, } - h.mu.Unlock() + h.pendingInputMu.Unlock() } func (h *Handler) processPendingInput(ctx context.Context, msg *models.Message, userID int64, pi *PendingInput) { @@ -412,26 +427,40 @@ func (h *Handler) backToMenu(ctx context.Context, chatID int64, msgID int, cbID } func (h *Handler) clearActiveJob(userID int64) { - h.mu.Lock() - defer h.mu.Unlock() - if job, ok := h.activeJobs[userID]; ok { + h.activeJobsMu.Lock() + job, ok := h.activeJobs[userID] + if ok { dl.CancelDownload(job) delete(h.activeJobs, userID) } + h.activeJobsMu.Unlock() + + h.cancelFuncsMu.Lock() + if cancel, ok := h.cancelFuncs[userID]; ok { + cancel() + delete(h.cancelFuncs, userID) + } + h.cancelFuncsMu.Unlock() } func (h *Handler) getActiveJob(userID int64) *domain.DownloadJob { - h.mu.Lock() - defer h.mu.Unlock() + h.activeJobsMu.Lock() + defer h.activeJobsMu.Unlock() return h.activeJobs[userID] } func (h *Handler) setActiveJob(userID int64, job *domain.DownloadJob) { - h.mu.Lock() - defer h.mu.Unlock() + h.activeJobsMu.Lock() + defer h.activeJobsMu.Unlock() h.activeJobs[userID] = job } +func (h *Handler) setCancelFunc(userID int64, cancel context.CancelFunc) { + h.cancelFuncsMu.Lock() + defer h.cancelFuncsMu.Unlock() + h.cancelFuncs[userID] = cancel +} + func formatBytes(bytes int64) string { const unit = 1024 if bytes < unit { @@ -455,7 +484,7 @@ func formatBytes(bytes int64) string { } func (h *Handler) generateJobID() string { - return uuid.NewString()[:8] + return uuid.NewString() } // handleStart shows the main menu and cleans up any stale state. diff --git a/internal/handler/playlist.go b/internal/handler/playlist.go index b934ad0..6b9582f 100644 --- a/internal/handler/playlist.go +++ b/internal/handler/playlist.go @@ -4,7 +4,6 @@ import ( "context" "fmt" "strconv" - "time" "uptodownBot/internal/domain" ) @@ -164,6 +163,8 @@ func (h *Handler) handlePlaylistConfirm(ctx context.Context, chatID int64, msgID } } + dlCtx, dlCancel := context.WithCancel(context.Background()) + job := &domain.DownloadJob{ ID: h.generateJobID(), UserID: userID, @@ -176,19 +177,17 @@ func (h *Handler) handlePlaylistConfirm(ctx context.Context, chatID int64, msgID Status: domain.StatusDownloading, Title: entry.Title, FileSize: format.Filesize, + Done: make(chan struct{}), } h.setActiveJob(userID, job) - h.runDownload(context.Background(), job) + h.setCancelFunc(userID, dlCancel) - // Wait for active job to complete - for { - j := h.getActiveJob(userID) - if j == nil || j.Status == domain.StatusCompleted || j.Status == domain.StatusFailed || j.Status == domain.StatusCancelled { - break - } - time.Sleep(500 * time.Millisecond) - } + h.runDownload(dlCtx, job) + + // runDownload is synchronous — it blocks until the download completes. + // The Done channel is closed inside runDownload when it exits. + // No busy-polling needed. } h.sendText(ctx, chatID, "All playlist downloads completed.")