From 3f0ceab06e5320f62f593c49daa30d963dbc36f9 Mon Sep 17 00:00:00 2001
From: Chocobozzz <me@florianbigard.com>
Date: Mon, 22 May 2023 15:52:59 +0200
Subject: [PATCH] More robust quota check

Avoid concurrency issues with permanent lives
---
 server/lib/live/live-manager.ts           |  4 ++--
 server/lib/live/live-quota-store.ts       | 14 +++++++-------
 server/lib/live/shared/muxing-session.ts  |  2 +-
 server/tests/api/live/live-constraints.ts | 12 +++++++++---
 4 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/server/lib/live/live-manager.ts b/server/lib/live/live-manager.ts
index f062e2fd3..5c6e69806 100644
--- a/server/lib/live/live-manager.ts
+++ b/server/lib/live/live-manager.ts
@@ -313,7 +313,7 @@ class LiveManager {
     const liveSession = await this.saveStartingSession(videoLive)
 
     const user = await UserModel.loadByLiveId(videoLive.id)
-    LiveQuotaStore.Instance.addNewLive(user.id, videoLive.id)
+    LiveQuotaStore.Instance.addNewLive(user.id, sessionId)
 
     const muxingSession = new MuxingSession({
       context: this.getContext(),
@@ -359,7 +359,7 @@ class LiveManager {
     muxingSession.on('after-cleanup', ({ videoUUID }) => {
       this.muxingSessions.delete(sessionId)
 
-      LiveQuotaStore.Instance.removeLive(user.id, videoLive.id)
+      LiveQuotaStore.Instance.removeLive(user.id, sessionId)
 
       muxingSession.destroy()
 
diff --git a/server/lib/live/live-quota-store.ts b/server/lib/live/live-quota-store.ts
index 8ceccde98..44539faaa 100644
--- a/server/lib/live/live-quota-store.ts
+++ b/server/lib/live/live-quota-store.ts
@@ -2,31 +2,31 @@ class LiveQuotaStore {
 
   private static instance: LiveQuotaStore
 
-  private readonly livesPerUser = new Map<number, { liveId: number, size: number }[]>()
+  private readonly livesPerUser = new Map<number, { sessionId: string, size: number }[]>()
 
   private constructor () {
   }
 
-  addNewLive (userId: number, liveId: number) {
+  addNewLive (userId: number, sessionId: string) {
     if (!this.livesPerUser.has(userId)) {
       this.livesPerUser.set(userId, [])
     }
 
-    const currentUserLive = { liveId, size: 0 }
+    const currentUserLive = { sessionId, size: 0 }
     const livesOfUser = this.livesPerUser.get(userId)
     livesOfUser.push(currentUserLive)
   }
 
-  removeLive (userId: number, liveId: number) {
+  removeLive (userId: number, sessionId: string) {
     const newLivesPerUser = this.livesPerUser.get(userId)
-                                             .filter(o => o.liveId !== liveId)
+                                             .filter(o => o.sessionId !== sessionId)
 
     this.livesPerUser.set(userId, newLivesPerUser)
   }
 
-  addQuotaTo (userId: number, liveId: number, size: number) {
+  addQuotaTo (userId: number, sessionId: string, size: number) {
     const lives = this.livesPerUser.get(userId)
-    const live = lives.find(l => l.liveId === liveId)
+    const live = lives.find(l => l.sessionId === sessionId)
 
     live.size += size
   }
diff --git a/server/lib/live/shared/muxing-session.ts b/server/lib/live/shared/muxing-session.ts
index c672ec4d6..434b308e6 100644
--- a/server/lib/live/shared/muxing-session.ts
+++ b/server/lib/live/shared/muxing-session.ts
@@ -271,7 +271,7 @@ class MuxingSession extends EventEmitter {
     try {
       const segmentStat = await stat(segmentPath)
 
-      LiveQuotaStore.Instance.addQuotaTo(this.user.id, this.videoLive.id, segmentStat.size)
+      LiveQuotaStore.Instance.addQuotaTo(this.user.id, this.sessionId, segmentStat.size)
 
       const canUpload = await this.isAbleToUploadVideoWithCache(this.user.id)
 
diff --git a/server/tests/api/live/live-constraints.ts b/server/tests/api/live/live-constraints.ts
index 168241b79..7af6af193 100644
--- a/server/tests/api/live/live-constraints.ts
+++ b/server/tests/api/live/live-constraints.ts
@@ -2,7 +2,7 @@
 
 import { expect } from 'chai'
 import { wait } from '@shared/core-utils'
-import { LiveVideoError, VideoPrivacy } from '@shared/models'
+import { LiveVideoError, UserVideoQuota, VideoPrivacy } from '@shared/models'
 import {
   cleanupTests,
   ConfigCommand,
@@ -172,12 +172,18 @@ describe('Test live constraints', function () {
     const ffmpegCommand = await servers[0].live.sendRTMPStreamInVideo({ token: userAccessToken, videoId: userVideoLiveoId })
 
     await servers[0].live.waitUntilPublished({ videoId: userVideoLiveoId })
+    // Wait previous live cleanups
+    await wait(3000)
 
     const baseQuota = await servers[0].users.getMyQuotaUsed({ token: userAccessToken })
 
-    await wait(3000)
+    let quotaUser: UserVideoQuota
+
+    do {
+      await wait(500)
 
-    const quotaUser = await servers[0].users.getMyQuotaUsed({ token: userAccessToken })
+      quotaUser = await servers[0].users.getMyQuotaUsed({ token: userAccessToken })
+    } while (quotaUser.videoQuotaUsed < baseQuota.videoQuotaUsed)
 
     const { data } = await servers[0].users.list()
     const quotaAdmin = data.find(u => u.username === 'user1')
-- 
2.41.0