Add a flag to ./bin/worker to select tasks based on their failureCount

Summary:
I frequently run into a situation where I want to kill tasks that have accumulated a lot of failures regardless of what class they are. Or I'll want to kill every worker of a certain class but only if it has failed at least once. This change allows me to run `./bin/worker cancel --class <MYCLASS> --min-failure-count 5` to only kill tasks with at least 5 failed attempts.

The `--min-failure-count N` argument can be used by itself as well as with `--class CLASSNAME`. I don't think it makes sense for it to work with `--id ID`, but I'm not dead set on that or anything.

Test Plan: I ran the worker management workflow with and without the `--min-failure-count` argument and it worked as expected.

Reviewers: #blessed_reviewers, epriestley

Reviewed By: #blessed_reviewers, epriestley

Subscribers: Korvin, epriestley, yelirekim

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

src/__phutil_library_map__.php

@@ -3988,6 +3988,7 @@ phutil_register_library_map(array(
     'PhabricatorWordPressAuthProvider' => 'applications/auth/provider/PhabricatorWordPressAuthProvider.php',
     'PhabricatorWorker' => 'infrastructure/daemon/workers/PhabricatorWorker.php',
     'PhabricatorWorkerActiveTask' => 'infrastructure/daemon/workers/storage/PhabricatorWorkerActiveTask.php',
+    'PhabricatorWorkerActiveTaskQuery' => 'infrastructure/daemon/workers/query/PhabricatorWorkerActiveTaskQuery.php',
     'PhabricatorWorkerArchiveTask' => 'infrastructure/daemon/workers/storage/PhabricatorWorkerArchiveTask.php',
     'PhabricatorWorkerArchiveTaskQuery' => 'infrastructure/daemon/workers/query/PhabricatorWorkerArchiveTaskQuery.php',
     'PhabricatorWorkerBulkJob' => 'infrastructure/daemon/workers/storage/PhabricatorWorkerBulkJob.php',
@@ -4017,6 +4018,7 @@ phutil_register_library_map(array(
     'PhabricatorWorkerTask' => 'infrastructure/daemon/workers/storage/PhabricatorWorkerTask.php',
     'PhabricatorWorkerTaskData' => 'infrastructure/daemon/workers/storage/PhabricatorWorkerTaskData.php',
     'PhabricatorWorkerTaskDetailController' => 'applications/daemon/controller/PhabricatorWorkerTaskDetailController.php',
+    'PhabricatorWorkerTaskQuery' => 'infrastructure/daemon/workers/query/PhabricatorWorkerTaskQuery.php',
     'PhabricatorWorkerTestCase' => 'infrastructure/daemon/workers/__tests__/PhabricatorWorkerTestCase.php',
     'PhabricatorWorkerTrigger' => 'infrastructure/daemon/workers/storage/PhabricatorWorkerTrigger.php',
     'PhabricatorWorkerTriggerEvent' => 'infrastructure/daemon/workers/storage/PhabricatorWorkerTriggerEvent.php',
@@ -9200,8 +9202,9 @@ phutil_register_library_map(array(
     'PhabricatorWordPressAuthProvider' => 'PhabricatorOAuth2AuthProvider',
     'PhabricatorWorker' => 'Phobject',
     'PhabricatorWorkerActiveTask' => 'PhabricatorWorkerTask',
+    'PhabricatorWorkerActiveTaskQuery' => 'PhabricatorWorkerTaskQuery',
     'PhabricatorWorkerArchiveTask' => 'PhabricatorWorkerTask',
-    'PhabricatorWorkerArchiveTaskQuery' => 'PhabricatorQuery',
+    'PhabricatorWorkerArchiveTaskQuery' => 'PhabricatorWorkerTaskQuery',
     'PhabricatorWorkerBulkJob' => array(
       'PhabricatorWorkerDAO',
       'PhabricatorPolicyInterface',
@@ -9235,6 +9238,7 @@ phutil_register_library_map(array(
     'PhabricatorWorkerTask' => 'PhabricatorWorkerDAO',
     'PhabricatorWorkerTaskData' => 'PhabricatorWorkerDAO',
     'PhabricatorWorkerTaskDetailController' => 'PhabricatorDaemonController',
+    'PhabricatorWorkerTaskQuery' => 'PhabricatorQuery',
     'PhabricatorWorkerTestCase' => 'PhabricatorTestCase',
     'PhabricatorWorkerTrigger' => array(
       'PhabricatorWorkerDAO',

src/infrastructure/daemon/workers/management/PhabricatorWorkerManagementWorkflow.php

@@ -16,37 +16,47 @@ abstract class PhabricatorWorkerManagementWorkflow
         'param' => 'name',
         'help' => pht('Select all tasks of a given class.'),
       ),
+      array(
+        'name' => 'min-failure-count',
+        'param' => 'int',
+        'help' => pht('Limit to tasks with at least this many failures.'),
+      ),
     );
   }
 
   protected function loadTasks(PhutilArgumentParser $args) {
     $ids = $args->getArg('id');
     $class = $args->getArg('class');
+    $min_failures = $args->getArg('min-failure-count');
 
-    if (!$ids && !$class) {
+    if (!$ids && !$class && !$min_failures) {
       throw new PhutilArgumentUsageException(
-        pht('Use --id or --class to select tasks.'));
-    } if ($ids && $class) {
-      throw new PhutilArgumentUsageException(
-        pht('Use one of --id or --class to select tasks, but not both.'));
+        pht('Use --id, --class, or --min-failure-count to select tasks.'));
     }
 
+    $active_query = new PhabricatorWorkerActiveTaskQuery();
+    $archive_query = new PhabricatorWorkerArchiveTaskQuery();
+
     if ($ids) {
-      $active_tasks = id(new PhabricatorWorkerActiveTask())->loadAllWhere(
-        'id IN (%Ls)',
-        $ids);
-      $archive_tasks = id(new PhabricatorWorkerArchiveTaskQuery())
-        ->withIDs($ids)
-        ->execute();
-    } else {
-      $active_tasks = id(new PhabricatorWorkerActiveTask())->loadAllWhere(
-        'taskClass IN (%Ls)',
-        array($class));
-      $archive_tasks = id(new PhabricatorWorkerArchiveTaskQuery())
-        ->withClassNames(array($class))
-        ->execute();
+      $active_query = $active_query->withIDs($ids);
+      $archive_query = $archive_query->withIDs($ids);
     }
 
+    if ($class) {
+      $class_array = array($class);
+      $active_query = $active_query->withClassNames($class_array);
+      $archive_query = $archive_query->withClassNames($class_array);
+    }
+
+    if ($min_failures) {
+      $active_query = $active_query->withFailureCountBetween(
+        $min_failures, null);
+      $archive_query = $archive_query->withFailureCountBetween(
+        $min_failures, null);
+    }
+
+    $active_tasks = $active_query->execute();
+    $archive_tasks = $archive_query->execute();
     $tasks =
       mpull($active_tasks, null, 'getID') +
       mpull($archive_tasks, null, 'getID');
@@ -58,11 +68,24 @@ abstract class PhabricatorWorkerManagementWorkflow
             pht('No task exists with id "%s"!', $id));
         }
       }
-    } else {
+    }
+    if ($class && $min_failures) {
+      if (!$tasks) {
+        throw new PhutilArgumentUsageException(
+          pht('No task exists with class "%s" and at least %d failures!',
+            $class,
+            $min_failures));
+      }
+    } else if ($class) {
       if (!$tasks) {
         throw new PhutilArgumentUsageException(
           pht('No task exists with class "%s"!', $class));
       }
+    } else if ($min_failures) {
+      if (!$tasks) {
+        throw new PhutilArgumentUsageException(
+          pht('No tasks exist with at least %d failures!', $min_failures));
+      }
     }
 
     // When we lock tasks properly, this gets populated as a side effect. Just

src/infrastructure/daemon/workers/query/PhabricatorWorkerActiveTaskQuery.php

@@ -0,0 +1,21 @@
+<?php
+
+final class PhabricatorWorkerActiveTaskQuery
+  extends PhabricatorWorkerTaskQuery {
+
+  public function execute() {
+    $task_table = new PhabricatorWorkerActiveTask();
+
+    $conn_r = $task_table->establishConnection('r');
+
+    $rows = queryfx_all(
+      $conn_r,
+      'SELECT * FROM %T %Q %Q %Q',
+      $task_table->getTableName(),
+      $this->buildWhereClause($conn_r),
+      $this->buildOrderClause($conn_r),
+      $this->buildLimitClause($conn_r));
+
+    return $task_table->loadAllFromArray($rows);
+  }
+}

src/infrastructure/daemon/workers/query/PhabricatorWorkerArchiveTaskQuery.php

@@ -1,44 +1,7 @@
 <?php
 
 final class PhabricatorWorkerArchiveTaskQuery
-  extends PhabricatorQuery {
-
-  private $ids;
-  private $dateModifiedSince;
-  private $dateCreatedBefore;
-  private $objectPHIDs;
-  private $classNames;
-  private $limit;
-
-  public function withIDs(array $ids) {
-    $this->ids = $ids;
-    return $this;
-  }
-
-  public function withDateModifiedSince($timestamp) {
-    $this->dateModifiedSince = $timestamp;
-    return $this;
-  }
-
-  public function withDateCreatedBefore($timestamp) {
-    $this->dateCreatedBefore = $timestamp;
-    return $this;
-  }
-
-  public function withObjectPHIDs(array $phids) {
-    $this->objectPHIDs = $phids;
-    return $this;
-  }
-
-  public function withClassNames(array $names) {
-    $this->classNames = $names;
-    return $this;
-  }
-
-  public function setLimit($limit) {
-    $this->limit = $limit;
-    return $this;
-  }
+  extends PhabricatorWorkerTaskQuery {
 
   public function execute() {
     $task_table = new PhabricatorWorkerArchiveTask();
@@ -55,68 +18,4 @@ final class PhabricatorWorkerArchiveTaskQuery
 
     return $task_table->loadAllFromArray($rows);
   }
-
-  protected function buildWhereClause(AphrontDatabaseConnection $conn_r) {
-    $where = array();
-
-    if ($this->ids !== null) {
-      $where[] = qsprintf(
-        $conn_r,
-        'id in (%Ld)',
-        $this->ids);
-    }
-
-    if ($this->objectPHIDs !== null) {
-      $where[] = qsprintf(
-        $conn_r,
-        'objectPHID IN (%Ls)',
-        $this->objectPHIDs);
-    }
-
-    if ($this->dateModifiedSince !== null) {
-      $where[] = qsprintf(
-        $conn_r,
-        'dateModified > %d',
-        $this->dateModifiedSince);
-    }
-
-    if ($this->dateCreatedBefore !== null) {
-      $where[] = qsprintf(
-        $conn_r,
-        'dateCreated < %d',
-        $this->dateCreatedBefore);
-    }
-
-    if ($this->classNames !== null) {
-      $where[] = qsprintf(
-        $conn_r,
-        'taskClass IN (%Ls)',
-        $this->classNames);
-    }
-
-    return $this->formatWhereClause($where);
-  }
-
-  private function buildOrderClause(AphrontDatabaseConnection $conn_r) {
-    // NOTE: The garbage collector executes this query with a date constraint,
-    // and the query is inefficient if we don't use the same key for ordering.
-    // See T9808 for discussion.
-
-    if ($this->dateCreatedBefore) {
-      return qsprintf($conn_r, 'ORDER BY dateCreated DESC, id DESC');
-    } else if ($this->dateModifiedSince) {
-      return qsprintf($conn_r, 'ORDER BY dateModified DESC, id DESC');
-    } else {
-      return qsprintf($conn_r, 'ORDER BY id DESC');
-    }
-  }
-
-  private function buildLimitClause(AphrontDatabaseConnection $conn_r) {
-    $clause =  '';
-    if ($this->limit) {
-      $clause = qsprintf($conn_r, 'LIMIT %d', $this->limit);
-    }
-    return $clause;
-  }
-
 }

src/infrastructure/daemon/workers/query/PhabricatorWorkerTaskQuery.php

@@ -0,0 +1,128 @@
+<?php
+
+abstract class PhabricatorWorkerTaskQuery
+  extends PhabricatorQuery {
+
+  private $ids;
+  private $dateModifiedSince;
+  private $dateCreatedBefore;
+  private $objectPHIDs;
+  private $classNames;
+  private $limit;
+  private $minFailureCount;
+  private $maxFailureCount;
+
+  public function withIDs(array $ids) {
+    $this->ids = $ids;
+    return $this;
+  }
+
+  public function withDateModifiedSince($timestamp) {
+    $this->dateModifiedSince = $timestamp;
+    return $this;
+  }
+
+  public function withDateCreatedBefore($timestamp) {
+    $this->dateCreatedBefore = $timestamp;
+    return $this;
+  }
+
+  public function withObjectPHIDs(array $phids) {
+    $this->objectPHIDs = $phids;
+    return $this;
+  }
+
+  public function withClassNames(array $names) {
+    $this->classNames = $names;
+    return $this;
+  }
+
+  public function withFailureCountBetween($min, $max) {
+    $this->minFailureCount = $min;
+    $this->maxFailureCount = $max;
+    return $this;
+  }
+
+  public function setLimit($limit) {
+    $this->limit = $limit;
+    return $this;
+  }
+
+  protected function buildWhereClause(AphrontDatabaseConnection $conn_r) {
+    $where = array();
+
+    if ($this->ids !== null) {
+      $where[] = qsprintf(
+        $conn_r,
+        'id in (%Ld)',
+        $this->ids);
+    }
+
+    if ($this->objectPHIDs !== null) {
+      $where[] = qsprintf(
+        $conn_r,
+        'objectPHID IN (%Ls)',
+        $this->objectPHIDs);
+    }
+
+    if ($this->dateModifiedSince !== null) {
+      $where[] = qsprintf(
+        $conn_r,
+        'dateModified > %d',
+        $this->dateModifiedSince);
+    }
+
+    if ($this->dateCreatedBefore !== null) {
+      $where[] = qsprintf(
+        $conn_r,
+        'dateCreated < %d',
+        $this->dateCreatedBefore);
+    }
+
+    if ($this->classNames !== null) {
+      $where[] = qsprintf(
+        $conn_r,
+        'taskClass IN (%Ls)',
+        $this->classNames);
+    }
+
+    if ($this->minFailureCount !== null) {
+      $where[] = qsprintf(
+        $conn_r,
+        'failureCount >= %d',
+        $this->minFailureCount);
+    }
+
+    if ($this->maxFailureCount !== null) {
+      $where[] = qsprintf(
+        $conn_r,
+        'failureCount <= %d',
+        $this->maxFailureCount);
+    }
+
+    return $this->formatWhereClause($where);
+  }
+
+  protected function buildOrderClause(AphrontDatabaseConnection $conn_r) {
+    // NOTE: The garbage collector executes this query with a date constraint,
+    // and the query is inefficient if we don't use the same key for ordering.
+    // See T9808 for discussion.
+
+    if ($this->dateCreatedBefore) {
+      return qsprintf($conn_r, 'ORDER BY dateCreated DESC, id DESC');
+    } else if ($this->dateModifiedSince) {
+      return qsprintf($conn_r, 'ORDER BY dateModified DESC, id DESC');
+    } else {
+      return qsprintf($conn_r, 'ORDER BY id DESC');
+    }
+  }
+
+  protected function buildLimitClause(AphrontDatabaseConnection $conn_r) {
+    $clause =  '';
+    if ($this->limit) {
+      $clause = qsprintf($conn_r, 'LIMIT %d', $this->limit);
+    }
+    return $clause;
+  }
+
+}