Use the same logic in "bin/drydock lease" and LeaseUpdateWorker to identify candidate blueprints

Summary:
Ref T13676. Currently, "bin/drydock lease" just creates a lease that permits any blueprint.

To prepare for "use specific blueprint X", unify the logic between this workflow and LeaseUpdateWorker so we select only blueprints which at least have coarse compatibility (e.g., if we're leasing a host, only select enabled blueprints of classes that can allocate hosts).

Test Plan: Used `bin/drydock lease` to try to lease a nonsense type, got sensible error. Leased a host.

Subscribers: yelirekim, PHID-OPKG-gm6ozazyms6q6i22gyam

Maniphest Tasks: T13676

Differential Revision: https://secure.phabricator.com/D21801

src/applications/drydock/blueprint/DrydockBlueprintImplementation.php

@@ -329,6 +329,47 @@ abstract class DrydockBlueprintImplementation extends Phobject {
       ->execute();
   }
 
+
+  /**
+   * Get all the @{class:DrydockBlueprintImplementation}s which can possibly
+   * build a resource to satisfy a lease.
+   *
+   * This method returns blueprints which might, at some time, be able to
+   * build a resource which can satisfy the lease. They may not be able to
+   * build that resource right now.
+   *
+   * @param DrydockLease Requested lease.
+   * @return list<DrydockBlueprintImplementation> List of qualifying blueprint
+   *   implementations.
+   */
+  public static function getAllForAllocatingLease(
+    DrydockLease $lease) {
+
+    $impls = self::getAllBlueprintImplementations();
+
+    $keep = array();
+    foreach ($impls as $key => $impl) {
+      // Don't use disabled blueprint types.
+      if (!$impl->isEnabled()) {
+        continue;
+      }
+
+      // Don't use blueprint types which can't allocate the correct kind of
+      // resource.
+      if ($impl->getType() != $lease->getResourceType()) {
+        continue;
+      }
+
+      if (!$impl->canAnyBlueprintEverAllocateResourceForLease($lease)) {
+        continue;
+      }
+
+      $keep[$key] = $impl;
+    }
+
+    return $keep;
+  }
+
   public static function getNamedImplementation($class) {
     return idx(self::getAllBlueprintImplementations(), $class);
   }

src/applications/drydock/management/DrydockManagementLeaseWorkflow.php

@@ -79,6 +79,8 @@ final class DrydockManagementLeaseWorkflow
       $attributes = array();
     }
 
+    $blueprint_phids = null;
+
     $leases = array();
     for ($idx = 0; $idx < $count; $idx++) {
       $lease = id(new DrydockLease())
@@ -91,20 +93,11 @@ final class DrydockManagementLeaseWorkflow
         $lease->setAttributes($attributes);
       }
 
-      // TODO: This is not hugely scalable, although this is a debugging
-      // workflow so maybe it's fine. Do we even need `bin/drydock lease` in
-      // the long run?
-      $all_blueprints = id(new DrydockBlueprintQuery())
-        ->setViewer($viewer)
-        ->execute();
-      $allowed_phids = mpull($all_blueprints, 'getPHID');
-      if (!$allowed_phids) {
-        throw new Exception(
-          pht(
-            'No blueprints exist which can plausibly allocate resources to '.
-            'satisfy the requested lease.'));
+      if ($blueprint_phids === null) {
+        $blueprint_phids = $this->newAllowedBlueprintPHIDs($lease);
       }
-      $lease->setAllowedBlueprintPHIDs($allowed_phids);
+
+      $lease->setAllowedBlueprintPHIDs($blueprint_phids);
 
       if ($until) {
         $lease->setUntil($until);
@@ -260,4 +253,36 @@ final class DrydockManagementLeaseWorkflow
     }
   }
 
+  private function newAllowedBlueprintPHIDs(DrydockLease $lease) {
+    $viewer = $this->getViewer();
+
+    $impls = DrydockBlueprintImplementation::getAllForAllocatingLease($lease);
+
+    if (!$impls) {
+      throw new PhutilArgumentUsageException(
+        pht(
+          'No known blueprint class can ever allocate the specified '.
+          'lease. Check that the resource type is spelled correctly.'));
+    }
+
+    $classes = array_keys($impls);
+
+    $blueprints = id(new DrydockBlueprintQuery())
+      ->setViewer($viewer)
+      ->withBlueprintClasses($classes)
+      ->withDisabled(false)
+      ->execute();
+
+    if (!$blueprints) {
+      throw new PhutilArgumentUsageException(
+        pht(
+          'No enabled blueprints exist with a blueprint class that can '.
+          'plausibly allocate resources to satisfy the requested lease.'));
+    }
+
+    $phids = mpull($blueprints, 'getPHID');
+
+    return $phids;
+  }
+
 }

src/applications/drydock/worker/DrydockLeaseUpdateWorker.php

@@ -343,48 +343,6 @@ final class DrydockLeaseUpdateWorker extends DrydockWorker {
   }
 
 
-  /**
-   * Get all the @{class:DrydockBlueprintImplementation}s which can possibly
-   * build a resource to satisfy a lease.
-   *
-   * This method returns blueprints which might, at some time, be able to
-   * build a resource which can satisfy the lease. They may not be able to
-   * build that resource right now.
-   *
-   * @param DrydockLease Requested lease.
-   * @return list<DrydockBlueprintImplementation> List of qualifying blueprint
-   *   implementations.
-   * @task allocator
-   */
-  private function loadBlueprintImplementationsForAllocatingLease(
-    DrydockLease $lease) {
-
-    $impls = DrydockBlueprintImplementation::getAllBlueprintImplementations();
-
-    $keep = array();
-    foreach ($impls as $key => $impl) {
-      // Don't use disabled blueprint types.
-      if (!$impl->isEnabled()) {
-        continue;
-      }
-
-      // Don't use blueprint types which can't allocate the correct kind of
-      // resource.
-      if ($impl->getType() != $lease->getResourceType()) {
-        continue;
-      }
-
-      if (!$impl->canAnyBlueprintEverAllocateResourceForLease($lease)) {
-        continue;
-      }
-
-      $keep[$key] = $impl;
-    }
-
-    return $keep;
-  }
-
-
   /**
    * Get all the concrete @{class:DrydockBlueprint}s which can possibly
    * build a resource to satisfy a lease.
@@ -397,7 +355,7 @@ final class DrydockLeaseUpdateWorker extends DrydockWorker {
     DrydockLease $lease) {
     $viewer = $this->getViewer();
 
-    $impls = $this->loadBlueprintImplementationsForAllocatingLease($lease);
+    $impls = DrydockBlueprintImplementation::getAllForAllocatingLease($lease);
     if (!$impls) {
       return array();
     }