refactor: split mutex, add download queue, wire context cancellation
All checks were successful
CI / build (push) Successful in 53s

- 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
This commit is contained in:
2026-06-25 23:23:19 +03:30
parent 3290bb55c0
commit dfe946a41b
8 changed files with 233 additions and 111 deletions

View File

@@ -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=

View File

@@ -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 |

View File

@@ -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 == "" {

View File

@@ -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

View File

@@ -129,6 +129,7 @@ type DownloadJob struct {
Title string
QuotaApplied bool
StderrLog string
Done chan struct{}
CancelCh chan struct{}
Cmd *exec.Cmd
}

View File

@@ -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"

View File

@@ -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.

View File

@@ -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.")