Compare commits

...

2 Commits

Author SHA1 Message Date
dfe946a41b 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
2026-06-25 23:23:19 +03:30
3290bb55c0 feat: add in-memory media info cache and YouTube search feature
- Add in-memory LRU cache (200 entries, 30 min TTL) to yt-dlp Client for
  GetMediaInfo and GetPlaylistInfo results, avoiding redundant yt-dlp calls
  for the same URL within a session
- Add YouTube (/search) and YouTube Music (/music) search with paginated
  results (5 per page), user selects a result to start the download flow
- Add SearchState type, new search.go handler, search results keyboard with
  < > navigation and New Search button
- Add Search button to main menu keyboard
- Update help text with new commands
2026-06-25 22:58:20 +03:30
11 changed files with 479 additions and 114 deletions

View File

@@ -3,7 +3,7 @@ HTTP_PROXY=
HTTPS_PROXY= HTTPS_PROXY=
DB_PATH=data/uptodown.db DB_PATH=data/uptodown.db
DOWNLOAD_DIR=/tmp/uptodown DOWNLOAD_DIR=/tmp/uptodown
MAX_FILE_SIZE_MB=2000 MAX_FILE_SIZE_MB=50
COOKIES_FILE= COOKIES_FILE=
TZ=UTC TZ=UTC
BOT_ALLOWED_USERS= BOT_ALLOWED_USERS=

View File

