refactor: penalty calculation for best node

hmes98318 · Sep 15, 2023 · 264f22b · 264f22b
1 parent b317579
commit 264f22b
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 30 deletions.
diff --git a/src/lib/LavaShark.ts b/src/lib/LavaShark.ts
@@ -172,7 +172,8 @@ export class LavaShark extends EventEmitter {
         try {
             await node.checkNodeSession();
         } catch (_) {
-            throw new Error('Session not found.');
+            this.lastNodeSorting = 0;
+            return this.bestNode();
         }
 
         return node;
@@ -385,24 +386,7 @@ export class LavaShark extends EventEmitter {
             throw new Error('Invalid node parameter. Expected Node instance.');
         }
 
-        try {
-            const startTime = Date.now();
-
-            await Promise.race([
-                node.getVersion(),
-                new Promise<void>((_, reject) => {
-                    setTimeout(() => {
-                        reject(new Error('Send ping data timed out.'));
-                    }, timeout);
-                })
-            ]);
-
-            const endTime = Date.now();
-            const ping = endTime - startTime;   
-            return ping;
-        } catch (_) {
-            return -1;
-        }
+        return await node.getPing(timeout);
     }
 
     /**
@@ -418,7 +402,7 @@ export class LavaShark extends EventEmitter {
         }
 
         const pingPromises = nodes.map(async (node) => {
-            return this.nodePing(node, timeout);
+            return node.getPing(timeout);
         });
 
         const pingList = await Promise.all(pingPromises);

diff --git a/src/lib/Node.ts b/src/lib/Node.ts
@@ -166,19 +166,29 @@ export default class Node {
     }
 
     private calcPenalties() {
-        // https://github.com/freyacodes/Lavalink-Client/blob/d71003475376d440baf081faa79577bd11221684/src/main/java/lavalink/client/io/LavalinkLoadBalancer.java#L131-L141
+        const totalWeight = 100; // Total weight is 100%
 
-        const cpuPenalty = Math.pow(1.05, 100 * this.stats.cpu.systemLoad) * 10 - 10;
+        // Weights for each penalty category
+        const weightCpuPenalty = 25;
+        const weightDeficitFramePenalty = 30;
+        const weightNullFramePenalty = 15;
+        const weightPlayingPlayers = 30;
 
+        const cpuPenalty = Math.pow(1.05, 100 * this.stats.cpu.systemLoad) * 10 - 10;
         let deficitFramePenalty = 0, nullFramePenalty = 0;
 
         if (this.stats.frameStats) {
-            deficitFramePenalty = Math.pow(1.03, 500 * this.stats.frameStats.deficit / 3000) * 600 - 600;
-            nullFramePenalty = Math.pow(1.03, 500 * this.stats.frameStats.nulled / 3000) * 300 - 300;
-            nullFramePenalty *= 2;
+            deficitFramePenalty = Math.pow(1.03, this.stats.frameStats.deficit * 2);
+            nullFramePenalty = Math.pow(1.03, this.stats.frameStats.nulled);
         }
 
-        this.penalties = ~~(cpuPenalty + deficitFramePenalty + nullFramePenalty + this.stats.playingPlayers);
+        // Calculate weighted values for each penalty category
+        const weightedPlayingPlayers = (weightPlayingPlayers / totalWeight) * this.stats.playingPlayers;
+        const weightedDeficitFramePenalty = (weightDeficitFramePenalty / totalWeight) * deficitFramePenalty;
+        const weightedNullFramePenalty = (weightNullFramePenalty / totalWeight) * nullFramePenalty;
+        const weightedCpuPenalty = (weightCpuPenalty / totalWeight) * cpuPenalty;
+
+        this.penalties = ~~((weightedPlayingPlayers + weightedDeficitFramePenalty + weightedNullFramePenalty + weightedCpuPenalty) * 100);
     }
 
     /**
@@ -290,6 +300,34 @@ export default class Node {
         return stats;
     }
 
+    /**
+     * Get the ping for the node
+     * @param {number} timeout - Timeout value in milliseconds
+     * @returns {Promise<number>} - Node latency, in milliseconds
+     */
+    public async getPing(timeout: number = 1500): Promise<number> {
+        try {
+            const startTime = Date.now();
+
+            await Promise.race([
+                this.getStats(),
+                new Promise<void>((_, reject) => {
+                    setTimeout(() => {
+                        reject(new Error('Update stats timed out'));
+                    }, timeout);
+                })
+            ]);
+
+            const endTime = Date.now();
+            const ping = endTime - startTime;
+            return ping;
+        } catch (_) {
+            this.disconnect();
+            this.lavashark.emit('error', this, new Error(`An error occurred while updating stats: Unable to connect to the node`));
+            return -1;
+        }
+    }
+
     /**
      * Gets the route planner status
      * @returns {Promise<Object>}
@@ -316,20 +354,20 @@ export default class Node {
     /**
      * Update node stats
      */
-    public async updateStats(): Promise<void> {
+    public async updateStats(timeout: number = 1500): Promise<void> {
         try {
             await Promise.race([
                 this.getStats(),
                 new Promise<void>((_, reject) => {
                     setTimeout(() => {
                         reject(new Error('Update stats timed out'));
-                    }, 1500);
+                    }, timeout);
                 })
             ]);
             this.calcPenalties();
         } catch (_) {
             this.disconnect();
-            this.lavashark.emit('error', this, new Error(`An error occurred while updating stats: Unable to connect to the node`));
+            this.lavashark.emit('error', this, new Error(`An error occurred while updating stats: Unable to connect to the node "${this.identifier}"`));
         }
     }
 
@@ -487,6 +525,7 @@ export default class Node {
                 delete payload.op;
                 this.stats = payload as NodeStats;
                 this.calcPenalties();
+                this.lavashark.emit('debug', `Node "${this.identifier}" penalties: ${this.penalties ?? 0}`);
                 break;
             }
             case 'playerUpdate': {

diff --git a/typings/src/lib/Node.d.ts b/typings/src/lib/Node.d.ts
@@ -74,6 +74,12 @@ export default class Node {
      * @returns {Promise<NodeStats>}
      */
     getStats(): Promise<NodeStats>;
+    /**
+     * Get the ping for the node
+     * @param {number} timeout - Timeout value in milliseconds
+     * @returns {Promise<number>} - Node latency, in milliseconds
+     */
+    getPing(timeout?: number): Promise<number>;
     /**
      * Gets the route planner status
      * @returns {Promise<Object>}
@@ -91,7 +97,7 @@ export default class Node {
     /**
      * Update node stats
      */
-    updateStats(): Promise<void>;
+    updateStats(timeout?: number): Promise<void>;
     private pollTrack;
     private handlePlayerEvent;
     private handleTrackStart;