Allow Drydock leases to expire after a time limit

Summary: Ref T6569. If a lease is activated with an expiration date, schedule a task to try to clean it up after that time.

Test Plan:
  - Used `bin/drydock lease ... --until ...` to activate a lease in the near future.
  - Waited for a bit.
  - Saw it expire and get destroyed at the scheduled time.

Reviewers: hach-que, chad

Reviewed By: chad

Maniphest Tasks: T6569

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

src/applications/drydock/controller/DrydockLeaseViewController.php

@@ -97,6 +97,7 @@ final class DrydockLeaseViewController extends DrydockLeaseController {
   private function buildPropertyListView(
     DrydockLease $lease,
     PhabricatorActionListView $actions) {
+    $viewer = $this->getViewer();
 
     $view = new PHUIPropertyListView();
     $view->setActionList($actions);
@@ -145,6 +146,14 @@ final class DrydockLeaseViewController extends DrydockLeaseController {
         pht('No Resource'));
     }
 
+    $until = $lease->getUntil();
+    if ($until) {
+      $until_display = phabricator_datetime($until, $viewer);
+    } else {
+      $until_display = phutil_tag('em', array(), pht('Never'));
+    }
+    $view->addProperty(pht('Expires'), $until_display);
+
     $attributes = $lease->getAttributes();
     if ($attributes) {
       $view->addSectionHeader(

src/applications/drydock/management/DrydockManagementLeaseWorkflow.php

@@ -14,6 +14,11 @@ final class DrydockManagementLeaseWorkflow
             'param'     => 'resource_type',
             'help'      => pht('Resource type.'),
           ),
+          array(
+            'name' => 'until',
+            'param' => 'time',
+            'help' => pht('Set lease expiration time.'),
+          ),
           array(
             'name'      => 'attributes',
             'param'     => 'name=value,...',
@@ -33,6 +38,17 @@ final class DrydockManagementLeaseWorkflow
           '--type'));
     }
 
+    $until = $args->getArg('until');
+    if (strlen($until)) {
+      $until = strtotime($until);
+      if ($until <= 0) {
+        throw new PhutilArgumentUsageException(
+          pht(
+            'Unable to parse argument to "%s".',
+            '--until'));
+      }
+    }
+
     $attributes = $args->getArg('attributes');
     if ($attributes) {
       $options = new PhutilSimpleOptions();
@@ -42,9 +58,15 @@ final class DrydockManagementLeaseWorkflow
 
     $lease = id(new DrydockLease())
       ->setResourceType($resource_type);
+
     if ($attributes) {
       $lease->setAttributes($attributes);
     }
+
+    if ($until) {
+      $lease->setUntil($until);
+    }
+
     $lease->queueForActivation();
 
     echo tsprintf(

src/applications/drydock/storage/DrydockLease.php

@@ -295,14 +295,16 @@ final class DrydockLease extends DrydockDAO
     }
   }
 
-  public function scheduleUpdate() {
+  public function scheduleUpdate($epoch = null) {
     PhabricatorWorker::scheduleTask(
       'DrydockLeaseUpdateWorker',
       array(
         'leasePHID' => $this->getPHID(),
+        'isExpireTask' => ($epoch !== null),
       ),
       array(
         'objectPHID' => $this->getPHID(),
+        'delayUntil' => $epoch,
       ));
   }
 
@@ -322,6 +324,11 @@ final class DrydockLease extends DrydockDAO
     if ($need_update) {
       $this->scheduleUpdate();
     }
+
+    $expires = $this->getUntil();
+    if ($expires) {
+      $this->scheduleUpdate($expires);
+    }
   }
 
 

src/applications/drydock/worker/DrydockLeaseUpdateWorker.php

@@ -11,13 +11,38 @@ final class DrydockLeaseUpdateWorker extends DrydockWorker {
     $lock = PhabricatorGlobalLock::newLock($lock_key)
       ->lock(1);
 
-    $lease = $this->loadLease($lease_phid);
-    $this->updateLease($lease);
+    try {
+      $lease = $this->loadLease($lease_phid);
+      $this->updateLease($lease);
+    } catch (Exception $ex) {
+      $lock->unlock();
+      throw $ex;
+    }
 
     $lock->unlock();
   }
 
   private function updateLease(DrydockLease $lease) {
+    if ($lease->getStatus() != DrydockLeaseStatus::STATUS_ACTIVE) {
+      return;
+    }
+
+    $viewer = $this->getViewer();
+    $drydock_phid = id(new PhabricatorDrydockApplication())->getPHID();
+
+    // Check if the lease has expired. If it is, we're going to send it a
+    // release command. This command will be handled immediately below, it
+    // just generates a command log and improves consistency.
+    $now = PhabricatorTime::getNow();
+    $expires = $lease->getUntil();
+    if ($expires && ($expires <= $now)) {
+      $command = DrydockCommand::initializeNewCommand($viewer)
+        ->setTargetPHID($lease->getPHID())
+        ->setAuthorPHID($drydock_phid)
+        ->setCommand(DrydockCommand::COMMAND_RELEASE)
+        ->save();
+    }
+
     $commands = $this->loadCommands($lease->getPHID());
     foreach ($commands as $command) {
       if ($lease->getStatus() != DrydockLeaseStatus::STATUS_ACTIVE) {
@@ -27,10 +52,21 @@ final class DrydockLeaseUpdateWorker extends DrydockWorker {
       }
 
       $this->processCommand($lease, $command);
+
       $command
         ->setIsConsumed(true)
         ->save();
     }
+
+    // If this is the task which will eventually release the lease after it
+    // expires but it is still active, reschedule the task to run after the
+    // lease expires. This can happen if the lease's expiration was pushed
+    // forward.
+    if ($lease->getStatus() == DrydockLeaseStatus::STATUS_ACTIVE) {
+      if ($this->getTaskDataValue('isExpireTask') && $expires) {
+        throw new PhabricatorWorkerYieldException($expires - $now);
+      }
+    }
   }
 
   private function processCommand(

src/infrastructure/daemon/workers/PhabricatorWorker.php

@@ -117,6 +117,11 @@ abstract class PhabricatorWorker extends Phobject {
       ->setPriority($priority)
       ->setObjectPHID($object_phid);
 
+    $delay = idx($options, 'delayUntil');
+    if ($delay) {
+      $task->setLeaseExpires($delay);
+    }
+
     if (self::$runAllTasksInProcess) {
       // Do the work in-process.
       $worker = newv($task_class, array($data));