Improve Drydock's ability to allocate leases correctly

Summary:
Right now, Drydock gives out multiple leases to the same working copy and gives out leases to working copies with repository "P" in them when the user requested some other repository.

Add two callbacks:

  - `canAllocateLease()` - allows a blueprint to reject a lease on a resource because of a fundamental incompatibility, like "it's a working copy with Phabricator in it, but the lease wants a working copy with Javelin in it".
  - `shouldAllocateLease()` - allows a blueprint to reject a lease on a resource because of resource limits, like "only one active lease can own a working copy at a time".

Also cleaned up various other things.

Test Plan:
After implementing the callbacks, Drydock has the correct behavior:

  - It gives multiple leases on `localhost`, but only one lease per working-copy resource.
  - It does not grant leases on resources with repository X to requests for repository Y.

Ran `bin/drydock lease --type working-copy --repositoryID 12` and similar repeatedly and verified results in the web console.

Reviewers: btrahan

Reviewed By: btrahan

CC: aran

Maniphest Tasks: T2015

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

src/__phutil_library_map__.php

@@ -424,6 +424,7 @@ phutil_register_library_map(array(
     'DrydockAllocatorWorker' => 'applications/drydock/worker/DrydockAllocatorWorker.php',
     'DrydockApacheWebrootInterface' => 'applications/drydock/interface/webroot/DrydockApacheWebrootInterface.php',
     'DrydockBlueprint' => 'applications/drydock/blueprint/DrydockBlueprint.php',
+    'DrydockBlueprintScopeGuard' => 'applications/drydock/util/DrydockBlueprintScopeGuard.php',
     'DrydockCommandInterface' => 'applications/drydock/interface/command/DrydockCommandInterface.php',
     'DrydockConstants' => 'applications/drydock/constants/DrydockConstants.php',
     'DrydockController' => 'applications/drydock/controller/DrydockController.php',

src/applications/drydock/blueprint/DrydockBlueprint.php

@@ -1,9 +1,14 @@
 <?php
 
+/**
+ * @task lease      Lease Acquisition
+ * @task resource   Resource Allocation
+ * @task log        Logging
+ */
 abstract class DrydockBlueprint {
 
-  private $activeLease;
   private $activeResource;
+  private $activeLease;
 
   abstract public function getType();
   abstract public function getInterface(
@@ -17,48 +22,248 @@ abstract class DrydockBlueprint {
     return get_class($this);
   }
 
+
+/* -(  Lease Acquisition  )-------------------------------------------------- */
+
+
+  /**
+   * @task lease
+   */
+  final public function filterResource(
+    DrydockResource $resource,
+    DrydockLease $lease) {
+    return $this->canAllocateLease($resource, $lease);
+  }
+
+
+  /**
+   * Enforce basic checks on lease/resource compatibility. Allows resources to
+   * reject leases if they are incompatible, even if the resource types match.
+   *
+   * For example, if a resource represents a 32-bit host, this method might
+   * reject leases that need a 64-bit host. If a resource represents a working
+   * copy of repository "X", this method might reject leases which need a
+   * working copy of repository "Y". Generally, although the main types of
+   * a lease and resource may match (e.g., both "host"), it may not actually be
+   * possible to satisfy the lease with a specific resource.
+   *
+   * This method generally should not enforce limits or perform capacity
+   * checks. Perform those in @{method:shouldAllocateLease} instead. It also
+   * should not perform actual acquisition of the lease; perform that in
+   * @{method:executeAcquireLease} instead.
+   *
+   * @param   DrydockResource   Candidiate resource to allocate the lease on.
+   * @param   DrydockLease      Pending lease that wants to allocate here.
+   * @return  bool              True if the resource and lease are compatible.
+   * @task lease
+   */
+  abstract protected function canAllocateLease(
+    DrydockResource $resource,
+    DrydockLease $lease);
+
+
+  /**
+   * @task lease
+   */
+  final public function allocateLease(
+    DrydockResource $resource,
+    DrydockLease $lease) {
+
+    $scope = $this->pushActiveScope($resource, $lease);
+
+    $this->log('Trying to Allocate Lease');
+
+    $lease->setStatus(DrydockLeaseStatus::STATUS_ACQUIRING);
+    $lease->setResourceID($resource->getID());
+    $lease->attachResource($resource);
+
+    $ephemeral_lease = id(clone $lease)->makeEphemeral();
+
+    $allocated = false;
+    $allocation_exception = null;
+
+    $resource->openTransaction();
+      $resource->beginReadLocking();
+        $resource->reload();
+
+        $other_leases = id(new DrydockLease())->loadAllWhere(
+          'status IN (%Ld) AND resourceID = %d',
+          array(
+            DrydockLeaseStatus::STATUS_ACQUIRING,
+            DrydockLeaseStatus::STATUS_ACTIVE,
+          ),
+          $resource->getID());
+
+        try {
+          $allocated = $this->shouldAllocateLease(
+            $resource,
+            $ephemeral_lease,
+            $other_leases);
+        } catch (Exception $ex) {
+          $allocation_exception = $ex;
+        }
+
+        if ($allocated) {
+          $lease->save();
+        }
+      $resource->endReadLocking();
+    if ($allocated) {
+      $resource->saveTransaction();
+      $this->log('Allocated Lease');
+    } else {
+      $resource->killTransaction();
+      $this->log('Failed to Allocate Lease');
+    }
+
+    if ($allocation_exception) {
+      $this->logException($allocation_exception);
+    }
+
+    return $allocated;
+  }
+
+
+  /**
+   * Enforce lease limits on resources. Allows resources to reject leases if
+   * they would become over-allocated by accepting them.
+   *
+   * For example, if a resource represents disk space, this method might check
+   * how much space the lease is asking for (say, 200MB) and how much space is
+   * left unallocated on the resource. It could grant the lease (return true)
+   * if it has enough remaining space (more than 200MB), and reject the lease
+   * (return false) if it does not (less than 200MB).
+   *
+   * A resource might also allow only exclusive leases. In this case it could
+   * accept a new lease (return true) if there are no active leases, or reject
+   * the new lease (return false) if there any other leases.
+   *
+   * A lock is held on the resource while this method executes to prevent
+   * multiple processes from allocating leases on the resource simultaneously.
+   * However, this means you should implement the method as cheaply as possible.
+   * In particular, do not perform any actual acquisition or setup in this
+   * method.
+   *
+   * If allocation is permitted, the lease will be moved to `ACQUIRING` status
+   * and @{method:executeAcquireLease} will be called to actually perform
+   * acquisition.
+   *
+   * General compatibility checks unrelated to resource limits and capacity are
+   * better implemented in @{method:canAllocateLease}, which serves as a
+   * cheap filter before lock acquisition.
+   *
+   * @param   DrydockResource     Candidate resource to allocate the lease on.
+   * @param   DrydockLease        Pending lease that wants to allocate here.
+   * @param   list<DrydockLease>  Other allocated and acquired leases on the
+   *                              resource. The implementation can inspect them
+   *                              to verify it can safely add the new lease.
+   * @return  bool                True to allocate the lease on the resource;
+   *                              false to reject it.
+   * @task lease
+   */
+  abstract protected function shouldAllocateLease(
+    DrydockResource $resource,
+    DrydockLease $lease,
+    array $other_leases);
+
+
+  /**
+   * @task lease
+   */
+  final public function acquireLease(
+    DrydockResource $resource,
+    DrydockLease $lease) {
+
+    $scope = $this->pushActiveScope($resource, $lease);
+
+    $this->log('Acquiring Lease');
+    $lease->setStatus(DrydockLeaseStatus::STATUS_ACTIVE);
+    $lease->setResourceID($resource->getID());
+    $lease->attachResource($resource);
+
+    $ephemeral_lease = id(clone $lease)->makeEphemeral();
+
+    try {
+      $this->executeAcquireLease($resource, $ephemeral_lease);
+    } catch (Exception $ex) {
+      $this->logException($ex);
+      throw $ex;
+    }
+
+    $lease->setAttributes($ephemeral_lease->getAttributes());
+    $lease->save();
+    $this->log('Acquired Lease');
+  }
+
+
+  /**
+   * Acquire and activate an allocated lease. Allows resources to peform setup
+   * as leases are brought online.
+   *
+   * Following a successful call to @{method:canAllocateLease}, a lease is moved
+   * to `ACQUIRING` status and this method is called after resource locks are
+   * released. Nothing is locked while this method executes; the implementation
+   * is free to perform expensive operations like writing files and directories,
+   * executing commands, etc.
+   *
+   * After this method executes, the lease status is moved to `ACTIVE` and the
+   * original leasee may access it.
+   *
+   * If acquisition fails, throw an exception.
+   *
+   * @param   DrydockResource   Resource to acquire a lease on.
+   * @param   DrydockLease      Lease to acquire.
+   * @return  void
+   */
+  abstract protected function executeAcquireLease(
+    DrydockResource $resource,
+    DrydockLease $lease);
+
+
+/* -(  Resource Allocation  )------------------------------------------------ */
+
+
   public function canAllocateMoreResources(array $pool) {
     return true;
   }
 
   abstract protected function executeAllocateResource(DrydockLease $lease);
 
-  abstract protected function executeAcquireLease(
-    DrydockResource $resource,
-    DrydockLease $lease);
 
-  final public function acquireLease(
-    DrydockResource $resource,
-    DrydockLease $lease) {
+  final public function allocateResource(DrydockLease $lease) {
+    $scope = $this->pushActiveScope(null, $lease);
 
-    $this->activeResource   = $resource;
-    $this->activeLease      = $lease;
+    $this->log(
+      pht(
+        "Blueprint '%s': Allocating Resource for '%s'",
+        $this->getBlueprintClass(),
+        $lease->getLeaseName()));
 
-    $this->log('Acquiring Lease');
     try {
-      $lease->setStatus(DrydockLeaseStatus::STATUS_ACTIVE);
-      $lease->setResourceID($resource->getID());
-      $lease->attachResource($resource);
-
-      $this->executeAcquireLease($resource, $lease);
+      $resource = $this->executeAllocateResource($lease);
+      $this->validateAllocatedResource($resource);
     } catch (Exception $ex) {
       $this->logException($ex);
-      $this->activeResource   = null;
-      $this->activeLease      = null;
-
       throw $ex;
     }
 
-    $lease->save();
-
-    $this->activeResource   = null;
-    $this->activeLease      = null;
+    return $resource;
   }
 
+
+/* -(  Logging  )------------------------------------------------------------ */
+
+
+  /**
+   * @task log
+   */
   protected function logException(Exception $ex) {
     $this->log($ex->getMessage());
   }
 
+
+  /**
+   * @task log
+   */
   protected function log($message) {
     self::writeLog(
       $this->activeResource,
@@ -66,6 +271,10 @@ abstract class DrydockBlueprint {
       $message);
   }
 
+
+  /**
+   * @task log
+   */
   public static function writeLog(
     DrydockResource $resource = null,
     DrydockLease $lease = null,
@@ -86,28 +295,6 @@ abstract class DrydockBlueprint {
     $log->save();
   }
 
-  final public function allocateResource(DrydockLease $lease) {
-    $this->activeLease = $lease;
-    $this->activeResource = null;
-
-    $this->log(
-      pht(
-        "Blueprint '%s': Allocating Resource for '%s'",
-        $this->getBlueprintClass(),
-        $lease->getLeaseName()));
-
-    try {
-      $resource = $this->executeAllocateResource($lease);
-      $this->validateAllocatedResource($resource);
-    } catch (Exception $ex) {
-      $this->logException($ex);
-      $this->activeResource = null;
-
-      throw $ex;
-    }
-
-    return $resource;
-  }
 
   public static function getAllBlueprints() {
     static $list = null;
@@ -144,6 +331,7 @@ abstract class DrydockBlueprint {
     $resource->save();
 
     $this->activeResource = $resource;
+
     $this->log(
       pht(
         "Blueprint '%s': Created New Template",
@@ -178,4 +366,24 @@ abstract class DrydockBlueprint {
     }
   }
 
+  private function pushActiveScope(
+    DrydockResource $resource = null,
+    DrydockLease $lease = null) {
+
+    if (($this->activeResource !== null) ||
+        ($this->activeLease !== null)) {
+      throw new Exception("There is already an active resource or lease!");
+    }
+
+    $this->activeResource = $resource;
+    $this->activeLease = $lease;
+
+    return new DrydockBlueprintScopeGuard($this);
+  }
+
+  public function popActiveScope() {
+    $this->activeResource = null;
+    $this->activeLease = null;
+  }
+
 }

src/applications/drydock/blueprint/DrydockLocalHostBlueprint.php

@@ -36,6 +36,19 @@ final class DrydockLocalHostBlueprint extends DrydockBlueprint {
     return $resource;
   }
 
+  protected function canAllocateLease(
+    DrydockResource $resource,
+    DrydockLease $lease) {
+    return true;
+  }
+
+  protected function shouldAllocateLease(
+    DrydockResource $resource,
+    DrydockLease $lease,
+    array $other_leases) {
+    return true;
+  }
+
   protected function executeAcquireLease(
     DrydockResource $resource,
     DrydockLease $lease) {
@@ -45,10 +58,7 @@ final class DrydockLocalHostBlueprint extends DrydockBlueprint {
     $cmd = $lease->getInterface('command');
     $cmd->execx('mkdir %s', $lease_id);
 
-    $lease->setAttribute('path', $resource->getAttribute('path').'/'.$lease_id);
-    $lease->save();
-
-    return;
+    $lease->setAttribute('path', $resource->getAttribute('path').$lease_id.'/');
   }
 
   public function getType() {

src/applications/drydock/blueprint/DrydockWorkingCopyBlueprint.php

@@ -3,7 +3,25 @@
 final class DrydockWorkingCopyBlueprint extends DrydockBlueprint {
 
   public function isEnabled() {
-    return PhabricatorEnv::getEnvConfig('drydock.localhost.enabled');
+    return true;
+  }
+
+  protected function canAllocateLease(
+    DrydockResource $resource,
+    DrydockLease $lease) {
+
+    $resource_repo = $resource->getAttribute('repositoryID');
+    $lease_repo = $lease->getAttribute('repositoryID');
+
+    return ($resource_repo && $lease_repo && $resource_repo == $lease_repo);
+  }
+
+  protected function shouldAllocateLease(
+    DrydockResource $resource,
+    DrydockLease $lease,
+    array $other_leases) {
+
+    return !$other_leases;
   }
 
   protected function executeAllocateResource(DrydockLease $lease) {
@@ -31,7 +49,10 @@ final class DrydockWorkingCopyBlueprint extends DrydockBlueprint {
       ->setResourceType('host')
       ->waitUntilActive();
 
-    $path = $host_lease->getAttribute('path').'/'.$repository->getCallsign();
+    $path = $host_lease->getAttribute('path').$repository->getCallsign();
+
+    $this->log(
+      pht('Cloning %s into %s....', $repository->getCallsign(), $path));
 
     $cmd = $host_lease->getInterface('command');
     $cmd->execx(
@@ -39,6 +60,8 @@ final class DrydockWorkingCopyBlueprint extends DrydockBlueprint {
       $repository->getRemoteURI(),
       $path);
 
+    $this->log(pht('Complete.'));
+
     $resource = $this->newResourceTemplate($repository->getCallsign());
     $resource->setStatus(DrydockResourceStatus::STATUS_OPEN);
     $resource->setAttribute('lease.host', $host_lease->getID());

src/applications/drydock/constants/DrydockLeaseStatus.php

@@ -3,6 +3,7 @@
 final class DrydockLeaseStatus extends DrydockConstants {
 
   const STATUS_PENDING      = 0;
+  const STATUS_ACQUIRING    = 5;
   const STATUS_ACTIVE       = 1;
   const STATUS_RELEASED     = 2;
   const STATUS_BROKEN       = 3;
@@ -10,11 +11,12 @@ final class DrydockLeaseStatus extends DrydockConstants {
 
   public static function getNameForStatus($status) {
     static $map = array(
-      self::STATUS_PENDING  => 'Pending',
-      self::STATUS_ACTIVE   => 'Active',
-      self::STATUS_RELEASED => 'Released',
-      self::STATUS_BROKEN   => 'Broken',
-      self::STATUS_EXPIRED  => 'Expired',
+      self::STATUS_PENDING    => 'Pending',
+      self::STATUS_ACQUIRING  => 'Acquiring',
+      self::STATUS_ACTIVE     => 'Active',
+      self::STATUS_RELEASED   => 'Released',
+      self::STATUS_BROKEN     => 'Broken',
+      self::STATUS_EXPIRED    => 'Expired',
     );
 
     return idx($map, $status, 'Unknown');

src/applications/drydock/storage/DrydockLease.php

@@ -104,7 +104,8 @@ final class DrydockLease extends DrydockDAO {
   }
 
   private function assertActive() {
-    if ($this->status != DrydockLeaseStatus::STATUS_ACTIVE) {
+    if (($this->status != DrydockLeaseStatus::STATUS_ACTIVE) &&
+        ($this->status != DrydockLeaseStatus::STATUS_ACQUIRING)) {
       throw new Exception(
         "Lease is not active! You can not interact with resources through ".
         "an inactive lease.");
@@ -133,6 +134,7 @@ final class DrydockLease extends DrydockDAO {
           case DrydockLeaseStatus::STATUS_BROKEN:
             throw new Exception("Lease has been broken!");
           case DrydockLeaseStatus::STATUS_PENDING:
+          case DrydockLeaseStatus::STATUS_ACQUIRING:
             break;
         }
       }

src/applications/drydock/util/DrydockBlueprintScopeGuard.php

@@ -0,0 +1,13 @@
+<?php
+
+final class DrydockBlueprintScopeGuard {
+
+  public function __construct(DrydockBlueprint $blueprint) {
+    $this->blueprint = $blueprint;
+  }
+
+  public function __destruct() {
+    $this->blueprint->popActiveScope();
+  }
+
+}

src/applications/drydock/worker/DrydockAllocatorWorker.php

@@ -66,30 +66,39 @@ final class DrydockAllocatorWorker extends PhabricatorWorker {
     $candidates = array();
     foreach ($pool as $key => $candidate) {
       try {
-        $candidate->getBlueprint();
+        $blueprint = $candidate->getBlueprint();
       } catch (Exception $ex) {
         unset($pool[$key]);
+        continue;
       }
 
-      // TODO: Filter candidates according to ability to satisfy the lease.
-
-      $candidates[] = $candidate;
+      if ($blueprint->filterResource($candidate, $lease)) {
+        $candidates[] = $candidate;
+      }
     }
 
-    $this->log(
-      pht('%d Open Resource(s) Remain', count($candidates)));
+    $this->log(pht('%d Open Resource(s) Remain', count($candidates)));
 
+    $resource = null;
     if ($candidates) {
       shuffle($candidates);
-      $resource = head($candidates);
-    } else {
+      foreach ($candidates as $candidate_resource) {
+        $blueprint = $candidate_resource->getBlueprint();
+        if ($blueprint->allocateLease($candidate_resource, $lease)) {
+          $resource = $candidate_resource;
+          break;
+        }
+      }
+    }
+
+    if (!$resource) {
       $blueprints = DrydockBlueprint::getAllBlueprintsForResource($type);
 
       $this->log(
         pht('Found %d Blueprints', count($blueprints)));
 
-      foreach ($blueprints as $key => $blueprint) {
-        if (!$blueprint->isEnabled()) {
+      foreach ($blueprints as $key => $candidate_blueprint) {
+        if (!$candidate_blueprint->isEnabled()) {
           unset($blueprints[$key]);
           continue;
         }
@@ -98,8 +107,8 @@ final class DrydockAllocatorWorker extends PhabricatorWorker {
       $this->log(
         pht('%d Blueprints Enabled', count($blueprints)));
 
-      foreach ($blueprints as $key => $blueprint) {
-        if (!$blueprint->canAllocateMoreResources($pool)) {
+      foreach ($blueprints as $key => $candidate_blueprint) {
+        if (!$candidate_blueprint->canAllocateMoreResources($pool)) {
           unset($blueprints[$key]);
           continue;
         }
@@ -124,6 +133,24 @@ final class DrydockAllocatorWorker extends PhabricatorWorker {
 
       $blueprint = head($blueprints);
       $resource = $blueprint->allocateResource($lease);
+
+      if (!$blueprint->allocateLease($resource, $lease)) {
+        // TODO: This "should" happen only if we lost a race with another lease,
+        // which happened to acquire this resource immediately after we
+        // allocated it. In this case, the right behavior is to retry
+        // immediately. However, other things like a blueprint allocating a
+        // resource it can't actually allocate the lease on might be happening
+        // too, in which case we'd just allocate infinite resources. Probably
+        // what we should do is test for an active or allocated lease and retry
+        // if we find one (although it might have already been released by now)
+        // and fail really hard ("your configuration is a huge broken mess")
+        // otherwise. But just throw for now since this stuff is all edge-casey.
+        // Alternatively we could bring resources up in a "BESPOKE" status
+        // and then switch them to "OPEN" only after the allocating lease gets
+        // its grubby mitts on the resource. This might make more sense but
+        // is a bit messy.
+        throw new Exception("Lost an allocation race?");
+      }
     }
 
     $blueprint = $resource->getBlueprint();