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

app/Contracts/Repository/NodeRepositoryInterface.php

@@ -54,15 +54,4 @@ interface NodeRepositoryInterface extends RepositoryInterface
      * @return \Illuminate\Support\Collection
      */
     public function getNodesForServerCreation(): Collection;
-
-    /**
-     * Return the IDs of all nodes that exist in the provided locations and have the space
-     * available to support the additional disk and memory provided.
-     *
-     * @param array $locations
-     * @param int $disk
-     * @param int $memory
-     * @return \Illuminate\Support\LazyCollection
-     */
-    public function getNodesWithResourceUse(array $locations, int $disk, int $memory): LazyCollection;
 }

app/Repositories/Eloquent/NodeRepository.php

@@ -171,28 +171,4 @@ class NodeRepository extends EloquentRepository implements NodeRepositoryInterfa
 
         return $instance->first();
     }
-
-    /**
-     * Return the IDs of all nodes that exist in the provided locations and have the space
-     * available to support the additional disk and memory provided.
-     *
-     * @param array $locations
-     * @param int $disk
-     * @param int $memory
-     * @return \Illuminate\Support\LazyCollection
-     */
-    public function getNodesWithResourceUse(array $locations, int $disk, int $memory): LazyCollection
-    {
-        $instance = $this->getBuilder()
-            ->select(['nodes.id', 'nodes.memory', 'nodes.disk', 'nodes.memory_overallocate', 'nodes.disk_overallocate'])
-            ->selectRaw('IFNULL(SUM(servers.memory), 0) as sum_memory, IFNULL(SUM(servers.disk), 0) as sum_disk')
-            ->leftJoin('servers', 'servers.node_id', '=', 'nodes.id')
-            ->where('nodes.public', 1);
-
-        if (! empty($locations)) {
-            $instance->whereIn('nodes.location_id', $locations);
-        }
-
-        return $instance->groupBy('nodes.id')->cursor();
-    }
 }

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->disk, 'Disk space must be an int, got %s');
-        Assert::integer($this->memory, 'Calls to ' . __METHOD__ . ' must have the memory usage set as an integer, received %s');
+        Assert::integer($this->memory, 'Memory usage must be an int, got %s');
 
-        $nodes = $this->repository->getNodesWithResourceUse($this->locations, $this->disk, $this->memory);
+        $query = Node::query()->select('nodes.*')
-        $viable = [];
+            ->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) {
+        if (! empty($this->locations)) {
-            $memoryLimit = $node->memory * (1 + ($node->memory_overallocate / 100));
+            $query = $query->whereIn('nodes.location_id', $this->locations);
-            $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($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();
     }

tests/Integration/Services/Deployment/FindViableNodesServiceTest.php

@@ -0,0 +1,115 @@
+<?php
+
+namespace Pterodactyl\Tests\Integration\Services\Deployment;
+
+use Pterodactyl\Models\Node;
+use InvalidArgumentException;
+use Pterodactyl\Models\Location;
+use Illuminate\Support\Collection;
+use Pterodactyl\Tests\Integration\IntegrationTestCase;
+use Pterodactyl\Services\Deployment\FindViableNodesService;
+use Pterodactyl\Exceptions\Service\Deployment\NoViableNodeException;
+
+class FindViableNodesServiceTest extends IntegrationTestCase
+{
+    public function testExceptionIsThrownIfNoDiskSpaceHasBeenSet()
+    {
+        $this->expectException(InvalidArgumentException::class);
+        $this->expectExceptionMessage('Disk space must be an int, got NULL');
+
+        $this->getService()->handle();
+    }
+
+    public function testExceptionIsThrownIfNoMemoryHasBeenSet()
+    {
+        $this->expectException(InvalidArgumentException::class);
+        $this->expectExceptionMessage('Memory usage must be an int, got NULL');
+
+        $this->getService()->setDisk(10)->handle();
+    }
+
+    public function testExpectedNodeIsReturnedForLocation()
+    {
+        /** @var \Pterodactyl\Models\Location[] $locations */
+        $locations = factory(Location::class)->times(2)->create();
+
+        /** @var \Pterodactyl\Models\Node[] $nodes */
+        $nodes = [
+            // This node should never be returned.
+            factory(Node::class)->create([
+                'location_id' => $locations[0]->id,
+                'memory' => 10240,
+                'disk' => 1024 * 100,
+            ]),
+            factory(Node::class)->create([
+                'location_id' => $locations[1]->id,
+                'memory' => 1024,
+                'disk' => 10240,
+                'disk_overallocate' => 10,
+            ]),
+            factory(Node::class)->create([
+                'location_id' => $locations[1]->id,
+                'memory' => 1024 * 4,
+                'memory_overallocate' => 50,
+                'disk' => 102400,
+            ]),
+        ];
+
+        $base = function () use ($locations) {
+            return $this->getService()->setLocations([ $locations[1]->id ])->setDisk(512);
+        };
+
+        $response = $base()->setMemory(512)->handle();
+        $this->assertInstanceOf(Collection::class, $response);
+        $this->assertFalse($response->isEmpty());
+        $this->assertSame(2, $response->count());
+        $this->assertSame(2, $response->where('location_id', $locations[1]->id)->count());
+
+        $response = $base()->setMemory(2048)->handle();
+        $this->assertSame(1, $response->count());
+        $this->assertSame($nodes[2]->id, $response[0]->id);
+
+        $response = $base()->setDisk(20480)->setMemory(256)->handle();
+        $this->assertSame(1, $response->count());
+        $this->assertSame($nodes[2]->id, $response[0]->id);
+
+        $response = $base()->setDisk(11263)->setMemory(256)->handle();
+        $this->assertSame(2, $response->count());
+
+        $servers = Collection::make([
+            $this->createServerModel(['node_id' => $nodes[1]->id, 'disk' => 5120]),
+            $this->createServerModel(['node_id' => $nodes[1]->id, 'disk' => 5120]),
+        ]);
+
+        $response = $base()->setDisk(1024)->setMemory(256)->handle();
+        $this->assertSame(1, $response->count());
+        $this->assertSame($nodes[2]->id, $response[0]->id);
+        $servers->each->delete();
+
+        $this->expectException(NoViableNodeException::class);
+        $base()->setMemory(10000)->handle();
+
+        Collection::make([
+            $this->createServerModel(['node_id' => $nodes[2]->id, 'memory' => 1024]),
+            $this->createServerModel(['node_id' => $nodes[2]->id, 'memory' => 1024]),
+            $this->createServerModel(['node_id' => $nodes[2]->id, 'memory' => 1024]),
+            $this->createServerModel(['node_id' => $nodes[2]->id, 'memory' => 1024]),
+        ]);
+
+        $response = $base()->setMemory(500)->handle();
+        $this->assertSame(2, $response->count());
+        $this->assertSame(2, $response->where('location_id', $locations[1]->id)->count());
+
+        $response = $base()->setMemory(512)->handle();
+        $this->assertSame(1, $response->count());
+        $this->assertSame($nodes[1]->id, $response[0]->id);
+    }
+
+    /**
+     * @return \Pterodactyl\Services\Deployment\FindViableNodesService
+     */
+    private function getService()
+    {
+        return $this->app->make(FindViableNodesService::class);
+    }
+}