Improve Conduit performance for custom fields

Summary:
Ref T11404. Depends on D16350.

Currently, custom fields can issue "N+1" queries in some cases, so querying 100 revisions issues 100 extra queries.

This affects all `*.search` endpoints for objects with custom fields, and some older endpoints (notably `differential.query`).

This change bulk loads "normal" custom fields, which gets rid of some of these queries. Instead of loading fields for each object, we build a big list of all fields and load them all at once.

The next change will tackle the remaining inefficient edge queries.

Test Plan:
  - Configured a custom field with normal database storage in Differential.
  - Ran `differential.query`, looking at custom fields in results for correctness.
  - Ran `differential.revision.search`, looking at custom fields in results for correctness.
  - In both cases, observed queries drop from `3N` to `2N` (all the "normal" custom field stuff got bulk loaded).

Reviewers: yelirekim, chad

Reviewed By: chad

Maniphest Tasks: T11404

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

src/__phutil_library_map__.php

@@ -2258,6 +2258,7 @@ phutil_register_library_map(array(
     'PhabricatorCustomFieldNumericIndexStorage' => 'infrastructure/customfield/storage/PhabricatorCustomFieldNumericIndexStorage.php',
     'PhabricatorCustomFieldSearchEngineExtension' => 'infrastructure/customfield/engineextension/PhabricatorCustomFieldSearchEngineExtension.php',
     'PhabricatorCustomFieldStorage' => 'infrastructure/customfield/storage/PhabricatorCustomFieldStorage.php',
+    'PhabricatorCustomFieldStorageQuery' => 'infrastructure/customfield/query/PhabricatorCustomFieldStorageQuery.php',
     'PhabricatorCustomFieldStringIndexStorage' => 'infrastructure/customfield/storage/PhabricatorCustomFieldStringIndexStorage.php',
     'PhabricatorCustomHeaderConfigType' => 'applications/config/custom/PhabricatorCustomHeaderConfigType.php',
     'PhabricatorDaemon' => 'infrastructure/daemon/PhabricatorDaemon.php',
@@ -6994,6 +6995,7 @@ phutil_register_library_map(array(
     'PhabricatorCustomFieldNumericIndexStorage' => 'PhabricatorCustomFieldIndexStorage',
     'PhabricatorCustomFieldSearchEngineExtension' => 'PhabricatorSearchEngineExtension',
     'PhabricatorCustomFieldStorage' => 'PhabricatorLiskDAO',
+    'PhabricatorCustomFieldStorageQuery' => 'Phobject',
     'PhabricatorCustomFieldStringIndexStorage' => 'PhabricatorCustomFieldIndexStorage',
     'PhabricatorCustomHeaderConfigType' => 'PhabricatorConfigOptionType',
     'PhabricatorDaemon' => 'PhutilDaemon',

src/applications/differential/conduit/DifferentialConduitAPIMethod.php

@@ -150,21 +150,38 @@ abstract class DifferentialConduitAPIMethod extends ConduitAPIMethod {
     array $revisions) {
     assert_instances_of($revisions, 'DifferentialRevision');
 
-    $results = array();
+    $field_lists = array();
     foreach ($revisions as $revision) {
-      // TODO: This is inefficient and issues a query for each object.
+      $revision_phid = $revision->getPHID();
+
       $field_list = PhabricatorCustomField::getObjectFields(
         $revision,
         PhabricatorCustomField::ROLE_CONDUIT);
 
       $field_list
         ->setViewer($viewer)
-        ->readFieldsFromStorage($revision);
+        ->readFieldsFromObject($revision);
 
+      $field_lists[$revision_phid] = $field_list;
+    }
+
+    $all_fields = array();
+    foreach ($field_lists as $field_list) {
+      foreach ($field_list->getFields() as $field) {
+        $all_fields[] = $field;
+      }
+    }
+
+    id(new PhabricatorCustomFieldStorageQuery())
+      ->addFields($all_fields)
+      ->execute();
+
+    $results = array();
+    foreach ($field_lists as $revision_phid => $field_list) {
       foreach ($field_list->getFields() as $field) {
         $field_key = $field->getFieldKeyForConduit();
         $value = $field->getConduitDictionaryValue();
-        $results[$revision->getPHID()][$field_key] = $value;
+        $results[$revision_phid][$field_key] = $value;
       }
     }
 

src/infrastructure/customfield/engineextension/PhabricatorCustomFieldSearchEngineExtension.php

@@ -80,17 +80,42 @@ final class PhabricatorCustomFieldSearchEngineExtension
     return $map;
   }
 
-  public function getFieldValuesForConduit($object, $data) {
+  public function loadExtensionConduitData(array $objects) {
-    // TODO: This is currently very inefficient. We should bulk-load these
+    $viewer = $this->getViewer();
-    // field values instead.
+
+    $field_map = array();
+    foreach ($objects as $object) {
+      $object_phid = $object->getPHID();
 
       $fields = PhabricatorCustomField::getObjectFields(
         $object,
         PhabricatorCustomField::ROLE_CONDUIT);
 
       $fields
-      ->setViewer($this->getViewer())
+        ->setViewer($viewer)
-      ->readFieldsFromStorage($object);
+        ->readFieldsFromObject($object);
+
+      $field_map[$object_phid] = $fields;
+    }
+
+    $all_fields = array();
+    foreach ($field_map as $field_list) {
+      foreach ($field_list->getFields() as $field) {
+        $all_fields[] = $field;
+      }
+    }
+
+    id(new PhabricatorCustomFieldStorageQuery())
+      ->addFields($all_fields)
+      ->execute();
+
+    return array(
+      'fields' => $field_map,
+    );
+  }
+
+  public function getFieldValuesForConduit($object, $data) {
+    $fields = $data['fields'][$object->getPHID()];
 
     $map = array();
     foreach ($fields->getFields() as $field) {

src/infrastructure/customfield/field/PhabricatorCustomFieldList.php

@@ -29,6 +29,19 @@ final class PhabricatorCustomFieldList extends Phobject {
     return $this;
   }
 
+  public function readFieldsFromObject(
+    PhabricatorCustomFieldInterface $object) {
+
+    $fields = $this->getFields();
+
+    foreach ($fields as $field) {
+      $field
+        ->setObject($object)
+        ->readValueFromObject($object);
+    }
+
+    return $this;
+  }
 
   /**
    * Read stored values for all fields which support storage.
@@ -39,48 +52,12 @@ final class PhabricatorCustomFieldList extends Phobject {
   public function readFieldsFromStorage(
     PhabricatorCustomFieldInterface $object) {
 
-    foreach ($this->fields as $field) {
+    $this->readFieldsFromObject($object);
-      $field->setObject($object);
-      $field->readValueFromObject($object);
-    }
 
-    $keys = array();
+    $fields = $this->getFields();
-    foreach ($this->fields as $field) {
+    id(new PhabricatorCustomFieldStorageQuery())
-      if ($field->shouldEnableForRole(PhabricatorCustomField::ROLE_STORAGE)) {
+      ->addFields($fields)
-        $keys[$field->getFieldIndex()] = $field;
+      ->execute();
-      }
-    }
-
-    if (!$keys) {
-      return $this;
-    }
-
-    // NOTE: We assume all fields share the same storage. This isn't guaranteed
-    // to be true, but always is for now.
-
-    $table = head($keys)->newStorageObject();
-
-    $objects = array();
-    if ($object->getPHID()) {
-      $objects = $table->loadAllWhere(
-        'objectPHID = %s AND fieldIndex IN (%Ls)',
-        $object->getPHID(),
-        array_keys($keys));
-      $objects = mpull($objects, null, 'getFieldIndex');
-    }
-
-    foreach ($keys as $key => $field) {
-      $storage = idx($objects, $key);
-      if ($storage) {
-        $field->setValueFromStorage($storage->getFieldValue());
-        $field->didSetValueFromStorage();
-      } else if ($object->getPHID()) {
-        // NOTE: We set this only if the object exists. Otherwise, we allow the
-        // field to retain any default value it may have.
-        $field->setValueFromStorage(null);
-        $field->didSetValueFromStorage();
-      }
-    }
 
     return $this;
   }

src/infrastructure/customfield/query/PhabricatorCustomFieldStorageQuery.php

@@ -0,0 +1,84 @@
+<?php
+
+/**
+ * Load custom field data from storage.
+ *
+ * This query loads the data directly into the field objects and does not
+ * return it to the caller. It can bulk load data for any list of fields,
+ * even if they have different objects or object types.
+ */
+final class PhabricatorCustomFieldStorageQuery extends Phobject {
+
+  private $fieldMap = array();
+  private $storageSources = array();
+
+  public function addFields(array $fields) {
+    assert_instances_of($fields, 'PhabricatorCustomField');
+
+    foreach ($fields as $field) {
+      $this->addField($field);
+    }
+
+    return $this;
+  }
+
+  public function addField(PhabricatorCustomField $field) {
+    $role_storage = PhabricatorCustomField::ROLE_STORAGE;
+
+    if (!$field->shouldEnableForRole($role_storage)) {
+      return $this;
+    }
+
+    $storage = $field->newStorageObject();
+    $source_key = $storage->getStorageSourceKey();
+
+    $this->fieldMap[$source_key][] = $field;
+
+    if (empty($this->storageSources[$source_key])) {
+      $this->storageSources[$source_key] = $storage;
+    }
+
+    return $this;
+  }
+
+  public function execute() {
+    foreach ($this->storageSources as $source_key => $storage) {
+      $fields = idx($this->fieldMap, $source_key, array());
+      $this->loadFieldsFromStorage($storage, $fields);
+    }
+  }
+
+  private function loadFieldsFromStorage($storage, array $fields) {
+    // Only try to load fields which have a persisted object.
+    $loadable = array();
+    foreach ($fields as $key => $field) {
+      $object = $field->getObject();
+      $phid = $object->getPHID();
+      if (!$phid) {
+        continue;
+      }
+
+      $loadable[$key] = $field;
+    }
+
+    if ($loadable) {
+      $data = $storage->loadStorageSourceData($loadable);
+    } else {
+      $data = array();
+    }
+
+    foreach ($fields as $key => $field) {
+      if (array_key_exists($key, $data)) {
+        $value = $data[$key];
+        $field->setValueFromStorage($value);
+        $field->didSetValueFromStorage();
+      } else if (isset($loadable[$key])) {
+        // NOTE: We set this only if the object exists. Otherwise, we allow
+        // the field to retain any default value it may have.
+        $field->setValueFromStorage(null);
+        $field->didSetValueFromStorage();
+      }
+    }
+  }
+
+}

src/infrastructure/customfield/storage/PhabricatorCustomFieldStorage.php

@@ -23,4 +23,68 @@ abstract class PhabricatorCustomFieldStorage
     ) + parent::getConfiguration();
   }
 
+
+  /**
+   * Get a key which uniquely identifies this storage source.
+   *
+   * When loading custom fields, fields using sources with the same source key
+   * are loaded in bulk.
+   *
+   * @return string Source identifier.
+   */
+  final public function getStorageSourceKey() {
+    return $this->getApplicationName().'/'.$this->getTableName();
+  }
+
+
+  /**
+   * Load stored data for custom fields.
+   *
+   * Given a map of fields, return a map with any stored data for those fields.
+   * The keys in the result should correspond to the keys in the input. The
+   * fields in the list may belong to different objects.
+   *
+   * @param map<string, PhabricatorCustomField> Map of fields.
+   * @return map<String, PhabricatorCustomField> Map of available field data.
+   */
+  final public function loadStorageSourceData(array $fields) {
+    $map = array();
+    $indexes = array();
+    $object_phids = array();
+
+    foreach ($fields as $key => $field) {
+      $index = $field->getFieldIndex();
+      $object_phid = $field->getObject()->getPHID();
+
+      $map[$index][$object_phid] = $key;
+      $indexes[$index] = $index;
+      $object_phids[$object_phid] = $object_phid;
+    }
+
+    if (!$indexes) {
+      return array();
+    }
+
+    $conn = $this->establishConnection('r');
+    $rows = queryfx_all(
+      $conn,
+      'SELECT objectPHID, fieldIndex, fieldValue FROM %T
+        WHERE objectPHID IN (%Ls) AND fieldIndex IN (%Ls)',
+      $this->getTableName(),
+      $object_phids,
+      $indexes);
+
+    $result = array();
+    foreach ($rows as $row) {
+      $index = $row['fieldIndex'];
+      $object_phid = $row['objectPHID'];
+      $value = $row['fieldValue'];
+
+      $key = $map[$index][$object_phid];
+      $result[$key] = $value;
+    }
+
+    return $result;
+  }
+
 }