Don't issue a bazillion queries to load Differential object lists

Summary:
Ref T3496. Currently, we call loadAssets() on each revision table, which invokes a new revision query and a pile of subqueries.

Instead, add `needFlags()` and `needDrafts()` to `RevisionQuery`. Some day these could perhaps be more generic.

Test Plan:
  - Viewed home, differential, etc., no longer saw 9203809238 queries being run for no reason.
  - Drafts and flags still appear properly.

Reviewers: btrahan

Reviewed By: btrahan

CC: aran

Maniphest Tasks: T3496

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

src/applications/differential/controller/DifferentialRevisionListController.php

@@ -46,8 +46,7 @@ final class DifferentialRevisionListController extends DifferentialController
           pht('No revisions are blocked on your action.'))
         ->setHighlightAge(true)
         ->setRevisions($blocking)
-        ->setHandles(array())
+        ->setHandles(array());
-        ->loadAssets();
 
       $views[] = id(clone $template)
         ->setHeader(pht('Action Required'))
@@ -55,21 +54,18 @@ final class DifferentialRevisionListController extends DifferentialController
           pht('No revisions require your action.'))
         ->setHighlightAge(true)
         ->setRevisions($active)
-        ->setHandles(array())
+        ->setHandles(array());
-        ->loadAssets();
 
       $views[] = id(clone $template)
         ->setHeader(pht('Waiting on Others'))
         ->setNoDataString(
           pht('You have no revisions waiting on others.'))
         ->setRevisions($waiting)
-        ->setHandles(array())
+        ->setHandles(array());
-        ->loadAssets();
     } else {
       $views[] = id(clone $template)
         ->setRevisions($revisions)
-        ->setHandles(array())
+        ->setHandles(array());
-        ->loadAssets();
     }
 
     $phids = array_mergev(mpull($views, 'getRequiredHandlePHIDs'));

src/applications/differential/controller/DifferentialRevisionViewController.php

