refactor: split mutex, add download queue, wire context cancellation
All checks were successful
CI / build (push) Successful in 53s
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:
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user