From 8b820f6bb8e0ad45c71743ff699142ac343341bb Mon Sep 17 00:00:00 2001
From: Nikita Karpukhin <gray@graynk.space>
Date: Sun, 21 Jul 2024 12:29:31 +0200
Subject: [PATCH] add maintenance mode and queue upper limit

---
 Dockerfile                         |  2 +-
 app/distortioner.go                | 10 ++++++++++
 app/queue/honest_priority_queue.go | 28 +++++++++++++++++++++++-----
 app/tools/video_worker.go          |  4 ++++
 4 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 3f71f62..cef2c36 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,4 +1,4 @@
-FROM golang:1.21-bullseye as build
+FROM golang:1.22-bullseye as build
 WORKDIR /go/src/distortioner
 COPY app .
 RUN go test ./...
diff --git a/app/distortioner.go b/app/distortioner.go
index fcb19d0..994f81a 100644
--- a/app/distortioner.go
+++ b/app/distortioner.go
@@ -307,6 +307,14 @@ func (d DistorterBot) handleQueueStats(c tb.Context) error {
 	return c.Reply(fmt.Sprintf("Currently in queue: %d requests from %d users", length, users))
 }
 
+func (d DistorterBot) handleMaintenance(c tb.Context) error {
+	if c.Message().Sender.ID != d.adminID {
+		return nil
+	}
+	currentMode := d.videoWorker.ToggleMaintenance()
+	return c.Reply(fmt.Sprintf("Maintenance on: %v", currentMode))
+}
+
 func main() {
 	lg, err := zap.NewProduction()
 	if err != nil {
@@ -403,6 +411,8 @@ func main() {
 
 	b.Handle("/queue", d.handleQueueStats)
 
+	b.Handle("/maintenance", d.handleMaintenance)
+
 	b.Handle("/distort", d.ApplyShutdownMiddleware(d.handleReplyDistortion))
 	b.Handle(tb.OnAnimation, d.ApplyShutdownMiddleware(d.handleAnimationDistortion))
 	b.Handle(tb.OnSticker, d.ApplyShutdownMiddleware(d.handleStickerDistortion))
diff --git a/app/queue/honest_priority_queue.go b/app/queue/honest_priority_queue.go
index 9d8884d..9237417 100644
--- a/app/queue/honest_priority_queue.go
+++ b/app/queue/honest_priority_queue.go
@@ -11,10 +11,11 @@ import (
 // Wraps PriorityQueue to make it thread-safe. Manages priorities.
 // Extremely inefficient, but works for my use-case (very slow jobs and small queue sizes)
 type HonestJobQueue struct {
-	mu     *sync.RWMutex
-	queue  PriorityQueue
-	users  map[int64]int // Tracks the amount of job per-user currently in the queue. Used to calculate priority
-	banned map[int64]any // Drop jobs from these users
+	mu          *sync.RWMutex
+	queue       PriorityQueue
+	users       map[int64]int // Tracks the amount of job per-user currently in the queue. Used to calculate priority
+	banned      map[int64]any // Drop jobs from these users
+	maintenance bool
 }
 
 func NewHonestJobQueue(initialCapacity int) *HonestJobQueue {
@@ -96,15 +97,32 @@ func (hjq *HonestJobQueue) Pop() *Job {
 	return job
 }
 
+func (hjq *HonestJobQueue) ToggleMaintenance() bool {
+	hjq.mu.Lock()
+	defer hjq.mu.Unlock()
+
+	hjq.maintenance = !hjq.maintenance
+
+	return hjq.maintenance
+}
+
 func (hjq *HonestJobQueue) Push(userID int64, runnable func()) error {
 	hjq.mu.Lock()
 	defer hjq.mu.Unlock()
 
+	if hjq.maintenance {
+		return errors.New("The server is on temporary maintenance, no new videos are being processed at the moment, try again later")
+	}
+
+	if hjq.queue.Len() > 2000 {
+		return errors.New("There are too many items queued already, try again later")
+	}
+
 	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")
+		return errors.New("You're distorting videos too often, wait until the previous ones have been processed")
 	}
 
 	hjq.users[userID] = priority + 1
diff --git a/app/tools/video_worker.go b/app/tools/video_worker.go
index accc9ed..240dd26 100644
--- a/app/tools/video_worker.go
+++ b/app/tools/video_worker.go
@@ -53,3 +53,7 @@ func (vw *VideoWorker) QueueStats() (int, int) {
 func (vw *VideoWorker) IsBusy() bool {
 	return vw.queue.Len() > vw.workerCount
 }
+
+func (vw *VideoWorker) ToggleMaintenance() bool {
+	return vw.queue.ToggleMaintenance()
+}