@@ -810,6 +810,8 @@ final class DifferentialRevisionViewController extends DifferentialController {
       ->withStatus(DifferentialRevisionQuery::STATUS_OPEN)
       ->setOrder(DifferentialRevisionQuery::ORDER_PATH_MODIFIED)
       ->setLimit(10)
+      ->needFlags(true)
+      ->needDrafts(true)
       ->needRelationships(true);
 
     foreach ($path_map as $path => $path_id) {
@@ -836,8 +838,7 @@ final class DifferentialRevisionViewController extends DifferentialController {
     $view = id(new DifferentialRevisionListView())
       ->setRevisions($revisions)
       ->setFields(DifferentialRevisionListView::getDefaultFields($user))
-      ->setUser($user)
+      ->setUser($user);
-      ->loadAssets();
 
     $phids = $view->getRequiredHandlePHIDs();
     $handles = $this->loadViewerHandles($phids);

src/applications/differential/query/DifferentialRevisionQuery.php

@@ -57,6 +57,8 @@ final class DifferentialRevisionQuery
   private $needHashes         = false;
   private $needReviewerStatus = false;
   private $needReviewerAuthority;
+  private $needDrafts;
+  private $needFlags;
 
   private $buildingGlobalOrder;
 
@@ -345,6 +347,16 @@ final class DifferentialRevisionQuery
     return $this;
   }
 
+  public function needFlags($need_flags) {
+    $this->needFlags = $need_flags;
+    return $this;
+  }
+
+  public function needDrafts($need_drafts) {
+    $this->needDrafts = $need_drafts;
+    return $this;
+  }
+
 
 /* -(  Query Execution  )---------------------------------------------------- */
 
@@ -463,6 +475,39 @@ final class DifferentialRevisionQuery
     return $revisions;
   }
 
+  protected function didFilterPage(array $revisions) {
+    $viewer = $this->getViewer();
+
+    if ($this->needFlags) {
+      $flags = id(new PhabricatorFlagQuery())
+        ->setViewer($viewer)
+        ->withOwnerPHIDs(array($viewer->getPHID()))
+        ->withObjectPHIDs(mpull($revisions, 'getPHID'))
+        ->execute();
+      $flags = mpull($flags, null, 'getObjectPHID');
+      foreach ($revisions as $revision) {
+        $revision->attachFlag(
+          $viewer,
+          idx($flags, $revision->getPHID()));
+      }
+    }
+
+    if ($this->needDrafts) {
+      $drafts = id(new DifferentialDraft())->loadAllWhere(
+        'authorPHID = %s AND objectPHID IN (%Ls)',
+        $viewer->getPHID(),
+        mpull($revisions, 'getPHID'));
+      $drafts = mgroup($drafts, 'getObjectPHID');
+      foreach ($revisions as $revision) {
+        $revision->attachDrafts(
+          $viewer,
+          idx($drafts, $revision->getPHID(), array()));
+      }
+    }
+
+    return $revisions;
+  }
+
   private function loadData() {
     $table = new DifferentialRevision();
     $conn_r = $table->establishConnection('r');

src/applications/differential/query/DifferentialRevisionSearchEngine.php

@@ -55,6 +55,8 @@ final class DifferentialRevisionSearchEngine
 
   public function buildQueryFromSavedQuery(PhabricatorSavedQuery $saved) {
     $query = id(new DifferentialRevisionQuery())
+      ->needFlags(true)
+      ->needDrafts(true)
       ->needRelationships(true);
 
     $responsible_phids = $saved->getParameter('responsiblePHIDs', array());

src/applications/differential/storage/DifferentialRevision.php

@@ -41,6 +41,8 @@ final class DifferentialRevision extends DifferentialDAO
 
   private $reviewerStatus = self::ATTACHABLE;
   private $customFields = self::ATTACHABLE;
+  private $drafts = array();
+  private $flags = array();
 
   const TABLE_COMMIT          = 'differential_commit';
 
@@ -405,6 +407,26 @@ final class DifferentialRevision extends DifferentialDAO
     return DifferentialRevisionStatus::isClosedStatus($this->getStatus());
   }
 
+  public function getFlag(PhabricatorUser $viewer) {
+    return $this->assertAttachedKey($this->flags, $viewer->getPHID());
+  }
+
+  public function attachFlag(
+    PhabricatorUser $viewer,
+    PhabricatorFlag $flag = null) {
+    $this->flags[$viewer->getPHID()] = $flag;
+    return $this;
+  }
+
+  public function getDrafts(PhabricatorUser $viewer) {
+    return $this->assertAttachedKey($this->drafts, $viewer->getPHID());
+  }
+
+  public function attachDrafts(PhabricatorUser $viewer, array $drafts) {
+    $this->drafts[$viewer->getPHID()] = $drafts;
+    return $this;
+  }
+
 
 /* -(  HarbormasterBuildableInterface  )------------------------------------- */
 

src/applications/differential/view/DifferentialRevisionListView.php

@@ -6,8 +6,6 @@
 final class DifferentialRevisionListView extends AphrontView {
 
   private $revisions;
-  private $flags = array();
-  private $drafts = array();
   private $handles;
   private $fields;
   private $highlightAge;
@@ -57,30 +55,6 @@ final class DifferentialRevisionListView extends AphrontView {
     return $this;
   }
 
-  public function loadAssets() {
-    $user = $this->user;
-    if (!$user) {
-      throw new Exception("Call setUser() before loadAssets()!");
-    }
-    if ($this->revisions === null) {
-      throw new Exception("Call setRevisions() before loadAssets()!");
-    }
-
-    $this->flags = id(new PhabricatorFlagQuery())
-      ->setViewer($user)
-      ->withOwnerPHIDs(array($user->getPHID()))
-      ->withObjectPHIDs(mpull($this->revisions, 'getPHID'))
-      ->execute();
-
-    $this->drafts = id(new DifferentialRevisionQuery())
-      ->setViewer($user)
-      ->withIDs(mpull($this->revisions, 'getID'))
-      ->withDraftRepliesByAuthors(array($user->getPHID()))
-      ->execute();
-
-    return $this;
-  }
-
   public function render() {
 
     $user = $this->user;
@@ -105,8 +79,6 @@ final class DifferentialRevisionListView extends AphrontView {
     $this->initBehavior('phabricator-tooltips', array());
     $this->requireResource('aphront-tooltip-css');
 
-    $flagged = mpull($this->flags, null, 'getObjectPHID');
-
     foreach ($this->fields as $field) {
       $field->setHandles($this->handles);
     }
@@ -122,8 +94,8 @@ final class DifferentialRevisionListView extends AphrontView {
       $icons = array();
 
       $phid = $revision->getPHID();
-      if (isset($flagged[$phid])) {
+      $flag = $revision->getFlag($user);
-        $flag = $flagged[$phid];
+      if ($flag) {
         $flag_class = PhabricatorFlagColor::getCSSClass($flag->getColor());
         $icons['flag'] = phutil_tag(
           'div',
@@ -132,7 +104,8 @@ final class DifferentialRevisionListView extends AphrontView {
           ),
           '');
       }
-      if (array_key_exists($revision->getID(), $this->drafts)) {
+
+      if ($revision->getDrafts($user)) {
         $icons['draft'] = true;
       }
 

src/applications/diffusion/controller/DiffusionBrowseController.php

@@ -201,6 +201,8 @@ abstract class DiffusionBrowseController extends DiffusionController {
       ->setOrder(DifferentialRevisionQuery::ORDER_PATH_MODIFIED)
       ->setLimit(10)
       ->needRelationships(true)
+      ->needFlags(true)
+      ->needDrafts(true)
       ->execute();
 
     if (!$revisions) {
@@ -210,8 +212,7 @@ abstract class DiffusionBrowseController extends DiffusionController {
     $view = id(new DifferentialRevisionListView())
       ->setRevisions($revisions)
       ->setFields(DifferentialRevisionListView::getDefaultFields($user))
-      ->setUser($user)
+      ->setUser($user);
-      ->loadAssets();
 
     $phids = $view->getRequiredHandlePHIDs();
     $handles = $this->loadViewerHandles($phids);

src/applications/home/controller/PhabricatorHomeMainController.php

@@ -193,7 +193,9 @@ final class PhabricatorHomeMainController
       ->setViewer($user)
       ->withStatus(DifferentialRevisionQuery::STATUS_OPEN)
       ->withResponsibleUsers(array($user_phid))
-      ->needRelationships(true);
+      ->needRelationships(true)
+      ->needFlags(true)
+      ->needDrafts(true);
 
     $revisions = $revision_query->execute();
 
@@ -216,8 +218,7 @@ final class PhabricatorHomeMainController
       ->setHighlightAge(true)
       ->setRevisions(array_merge($blocking, $active))
       ->setFields(DifferentialRevisionListView::getDefaultFields($user))
-      ->setUser($user)
+      ->setUser($user);
-      ->loadAssets();
     $phids = array_merge(
       array($user_phid),
       $revision_view->getRequiredHandlePHIDs());