From 3460da5f3402d69be2e50d220b35e17916b10098 Mon Sep 17 00:00:00 2001
From: epriestley <git@epriestley.com>
Date: Wed, 8 Aug 2012 18:55:07 -0700
Subject: [PATCH] Fix limits in queries

Summary: I think this is simpler? Includes test cases.

Test Plan: Ran tests. Loaded /paste/.

Reviewers: vrana, nh

Reviewed By: vrana

CC: aran

Differential Revision: https://secure.phabricator.com/D3209
---
 .../__tests__/PhabricatorPolicyTestCase.php   | 30 +++++++++++++++++
 .../query/policy/PhabricatorPolicyQuery.php   | 32 ++++---------------
 2 files changed, 37 insertions(+), 25 deletions(-)

diff --git a/src/applications/policy/__tests__/PhabricatorPolicyTestCase.php b/src/applications/policy/__tests__/PhabricatorPolicyTestCase.php
index 24fb6a08b5..08a793b8b7 100644
--- a/src/applications/policy/__tests__/PhabricatorPolicyTestCase.php
+++ b/src/applications/policy/__tests__/PhabricatorPolicyTestCase.php
@@ -136,6 +136,36 @@ final class PhabricatorPolicyTestCase extends PhabricatorTestCase {
   }
 
 
+  /**
+   * Test limits.
+   */
+  public function testLimits() {
+    $results = array(
+      $this->buildObject(PhabricatorPolicies::POLICY_USER),
+      $this->buildObject(PhabricatorPolicies::POLICY_USER),
+      $this->buildObject(PhabricatorPolicies::POLICY_USER),
+      $this->buildObject(PhabricatorPolicies::POLICY_USER),
+      $this->buildObject(PhabricatorPolicies::POLICY_USER),
+      $this->buildObject(PhabricatorPolicies::POLICY_USER),
+    );
+
+    $query = new PhabricatorPolicyTestQuery();
+    $query->setResults($results);
+    $query->setViewer($this->buildUser('user'));
+
+    $this->assertEqual(
+      3,
+      count($query->setLimit(3)->setOffset(0)->execute()),
+      'Limits work.');
+
+    $this->assertEqual(
+      2,
+      count($query->setLimit(3)->setOffset(4)->execute()),
+      'Limit + offset work.');
+  }
+
+
+
   /**
    * Test an object for visibility across multiple user specifications.
    */
diff --git a/src/infrastructure/query/policy/PhabricatorPolicyQuery.php b/src/infrastructure/query/policy/PhabricatorPolicyQuery.php
index ffdd362789..d9c29e1350 100644
--- a/src/infrastructure/query/policy/PhabricatorPolicyQuery.php
+++ b/src/infrastructure/query/policy/PhabricatorPolicyQuery.php
@@ -153,34 +153,16 @@ abstract class PhabricatorPolicyQuery extends PhabricatorOffsetPagedQuery {
     $limit  = (int)$this->getLimit();
     $count  = 0;
 
-    $need = null;
-    if ($offset) {
-      $need = $offset + $limit;
-    }
+    $need = $offset + $limit;
 
     $this->willExecute();
 
     do {
-
-      // Figure out how many results to load. "0" means "all results".
-      $load = 0;
-      if ($need && ($count < $offset)) {
-        // This cap is just an arbitrary limit to keep memory usage from going
-        // crazy for large offsets; we can't execute them efficiently, but
-        // it should be possible to execute them without crashing.
-        $load = min($need, 1024);
-      } else if ($limit) {
-        // Otherwise, just load the number of rows we're after. Note that it
-        // might be more efficient to load more rows than this (if we expect
-        // about 5% of objects to be filtered, loading 105% of the limit might
-        // be better) or fewer rows than this (if we already have 95 rows and
-        // only need 100, loading only 5 rows might be better), but we currently
-        // just use the simplest heuristic since we don't have enough data
-        // about policy queries in the real world to tweak it.
-        $load = $limit;
+      if ($need) {
+        $this->rawResultLimit = min($need - $count, 1024);
+      } else {
+        $this->rawResultLimit = 0;
       }
-      $this->rawResultLimit = $load;
-
 
       $page = $this->loadPage();
 
@@ -202,13 +184,13 @@ abstract class PhabricatorPolicyQuery extends PhabricatorOffsetPagedQuery {
         }
       }
 
-      if (!$load) {
+      if (!$this->rawResultLimit) {
         // If we don't have a load count, we loaded all the results. We do
         // not need to load another page.
         break;
       }
 
-      if (count($page) < $load) {
+      if (count($page) < $this->rawResultLimit) {
         // If we have a load count but the unfiltered results contained fewer
         // objects, we know this was the last page of objects; we do not need
         // to load another page because we can deduce it would be empty.