diff --git a/app/distortioner.go b/app/distortioner.go index fb5925d..e8eebb0 100644 --- a/app/distortioner.go +++ b/app/distortioner.go @@ -43,7 +43,7 @@ func (d DistorterBot) handleAnimationDistortion(c tb.Context) error { } //TODO: Jesus, just find the time to refactor all of this already - d.videoWorker.Submit(m.Chat.ID, func() { + err := d.videoWorker.Submit(m.Chat.ID, func() { progressMessage, filename, output, err := d.HandleAnimationCommon(c) failed := err != nil if failed { @@ -64,6 +64,10 @@ func (d DistorterBot) handleAnimationDistortion(c tb.Context) error { err = d.SendMessageWithRepeater(c, distorted) d.DoneMessageWithRepeater(b, progressMessage, failed) }) + if err != nil { + d.SendMessageWithRepeater(c, err.Error()) + return nil + } if d.videoWorker.IsBusy() { d.SendMessageWithRepeater(c, distorters.Queued) } @@ -138,7 +142,7 @@ func (d DistorterBot) handleVideoDistortion(c tb.Context) error { return d.SendMessageWithRepeater(c, tools.FormatRateLimitResponse(diff)) } - d.videoWorker.Submit(m.Chat.ID, func() { + err := d.videoWorker.Submit(m.Chat.ID, func() { output, progressMessage, err := d.HandleVideoCommon(c) failed := err != nil if failed { @@ -157,6 +161,10 @@ func (d DistorterBot) handleVideoDistortion(c tb.Context) error { d.logger.Error(err) } }) + if err != nil { + d.SendMessageWithRepeater(c, err.Error()) + return nil + } if d.videoWorker.IsBusy() { d.SendMessageWithRepeater(c, distorters.Queued) } @@ -172,7 +180,7 @@ func (d DistorterBot) handleVideoNoteDistortion(c tb.Context) error { return d.SendMessageWithRepeater(c, tools.FormatRateLimitResponse(diff)) } - d.videoWorker.Submit(m.Chat.ID, func() { + err := d.videoWorker.Submit(m.Chat.ID, func() { output, progressMessage, err := d.HandleVideoCommon(c) failed := err != nil if failed { @@ -187,6 +195,10 @@ func (d DistorterBot) handleVideoNoteDistortion(c tb.Context) error { err = d.SendMessageWithRepeater(c, distorted) d.DoneMessageWithRepeater(b, progressMessage, failed) }) + if err != nil { + d.SendMessageWithRepeater(c, err.Error()) + return nil + } if d.videoWorker.IsBusy() { d.SendMessageWithRepeater(c, distorters.Queued) } diff --git a/app/queue/honest_priority_queue.go b/app/queue/honest_priority_queue.go index 17781fa..9d8884d 100644 --- a/app/queue/honest_priority_queue.go +++ b/app/queue/honest_priority_queue.go @@ -2,6 +2,7 @@ package queue import ( "container/heap" + "github.com/pkg/errors" "sync" "time" ) @@ -95,13 +96,21 @@ func (hjq *HonestJobQueue) Pop() *Job { return job } -func (hjq *HonestJobQueue) Push(userID int64, runnable func()) { +func (hjq *HonestJobQueue) Push(userID int64, runnable func()) error { hjq.mu.Lock() defer hjq.mu.Unlock() - hjq.users[userID]++ priority := hjq.users[userID] + if priority > 2 { + hjq.users[userID]-- + return errors.New("you're distorting videos too often, wait until the previous ones have been processed") + } + + hjq.users[userID] = priority + 1 + job := newJob(userID, priority, runnable) heap.Push(&hjq.queue, &job) + + return nil } diff --git a/app/tools/video_worker.go b/app/tools/video_worker.go index 796e2df..4da558a 100644 --- a/app/tools/video_worker.go +++ b/app/tools/video_worker.go @@ -33,9 +33,13 @@ func (vw *VideoWorker) run() { } } -func (vw *VideoWorker) Submit(userID int64, runnable func()) { - vw.queue.Push(userID, runnable) +func (vw *VideoWorker) Submit(userID int64, runnable func()) error { + err := vw.queue.Push(userID, runnable) + if err != nil { + return err + } vw.messenger <- nil // let goroutines know that there's something in the queue + return nil } func (vw *VideoWorker) Shutdown() { @@ -47,5 +51,5 @@ func (vw *VideoWorker) QueueStats() (int, int) { } func (vw *VideoWorker) IsBusy() bool { - return len(vw.messenger) > 0 + return vw.queue.Len() > 0 }