@@ -41,7 +41,7 @@ Copy `.env.example` to `.env` and configure:
| `HTTPS_PROXY` | empty | Outbound HTTPS proxy | | `HTTPS_PROXY` | empty | Outbound HTTPS proxy |
| `DB_PATH` | `data/uptodown.db` | SQLite database path | | `DB_PATH` | `data/uptodown.db` | SQLite database path |
| `DOWNLOAD_DIR` | `/tmp/uptodown` | Temporary download directory | | `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 | | `COOKIES_FILE` | empty | Path to cookies.txt for auth |
| `TZ` | `UTC` | Timezone | | `TZ` | `UTC` | Timezone |
| `BOT_ALLOWED_USERS` | empty | Comma-separated allowed user IDs | | `BOT_ALLOWED_USERS` | empty | Comma-separated allowed user IDs |

View File

@@ -29,7 +29,7 @@ func main() {
token := os.Getenv("BOT_TOKEN") token := os.Getenv("BOT_TOKEN")
dbPath := getEnvDefault("DB_PATH", "data/uptodown.db") dbPath := getEnvDefault("DB_PATH", "data/uptodown.db")
downloadDir := getEnvDefault("DOWNLOAD_DIR", "/tmp/uptodown") 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") cookiesFile := os.Getenv("COOKIES_FILE")
if token == "" { if token == "" {

View File

@@ -10,6 +10,7 @@ import (
"os/exec" "os/exec"
"path/filepath" "path/filepath"
"strings" "strings"
"sync"
"syscall" "syscall"
"time" "time"
@@ -20,6 +21,36 @@ import (
type Client struct { type Client struct {
binaryPath string binaryPath string
cookiesFile string cookiesFile string
cache map[string]*cachedMediaInfo
cacheMu sync.Mutex
}
type cachedMediaInfo struct {
info *domain.MediaInfo
ts time.Time
}
const cacheTTL = 30 * time.Minute
const cacheMaxEntries = 200
func (c *Client) getCached(url string) *domain.MediaInfo {
c.cacheMu.Lock()
defer c.cacheMu.Unlock()
entry, ok := c.cache[url]
if !ok || time.Since(entry.ts) > cacheTTL {
delete(c.cache, url)
return nil
}
return entry.info
}
func (c *Client) setCache(url string, info *domain.MediaInfo) {
c.cacheMu.Lock()
defer c.cacheMu.Unlock()
if len(c.cache) >= cacheMaxEntries {
c.cache = make(map[string]*cachedMediaInfo)
}
c.cache[url] = &cachedMediaInfo{info: info, ts: time.Now()}
} }
// NewClient creates a new yt-dlp client. It checks for the binary. // NewClient creates a new yt-dlp client. It checks for the binary.
@@ -29,7 +60,11 @@ func NewClient(cookiesFile string) (*Client, error) {
return nil, fmt.Errorf("yt-dlp not found in PATH: %w", err) return nil, fmt.Errorf("yt-dlp not found in PATH: %w", err)
} }
slog.Info("yt-dlp found", "path", path) slog.Info("yt-dlp found", "path", path)
return &Client{binaryPath: path, cookiesFile: cookiesFile}, nil return &Client{
binaryPath: path,
cookiesFile: cookiesFile,
cache: make(map[string]*cachedMediaInfo),
}, nil
} }
type ytdlpFormat struct { type ytdlpFormat struct {
@@ -70,6 +105,11 @@ type ytdlpPlaylistEntry struct {
// GetMediaInfo fetches format information for a given URL. // GetMediaInfo fetches format information for a given URL.
func (c *Client) GetMediaInfo(url string) (*domain.MediaInfo, error) { func (c *Client) GetMediaInfo(url string) (*domain.MediaInfo, error) {
if cached := c.getCached(url); cached != nil {
slog.Debug("cache hit", "url", url)
return cached, nil
}
args := []string{ args := []string{
"-J", "--no-download", "-J", "--no-download",
"--no-warnings", "--no-warnings",
@@ -106,6 +146,7 @@ func (c *Client) GetMediaInfo(url string) (*domain.MediaInfo, error) {
if mi.IsPlaylist { if mi.IsPlaylist {
mi.PlaylistCount = info.PlaylistCount mi.PlaylistCount = info.PlaylistCount
mi.PlaylistEntries = c.parsePlaylistEntries(info.Entries) mi.PlaylistEntries = c.parsePlaylistEntries(info.Entries)
c.setCache(url, mi)
return mi, nil return mi, nil
} }
@@ -124,11 +165,17 @@ func (c *Client) GetMediaInfo(url string) (*domain.MediaInfo, error) {
} }
mi.EstimatedSize = maxSize mi.EstimatedSize = maxSize
c.setCache(url, mi)
return mi, nil return mi, nil
} }
// GetPlaylistInfo fetches a flat playlist listing (fast, no format detail). // GetPlaylistInfo fetches a flat playlist listing (fast, no format detail).
func (c *Client) GetPlaylistInfo(url string) (*domain.MediaInfo, error) { func (c *Client) GetPlaylistInfo(url string) (*domain.MediaInfo, error) {
if cached := c.getCached(url); cached != nil {
slog.Debug("cache hit", "url", url)
return cached, nil
}
args := []string{ args := []string{
"-J", "--flat-playlist", "--no-download", "-J", "--flat-playlist", "--no-download",
"--no-warnings", "--no-warnings",
@@ -159,6 +206,7 @@ func (c *Client) GetPlaylistInfo(url string) (*domain.MediaInfo, error) {
PlaylistEntries: c.parsePlaylistEntries(info.Entries), PlaylistEntries: c.parsePlaylistEntries(info.Entries),
} }
c.setCache(url, mi)
return mi, nil return mi, nil
} }
@@ -376,6 +424,36 @@ func (c *Client) SearchYoutube(query string) (url, title string, err error) {
return info.Entries[0].URL, info.Entries[0].Title, nil return info.Entries[0].URL, info.Entries[0].Title, nil
} }
// Search does a yt-dlp search and returns entries.
func (c *Client) Search(query, source string, limit int) ([]domain.PlaylistEntry, error) {
searchPrefix := "ytsearch"
if source == "ytmusic" {
searchPrefix = "ytmusicsearch"
}
args := []string{
"--flat-playlist", "-J", "--no-download",
"--no-warnings",
fmt.Sprintf("%s%d:%s", searchPrefix, limit, query),
}
if c.cookiesFile != "" {
args = append(args, "--cookies", c.cookiesFile)
}
slog.Info("searching", "query", query, "source", source)
out, err := c.run(args)
if err != nil {
return nil, fmt.Errorf("search failed: %w", err)
}
var info ytdlpInfo
if err := json.Unmarshal(out, &info); err != nil {
return nil, fmt.Errorf("parse search result: %w", err)
}
return c.parsePlaylistEntries(info.Entries), nil
}
func (c *Client) run(args []string) ([]byte, error) { func (c *Client) run(args []string) ([]byte, error) {
cmd := exec.Command(c.binaryPath, args...) cmd := exec.Command(c.binaryPath, args...)
var stderr bytes.Buffer var stderr bytes.Buffer

View File

@@ -6,7 +6,7 @@ const (
RateLimitInterval = 1 * time.Second RateLimitInterval = 1 * time.Second
ProgressThreshold = 5.0 ProgressThreshold = 5.0
MaxFileSizeEnvKey = "MAX_FILE_SIZE_MB" MaxFileSizeEnvKey = "MAX_FILE_SIZE_MB"
DefaultMaxFileSize = 2000 DefaultMaxFileSize = 50
DefaultDailyQuotaMB = 100 DefaultDailyQuotaMB = 100
DefaultWeeklyQuotaMB = 500 DefaultWeeklyQuotaMB = 500
DefaultMonthlyQuotaMB = 2000 DefaultMonthlyQuotaMB = 2000

View File

@@ -129,6 +129,7 @@ type DownloadJob struct {
Title string Title string
QuotaApplied bool QuotaApplied bool
StderrLog string StderrLog string
Done chan struct{}
CancelCh chan struct{} CancelCh chan struct{}
Cmd *exec.Cmd Cmd *exec.Cmd
} }
@@ -177,3 +178,12 @@ type PlaylistState struct {
PerPage int PerPage int
Selected map[string]bool Selected map[string]bool
} }
// SearchState tracks the user's current search results flow.
type SearchState struct {
Query string
Entries []PlaylistEntry
Page int
PerPage int
Source string // "youtube" or "ytmusic"
}

View File

@@ -3,8 +3,10 @@ package handler
import ( import (
"context" "context"
"fmt" "fmt"
"io"
"log/slog" "log/slog"
"math" "math"
"net/http"
"os" "os"
"path/filepath" "path/filepath"
"strings" "strings"
@@ -29,23 +31,24 @@ type stepState struct {
Container string Container string
FormatIDs []string FormatIDs []string
PlaylistState *domain.PlaylistState PlaylistState *domain.PlaylistState
SearchState *domain.SearchState
} }
func (h *Handler) getUserStepState(chatID int64) *stepState { func (h *Handler) getUserStepState(chatID int64) *stepState {
h.mu.Lock() h.stepStatesMu.Lock()
defer h.mu.Unlock() defer h.stepStatesMu.Unlock()
return h.stepStates[chatID] return h.stepStates[chatID]
} }
func (h *Handler) setUserStepState(chatID int64, ss *stepState) { func (h *Handler) setUserStepState(chatID int64, ss *stepState) {
h.mu.Lock() h.stepStatesMu.Lock()
defer h.mu.Unlock() defer h.stepStatesMu.Unlock()
h.stepStates[chatID] = ss h.stepStates[chatID] = ss
} }
func (h *Handler) clearStepState(chatID int64) { func (h *Handler) clearStepState(chatID int64) {
h.mu.Lock() h.stepStatesMu.Lock()
defer h.mu.Unlock() defer h.stepStatesMu.Unlock()
delete(h.stepStates, chatID) delete(h.stepStates, chatID)
} }
@@ -354,6 +357,8 @@ func (h *Handler) handleContainerSelection(ctx context.Context, chatID int64, ms
return return
} }
dlCtx, dlCancel := context.WithCancel(context.Background())
job := &domain.DownloadJob{ job := &domain.DownloadJob{
ID: h.generateJobID(), ID: h.generateJobID(),
UserID: userID, UserID: userID,
@@ -368,17 +373,33 @@ func (h *Handler) handleContainerSelection(ctx context.Context, chatID int64, ms
Status: domain.StatusDownloading, Status: domain.StatusDownloading,
Title: ss.MediaInfo.Title, Title: ss.MediaInfo.Title,
FileSize: fileSize, FileSize: fileSize,
Done: make(chan struct{}),
} }
h.editText(ctx, chatID, msgID, "Processing...", nil) h.editText(ctx, chatID, msgID, "Processing...", nil)
h.setActiveJob(userID, job) h.setActiveJob(userID, job)
h.setCancelFunc(userID, dlCancel)
h.clearStepState(chatID) h.clearStepState(chatID)
go h.runDownload(context.Background(), job) go h.runDownload(dlCtx, job)
} }
// runDownload manages the download lifecycle, progress, and completion. // runDownload manages the download lifecycle, progress, and completion.
func (h *Handler) runDownload(ctx context.Context, job *domain.DownloadJob) { 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, "", "") kb := progressKeyboard(0, "", "")
msg, err := h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{ msg, err := h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{
ChatID: job.ChatID, ChatID: job.ChatID,
@@ -395,40 +416,45 @@ func (h *Handler) runDownload(ctx context.Context, job *domain.DownloadJob) {
if err != nil { if err != nil {
h.editText(ctx, job.ChatID, job.MessageID, fmt.Sprintf("Download failed: %s", err.Error()), nil) h.editText(ctx, job.ChatID, job.MessageID, fmt.Sprintf("Download failed: %s", err.Error()), nil)
job.Status = domain.StatusFailed job.Status = domain.StatusFailed
h.mu.Lock() h.activeJobsMu.Lock()
delete(h.activeJobs, job.UserID) delete(h.activeJobs, job.UserID)
h.mu.Unlock() h.activeJobsMu.Unlock()
return return
} }
lastProgress := 0.0 lastProgress := 0.0
for upd := range updates { loop:
for {
select {
case upd, ok := <-updates:
if !ok {
break loop
}
slog.Info("progress update", "pct", upd.Progress, "speed", upd.Speed, "status", upd.Status) slog.Info("progress update", "pct", upd.Progress, "speed", upd.Speed, "status", upd.Status)
h.mu.Lock() h.activeJobsMu.Lock()
currentJob := h.activeJobs[job.UserID] currentJob := h.activeJobs[job.UserID]
h.mu.Unlock() h.activeJobsMu.Unlock()
if currentJob == nil { if currentJob == nil {
return return
} }
if upd.Status == domain.StatusCancelled { if upd.Status == domain.StatusCancelled {
h.editText(ctx, job.ChatID, job.MessageID, "Download cancelled.", nil) h.editText(ctx, job.ChatID, job.MessageID, "Download cancelled.", nil)
h.mu.Lock() h.activeJobsMu.Lock()
delete(h.activeJobs, job.UserID) delete(h.activeJobs, job.UserID)
h.mu.Unlock() h.activeJobsMu.Unlock()
return return
} }
if upd.Status == domain.StatusFailed { if upd.Status == domain.StatusFailed {
h.editText(ctx, job.ChatID, job.MessageID, "Download failed.", nil) h.editText(ctx, job.ChatID, job.MessageID, "Download failed.", nil)
h.mu.Lock() h.activeJobsMu.Lock()
delete(h.activeJobs, job.UserID) delete(h.activeJobs, job.UserID)
h.mu.Unlock() h.activeJobsMu.Unlock()
return return
} }
// Always show the first real progress update, then throttle at 5% intervals
if upd.Progress > 0 && lastProgress == 0 { if upd.Progress > 0 && lastProgress == 0 {
lastProgress = upd.Progress lastProgress = upd.Progress
} else if upd.Progress-lastProgress < domain.ProgressThreshold && upd.Status != domain.StatusCompleted { } else if upd.Progress-lastProgress < domain.ProgressThreshold && upd.Status != domain.StatusCompleted {
@@ -445,11 +471,22 @@ func (h *Handler) runDownload(ctx context.Context, job *domain.DownloadJob) {
kb := progressKeyboard(upd.Progress, upd.Speed, upd.ETA) kb := progressKeyboard(upd.Progress, upd.Speed, upd.ETA)
text := fmt.Sprintf("Downloading: %s", job.Title) text := fmt.Sprintf("Downloading: %s", job.Title)
h.editText(ctx, job.ChatID, job.MessageID, text, &kb) h.editText(ctx, job.ChatID, job.MessageID, text, &kb)
case <-ctx.Done():
dl.CancelDownload(job)
for range updates {
}
break loop
}
} }
h.mu.Lock() h.activeJobsMu.Lock()
delete(h.activeJobs, job.UserID) 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 { if job.Status == domain.StatusCompleted {
pattern := filepath.Join(h.downloadDir, job.ID+"_*") pattern := filepath.Join(h.downloadDir, job.ID+"_*")
@@ -517,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) { func (h *Handler) handleCancelDownload(ctx context.Context, chatID int64, msgID int, cbID string, userID int64) {
defer h.answerCb(ctx, cbID) 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) job := h.getActiveJob(userID)
if job != nil { if job != nil {
dl.CancelDownload(job) dl.CancelDownload(job)
h.mu.Lock() h.activeJobsMu.Lock()
delete(h.activeJobs, userID) delete(h.activeJobs, userID)
h.mu.Unlock() h.activeJobsMu.Unlock()
h.editText(ctx, chatID, msgID, "Download cancelled.", nil) h.editText(ctx, chatID, msgID, "Download cancelled.", nil)
return return
} }
@@ -646,21 +690,71 @@ func (h *Handler) sendMediaPreview(ctx context.Context, chatID int64, info *doma
} }
if info.Thumbnail != "" { if info.Thumbnail != "" {
photo := &models.InputFileString{Data: info.Thumbnail} thumbPath, err := downloadFile(info.Thumbnail, h.downloadDir)
_, err := h.Bot.SendPhoto(ctx, &tgbot.SendPhotoParams{ if err != nil {
ChatID: chatID, slog.Warn("download thumbnail failed, falling back to text", "url", info.Thumbnail, "error", err)
Photo: photo, h.sendText(ctx, chatID, caption)
Caption: 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 { if err == nil {
return 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) 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 { func formatDuration(seconds float64) string {
if seconds <= 0 { if seconds <= 0 {
return "0:00" return "0:00"

View File

@@ -25,6 +25,7 @@ type PendingKind string
const ( const (
PendingURL PendingKind = "url" PendingURL PendingKind = "url"
PendingDRM PendingKind = "drm" PendingDRM PendingKind = "drm"
PendingSearch PendingKind = "search"
) )
type PendingInput struct { type PendingInput struct {
@@ -35,18 +36,29 @@ type PendingInput struct {
Data map[string]string Data map[string]string
} }
const maxConcurrentDownloads = 3
type Handler struct { type Handler struct {
Bot *tgbot.Bot Bot *tgbot.Bot
Repo repo.Repository Repo repo.Repository
YtDlp *dl.Client YtDlp *dl.Client
AllowedUsers map[int64]bool AllowedUsers map[int64]bool
activeJobs map[int64]*domain.DownloadJob activeJobs map[int64]*domain.DownloadJob
stepStates map[int64]*stepState stepStates map[int64]*stepState
lastMsgTime map[int64]time.Time lastMsgTime map[int64]time.Time
pendingInput map[int64]*PendingInput pendingInput map[int64]*PendingInput
mu sync.Mutex cancelFuncs map[int64]context.CancelFunc
activeJobsMu sync.Mutex
stepStatesMu sync.Mutex
rateLimiterMu sync.Mutex
pendingInputMu sync.Mutex
cancelFuncsMu sync.Mutex
downloadDir string downloadDir string
maxFileSize int64 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 { func NewHandler(bot *tgbot.Bot, store repo.Repository, ytdlp *dl.Client, allowed map[int64]bool, downloadDir string, maxFileSize int64) *Handler {
@@ -59,8 +71,10 @@ func NewHandler(bot *tgbot.Bot, store repo.Repository, ytdlp *dl.Client, allowed
stepStates: make(map[int64]*stepState), stepStates: make(map[int64]*stepState),
lastMsgTime: make(map[int64]time.Time), lastMsgTime: make(map[int64]time.Time),
pendingInput: make(map[int64]*PendingInput), pendingInput: make(map[int64]*PendingInput),
cancelFuncs: make(map[int64]context.CancelFunc),
downloadDir: downloadDir, downloadDir: downloadDir,
maxFileSize: maxFileSize, maxFileSize: maxFileSize,
downloadSem: make(chan struct{}, maxConcurrentDownloads),
} }
go h.periodicCleanup() go h.periodicCleanup()
return h return h
@@ -69,27 +83,29 @@ func NewHandler(bot *tgbot.Bot, store repo.Repository, ytdlp *dl.Client, allowed
func (h *Handler) periodicCleanup() { func (h *Handler) periodicCleanup() {
for { for {
time.Sleep(domain.CleanupInterval) time.Sleep(domain.CleanupInterval)
h.mu.Lock()
// Clean rate limit entries func() {
h.rateLimiterMu.Lock()
defer h.rateLimiterMu.Unlock()
cutoff := time.Now().Add(-domain.RateLimitInterval * 2) cutoff := time.Now().Add(-domain.RateLimitInterval * 2)
for uid, t := range h.lastMsgTime { for uid, t := range h.lastMsgTime {
if t.Before(cutoff) { if t.Before(cutoff) {
delete(h.lastMsgTime, uid) delete(h.lastMsgTime, uid)
} }
} }
}()
// Clean stale pending inputs func() {
h.pendingInputMu.Lock()
defer h.pendingInputMu.Unlock()
pendingCutoff := time.Now().Add(-5 * time.Minute) pendingCutoff := time.Now().Add(-5 * time.Minute)
for chatID, pi := range h.pendingInput { for chatID, pi := range h.pendingInput {
if pi.CreatedAt.Before(pendingCutoff) { if pi.CreatedAt.Before(pendingCutoff) {
delete(h.pendingInput, chatID) delete(h.pendingInput, chatID)
} }
} }
}()
h.mu.Unlock()
// Clean old temp files (>1h)
h.cleanOldTempFiles() h.cleanOldTempFiles()
} }
} }
@@ -115,8 +131,8 @@ func (h *Handler) cleanOldTempFiles() {
} }
func (h *Handler) isRateLimited(userID int64) bool { func (h *Handler) isRateLimited(userID int64) bool {
h.mu.Lock() h.rateLimiterMu.Lock()
defer h.mu.Unlock() defer h.rateLimiterMu.Unlock()
now := time.Now() now := time.Now()
if last, ok := h.lastMsgTime[userID]; ok && now.Sub(last) < domain.RateLimitInterval { if last, ok := h.lastMsgTime[userID]; ok && now.Sub(last) < domain.RateLimitInterval {
return true return true
@@ -146,19 +162,19 @@ func (h *Handler) HandleMessage(ctx context.Context, msg *models.Message) {
slog.Info("message", "chat_id", msg.Chat.ID, "text", msg.Text) slog.Info("message", "chat_id", msg.Chat.ID, "text", msg.Text)
if msg.Text[0] == '/' { if msg.Text[0] == '/' {
h.mu.Lock() h.pendingInputMu.Lock()
delete(h.pendingInput, msg.Chat.ID) delete(h.pendingInput, msg.Chat.ID)
h.mu.Unlock() h.pendingInputMu.Unlock()
h.routeCommand(ctx, msg, userID) h.routeCommand(ctx, msg, userID)
return return
} }
h.mu.Lock() h.pendingInputMu.Lock()
pi, hasPI := h.pendingInput[msg.Chat.ID] pi, hasPI := h.pendingInput[msg.Chat.ID]
if hasPI { if hasPI {
delete(h.pendingInput, msg.Chat.ID) delete(h.pendingInput, msg.Chat.ID)
} }
h.mu.Unlock() h.pendingInputMu.Unlock()
if hasPI { if hasPI {
h.processPendingInput(ctx, msg, userID, pi) h.processPendingInput(ctx, msg, userID, pi)
@@ -185,6 +201,10 @@ func (h *Handler) routeCommand(ctx context.Context, msg *models.Message, userID
h.handleHistory(ctx, msg, userID) h.handleHistory(ctx, msg, userID)
case "/download": case "/download":
h.handleDownloadCommand(ctx, msg, userID) h.handleDownloadCommand(ctx, msg, userID)
case "/search":
h.handleSearchCommand(ctx, msg, userID)
case "/music":
h.handleMusicSearchCommand(ctx, msg, userID)
default: default:
h.sendText(ctx, msg.Chat.ID, "Unknown command") h.sendText(ctx, msg.Chat.ID, "Unknown command")
} }
@@ -207,9 +227,9 @@ func (h *Handler) HandleCallback(ctx context.Context, cb *models.CallbackQuery)
} }
slog.Info("callback", "chat_id", chatID, "data", cb.Data) slog.Info("callback", "chat_id", chatID, "data", cb.Data)
h.mu.Lock() h.pendingInputMu.Lock()
delete(h.pendingInput, chatID) delete(h.pendingInput, chatID)
h.mu.Unlock() h.pendingInputMu.Unlock()
msgID := cb.Message.Message.ID msgID := cb.Message.Message.ID
data := cb.Data data := cb.Data
@@ -226,6 +246,9 @@ func (h *Handler) HandleCallback(ctx context.Context, cb *models.CallbackQuery)
case data == "download": case data == "download":
h.handleDownloadPrompt(ctx, chatID, msgID, cb.ID, userID) h.handleDownloadPrompt(ctx, chatID, msgID, cb.ID, userID)
h.answerCb(ctx, cb.ID) h.answerCb(ctx, cb.ID)
case data == "search":
h.promptForInput(ctx, chatID, msgID, userID, PendingSearch, "Enter a YouTube search query:", map[string]string{"source": "youtube"})
h.answerCb(ctx, cb.ID)
case data == "settings": case data == "settings":
h.settingsCallback(ctx, chatID, msgID, cb.ID, userID) h.settingsCallback(ctx, chatID, msgID, cb.ID, userID)
case data == "help": case data == "help":
@@ -280,11 +303,13 @@ func (h *Handler) routePrefixedCallback(ctx context.Context, chatID int64, msgID
h.handleSettingsSet(ctx, chatID, msgID, userID, data[13:]) h.handleSettingsSet(ctx, chatID, msgID, userID, data[13:])
case strings.HasPrefix(data, "settings_"): case strings.HasPrefix(data, "settings_"):
h.handleSettingsSelection(ctx, chatID, msgID, userID, data[9:]) h.handleSettingsSelection(ctx, chatID, msgID, userID, data[9:])
case strings.HasPrefix(data, "search_"):
h.handleSearchCallback(ctx, chatID, msgID, userID, data[7:])
} }
} }
func (h *Handler) setPending(ctx context.Context, chatID int64, msgID int, userID int64, kind PendingKind, data map[string]string) { 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{ h.pendingInput[chatID] = &PendingInput{
Kind: kind, Kind: kind,
UserID: userID, UserID: userID,
@@ -292,7 +317,7 @@ func (h *Handler) setPending(ctx context.Context, chatID int64, msgID int, userI
MsgID: msgID, MsgID: msgID,
Data: data, Data: data,
} }
h.mu.Unlock() h.pendingInputMu.Unlock()
} }
func (h *Handler) processPendingInput(ctx context.Context, msg *models.Message, userID int64, pi *PendingInput) { func (h *Handler) processPendingInput(ctx context.Context, msg *models.Message, userID int64, pi *PendingInput) {
@@ -302,6 +327,12 @@ func (h *Handler) processPendingInput(ctx context.Context, msg *models.Message,
h.handleURLInput(ctx, msg.Chat.ID, text, userID) h.handleURLInput(ctx, msg.Chat.ID, text, userID)
case PendingDRM: case PendingDRM:
h.handleDRMSearch(ctx, msg.Chat.ID, text, userID) h.handleDRMSearch(ctx, msg.Chat.ID, text, userID)
case PendingSearch:
source := "youtube"
if pi.Data != nil && pi.Data["source"] != "" {
source = pi.Data["source"]
}
h.handleSearchQuery(ctx, msg.Chat.ID, text, userID, source)
default: default:
h.sendText(ctx, msg.Chat.ID, "Unknown input type") h.sendText(ctx, msg.Chat.ID, "Unknown input type")
} }
@@ -396,26 +427,40 @@ func (h *Handler) backToMenu(ctx context.Context, chatID int64, msgID int, cbID
} }
func (h *Handler) clearActiveJob(userID int64) { func (h *Handler) clearActiveJob(userID int64) {
h.mu.Lock() h.activeJobsMu.Lock()
defer h.mu.Unlock() job, ok := h.activeJobs[userID]
if job, ok := h.activeJobs[userID]; ok { if ok {
dl.CancelDownload(job) dl.CancelDownload(job)
delete(h.activeJobs, userID) 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 { func (h *Handler) getActiveJob(userID int64) *domain.DownloadJob {
h.mu.Lock() h.activeJobsMu.Lock()
defer h.mu.Unlock() defer h.activeJobsMu.Unlock()
return h.activeJobs[userID] return h.activeJobs[userID]
} }
func (h *Handler) setActiveJob(userID int64, job *domain.DownloadJob) { func (h *Handler) setActiveJob(userID int64, job *domain.DownloadJob) {
h.mu.Lock() h.activeJobsMu.Lock()
defer h.mu.Unlock() defer h.activeJobsMu.Unlock()
h.activeJobs[userID] = job 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 { func formatBytes(bytes int64) string {
const unit = 1024 const unit = 1024
if bytes < unit { if bytes < unit {
@@ -439,7 +484,7 @@ func formatBytes(bytes int64) string {
} }
func (h *Handler) generateJobID() string { func (h *Handler) generateJobID() string {
return uuid.NewString()[:8] return uuid.NewString()
} }
// handleStart shows the main menu and cleans up any stale state. // handleStart shows the main menu and cleans up any stale state.
@@ -457,6 +502,8 @@ func (h *Handler) handleHelp(ctx context.Context, msg *models.Message) {
/settings - Open settings /settings - Open settings
/help - Show this help /help - Show this help
/history - View download history /history - View download history
/search - Search YouTube
/music - Search YouTube Music
Just send me a URL to start downloading.` Just send me a URL to start downloading.`
h.sendText(ctx, msg.Chat.ID, text) h.sendText(ctx, msg.Chat.ID, text)

View File

@@ -14,6 +14,7 @@ func mainKeyboard() models.InlineKeyboardMarkup {
InlineKeyboard: [][]models.InlineKeyboardButton{ InlineKeyboard: [][]models.InlineKeyboardButton{
{ {
{Text: "Download", CallbackData: "download"}, {Text: "Download", CallbackData: "download"},
{Text: "Search", CallbackData: "search"},
{Text: "Settings", CallbackData: "settings"}, {Text: "Settings", CallbackData: "settings"},
}, },
{ {

View File

@@ -4,7 +4,6 @@ import (
"context" "context"
"fmt" "fmt"
"strconv" "strconv"
"time"
"uptodownBot/internal/domain" "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{ job := &domain.DownloadJob{
ID: h.generateJobID(), ID: h.generateJobID(),
UserID: userID, UserID: userID,
@@ -176,19 +177,17 @@ func (h *Handler) handlePlaylistConfirm(ctx context.Context, chatID int64, msgID
Status: domain.StatusDownloading, Status: domain.StatusDownloading,
Title: entry.Title, Title: entry.Title,
FileSize: format.Filesize, FileSize: format.Filesize,
Done: make(chan struct{}),
} }
h.setActiveJob(userID, job) h.setActiveJob(userID, job)
h.runDownload(context.Background(), job) h.setCancelFunc(userID, dlCancel)
// Wait for active job to complete h.runDownload(dlCtx, job)
for {
j := h.getActiveJob(userID) // runDownload is synchronous — it blocks until the download completes.
if j == nil || j.Status == domain.StatusCompleted || j.Status == domain.StatusFailed || j.Status == domain.StatusCancelled { // The Done channel is closed inside runDownload when it exits.
break // No busy-polling needed.
}
time.Sleep(500 * time.Millisecond)
}
} }
h.sendText(ctx, chatID, "All playlist downloads completed.") h.sendText(ctx, chatID, "All playlist downloads completed.")

136
internal/handler/search.go Normal file
View File

@@ -0,0 +1,136 @@
package handler
import (
"context"
"fmt"
"strconv"
"strings"
"github.com/go-telegram/bot/models"
"uptodownBot/internal/domain"
)
func (h *Handler) handleSearchCommand(ctx context.Context, msg *models.Message, userID int64) {
h.promptForInput(ctx, msg.Chat.ID, 0, userID, PendingSearch,
"Enter a YouTube search query:", map[string]string{"source": "youtube"})
}
func (h *Handler) handleMusicSearchCommand(ctx context.Context, msg *models.Message, userID int64) {
h.promptForInput(ctx, msg.Chat.ID, 0, userID, PendingSearch,
"Enter a YouTube Music search query:", map[string]string{"source": "ytmusic"})
}
func (h *Handler) handleSearchQuery(ctx context.Context, chatID int64, text string, userID int64, source string) {
results, err := h.YtDlp.Search(text, source, 10)
if err != nil {
h.sendText(ctx, chatID, fmt.Sprintf("Search failed: %s", err.Error()))
return
}
if len(results) == 0 {
h.sendText(ctx, chatID, "No results found.")
return
}
state := &domain.SearchState{
Query: text,
Entries: results,
Page: 0,
PerPage: 5,
Source: source,
}
ss := &stepState{SearchState: state}
h.setUserStepState(chatID, ss)
kb := searchResultsKeyboard(state)
h.sendWithKB(ctx, chatID, fmt.Sprintf("Search results for '%s':", text), kb)
}
func (h *Handler) handleSearchCallback(ctx context.Context, chatID int64, msgID int, userID int64, data string) {
switch {
case data == "refresh":
ss := h.getUserStepState(chatID)
if ss == nil || ss.SearchState == nil {
return
}
h.handleSearchQuery(ctx, chatID, ss.SearchState.Query, userID, ss.SearchState.Source)
case strings.HasPrefix(data, "page_"):
h.handleSearchPage(ctx, chatID, msgID, userID, data[5:])
case strings.HasPrefix(data, "select_"):
h.handleSearchSelect(ctx, chatID, msgID, userID, data[7:])
}
}
func (h *Handler) handleSearchPage(ctx context.Context, chatID int64, msgID int, userID int64, pageStr string) {
ss := h.getUserStepState(chatID)
if ss == nil || ss.SearchState == nil {
return
}
page, err := strconv.Atoi(pageStr)
if err != nil {
return
}
totalPages := (len(ss.SearchState.Entries) + ss.SearchState.PerPage - 1) / ss.SearchState.PerPage
if page < 0 || page >= totalPages {
return
}
ss.SearchState.Page = page
h.setUserStepState(chatID, ss)
kb := searchResultsKeyboard(ss.SearchState)
h.editText(ctx, chatID, msgID, fmt.Sprintf("Search results for '%s':", ss.SearchState.Query), &kb)
}
func (h *Handler) handleSearchSelect(ctx context.Context, chatID int64, msgID int, userID int64, entryURL string) {
h.clearStepState(chatID)
h.editText(ctx, chatID, msgID, "Fetching media info...", nil)
h.handleURLInput(ctx, chatID, entryURL, userID)
}
func searchResultsKeyboard(state *domain.SearchState) models.InlineKeyboardMarkup {
totalPages := (len(state.Entries) + state.PerPage - 1) / state.PerPage
if totalPages < 1 {
totalPages = 1
}
rows := [][]models.InlineKeyboardButton{}
start := state.Page * state.PerPage
end := start + state.PerPage
if end > len(state.Entries) {
end = len(state.Entries)
}
for i := start; i < end; i++ {
e := state.Entries[i]
title := e.Title
if len(title) > 40 {
title = title[:37] + "..."
}
label := title
if e.Duration > 0 {
label += fmt.Sprintf(" (%s)", formatDuration(e.Duration))
}
rows = append(rows, []models.InlineKeyboardButton{
{Text: label, CallbackData: "search_select_" + e.URL},
})
}
navRow := []models.InlineKeyboardButton{}
navRow = append(navRow, models.InlineKeyboardButton{
Text: "<", CallbackData: fmt.Sprintf("search_page_%d", state.Page-1),
})
navRow = append(navRow, models.InlineKeyboardButton{
Text: fmt.Sprintf("Page %d/%d", state.Page+1, totalPages), CallbackData: "noop",
})
navRow = append(navRow, models.InlineKeyboardButton{
Text: ">", CallbackData: fmt.Sprintf("search_page_%d", state.Page+1),
})
rows = append(rows, navRow)
rows = append(rows, []models.InlineKeyboardButton{
{Text: "New Search", CallbackData: "search_refresh"},
{Text: "Cancel", CallbackData: "cancel_dl"},
})
return models.InlineKeyboardMarkup{InlineKeyboard: rows}
}