Update node finding service logic to be single query; add test coverage

app/Services/Deployment/FindViableNodesService.php

@@ -3,16 +3,12 @@
 namespace Pterodactyl\Services\Deployment;
 
 use Webmozart\Assert\Assert;
-use Pterodactyl\Contracts\Repository\NodeRepositoryInterface;
+use Pterodactyl\Models\Node;
+use Illuminate\Support\LazyCollection;
 use Pterodactyl\Exceptions\Service\Deployment\NoViableNodeException;
 
 class FindViableNodesService
 {
-    /**
-     * @var \Pterodactyl\Contracts\Repository\NodeRepositoryInterface
-     */
-    private $repository;
-
     /**
      * @var array
      */
@@ -28,16 +24,6 @@ class FindViableNodesService
      */
     protected $memory;
 
-    /**
-     * FindViableNodesService constructor.
-     *
-     * @param \Pterodactyl\Contracts\Repository\NodeRepositoryInterface $repository
-     */
-    public function __construct(NodeRepositoryInterface $repository)
-    {
-        $this->repository = $repository;
-    }
-
     /**
      * Set the locations that should be searched through to locate available nodes.
      *
@@ -46,6 +32,8 @@ class FindViableNodesService
      */
     public function setLocations(array $locations): self
     {
+        Assert::allInteger($locations, 'An array of location IDs should be provided when calling setLocations.');
+
         $this->locations = $locations;
 
         return $this;
@@ -90,32 +78,34 @@ class FindViableNodesService
      * are tossed out, as are any nodes marked as non-public, meaning automatic
      * deployments should not be done against them.
      *
-     * @return int[]
+     * @return \Pterodactyl\Models\Node[]|\Illuminate\Support\Collection
      * @throws \Pterodactyl\Exceptions\Service\Deployment\NoViableNodeException
      */
-    public function handle(): array
+    public function handle()
     {
-        Assert::integer($this->disk, 'Calls to ' . __METHOD__ . ' must have the disk space set as an integer, received %s');
-        Assert::integer($this->memory, 'Calls to ' . __METHOD__ . ' must have the memory usage set as an integer, received %s');
+        Assert::integer($this->disk, 'Disk space must be an int, got %s');
+        Assert::integer($this->memory, 'Memory usage must be an int, got %s');
 
-        $nodes = $this->repository->getNodesWithResourceUse($this->locations, $this->disk, $this->memory);
-        $viable = [];
+        $query = Node::query()->select('nodes.*')
+            ->selectRaw('IFNULL(SUM(servers.memory), 0) as sum_memory')
+            ->selectRaw('IFNULL(SUM(servers.disk), 0) as sum_disk')
+            ->leftJoin('servers', 'servers.node_id', '=', 'nodes.id')
+            ->where('nodes.public', 1);
 
-        foreach ($nodes as $node) {
-            $memoryLimit = $node->memory * (1 + ($node->memory_overallocate / 100));
-            $diskLimit = $node->disk * (1 + ($node->disk_overallocate / 100));
-
-            if (($node->sum_memory + $this->memory) > $memoryLimit || ($node->sum_disk + $this->disk) > $diskLimit) {
-                continue;
-            }
-
-            $viable[] = $node->id;
+        if (! empty($this->locations)) {
+            $query = $query->whereIn('nodes.location_id', $this->locations);
         }
 
-        if (empty($viable)) {
+        $results = $query->groupBy('nodes.id')
+            ->havingRaw('(IFNULL(SUM(servers.memory), 0) + ?) < (nodes.memory * (1 + (nodes.memory_overallocate / 100)))', [ $this->memory ])
+            ->havingRaw('(IFNULL(SUM(servers.disk), 0) + ?) < (nodes.disk * (1 + (nodes.disk_overallocate / 100)))', [ $this->disk ])
+            ->get()
+            ->toBase();
+
+        if ($results->isEmpty()) {
             throw new NoViableNodeException(trans('exceptions.deployment.no_viable_nodes'));
         }
 
-        return $viable;
+        return $results;
     }
 }

app/Services/Servers/ServerCreationService.php

@@ -203,7 +203,7 @@ class ServerCreationService
             ->handle();
 
         return $this->allocationSelectionService->setDedicated($deployment->isDedicated())
-            ->setNodes($nodes)
+            ->setNodes($nodes->pluck('id')->toArray())
             ->setPorts($deployment->getPorts())
             ->handle();
     }