Skip to content

Commit

Permalink
add maintenance mode and queue upper limit
Browse files Browse the repository at this point in the history
  • Loading branch information
graynk committed Jul 21, 2024
1 parent 637cf17 commit 8b820f6
Show file tree
Hide file tree
Showing 4 changed files with 38 additions and 6 deletions.
2 changes: 1 addition & 1 deletion Dockerfile
Original file line number Diff line number Diff line change
@@ -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 ./...
Expand Down
10 changes: 10 additions & 0 deletions app/distortioner.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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))
Expand Down
28 changes: 23 additions & 5 deletions app/queue/honest_priority_queue.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions app/tools/video_worker.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}

0 comments on commit 8b820f6

Please sign in to comment.