introduce some sort of limit

graynk · Jul 20, 2024 · abad426 · abad426
1 parent 02b7bef
commit abad426
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 8 deletions.
diff --git 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
@@ -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
@@ -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
 }