Add tree metadata prefetching to batch tree fetching path.

Summary:
Tree metadata prefetching is a second step after a normal fetch. Let's
add it to the batch tree fetching path. Ideally we'd unify this with the
non-batch tree fetching path a bit. Or just get rid of the non-batch path
entirely.

Reviewed By: chadaustin

Differential Revision: D28427267

fbshipit-source-id: e23fc667f609925e8671a71876165ad2cea161b9

eden/fs/store/hg/HgBackingStore.cpp

@@ -238,10 +238,54 @@ SemiFuture<unique_ptr<Tree>> HgBackingStore::getTree(
 void HgBackingStore::getTreeBatch(
     const std::vector<Hash>& ids,
     const std::vector<HgProxyHash>& hashes,
-    std::vector<folly::Promise<std::unique_ptr<Tree>>*> promises) {
+    std::vector<folly::Promise<std::unique_ptr<Tree>>*> promises,
+    bool prefetchMetadata) {
   auto writeBatch = localStore_->beginWrite();
 
-  datapackStore_.getTreeBatch(ids, hashes, writeBatch.get(), &promises);
+  std::vector<folly::Promise<std::unique_ptr<Tree>>> innerPromises;
+  innerPromises.reserve(promises.size());
+  std::vector<folly::SemiFuture<std::unique_ptr<TreeMetadata>>> metadataFutures;
+  metadataFutures.reserve(promises.size());
+
+  bool metadataEnabled =
+      metadataImporter_->metadataFetchingAvailable() && prefetchMetadata;
+
+  // Kick off all the fetching
+  auto proxyHash = hashes.begin();
+  auto id = ids.begin();
+  auto promise = promises.begin();
+  for (; promise != promises.end(); ++id, ++proxyHash, ++promise) {
+    innerPromises.emplace_back(folly::Promise<std::unique_ptr<Tree>>());
+
+    auto treeMetadataFuture =
+        folly::SemiFuture<std::unique_ptr<TreeMetadata>>::makeEmpty();
+    if (metadataEnabled) {
+      treeMetadataFuture =
+          metadataImporter_->getTreeMetadata(proxyHash->revHash(), *id);
+    }
+    metadataFutures.push_back(std::move(treeMetadataFuture));
+  }
+  datapackStore_.getTreeBatch(ids, hashes, writeBatch.get(), &innerPromises);
+
+  // Receive the fetches and tie the content and metadata together if needed.
+  auto innerPromise = innerPromises.begin();
+  promise = promises.begin();
+  auto treeMetadataFuture = std::make_move_iterator(metadataFutures.begin());
+  for (; innerPromise != innerPromises.end();
+       ++innerPromise, ++promise, ++treeMetadataFuture) {
+    // This innerPromise pattern is so we can retrieve the tree from the
+    // innerPromise and use it for tree metadata prefetching, without
+    // invalidating the passed in Promise.
+    if (!innerPromise->isFulfilled()) {
+      continue;
+    }
+
+    (*promise)->setWith([&]() mutable {
+      std::unique_ptr<Tree> tree = (*innerPromise).getFuture().get();
+      this->processTreeMetadata(std::move(*treeMetadataFuture), *tree);
+      return tree;
+    });
+  }
 }
 
 Future<unique_ptr<Tree>> HgBackingStore::importTreeImpl(

eden/fs/store/hg/HgBackingStore.h

@@ -79,7 +79,8 @@ class HgBackingStore {
   void getTreeBatch(
       const std::vector<Hash>& ids,
       const std::vector<HgProxyHash>& hashes,
-      std::vector<folly::Promise<std::unique_ptr<Tree>>*> promises);
+      std::vector<folly::Promise<std::unique_ptr<Tree>>*> promises,
+      bool prefetchMetadata);
   void processTreeMetadata(
       folly::SemiFuture<std::unique_ptr<TreeMetadata>>&& treeMetadataFuture,
       const Tree& tree);

eden/fs/store/hg/HgDatapackStore.cpp

@@ -171,7 +171,7 @@ void HgDatapackStore::getTreeBatch(
     const std::vector<Hash>& ids,
     const std::vector<HgProxyHash>& hashes,
     LocalStore::WriteBatch* writeBatch,
-    std::vector<folly::Promise<std::unique_ptr<Tree>>*>* promises) {
+    std::vector<folly::Promise<std::unique_ptr<Tree>>>* promises) {
   std::vector<Hash> treehashes;
   std::vector<std::pair<folly::ByteRange, folly::ByteRange>> requests;
 
@@ -201,7 +201,7 @@ void HgDatapackStore::getTreeBatch(
       false,
       [promises = promises, ids = ids, hashes = hashes, writeBatch, requests](
           size_t index, std::shared_ptr<RustTree> content) mutable {
-        (*promises)[index]->setWith([&] {
+        (*promises)[index].setWith([&] {
           XLOGF(
               DBG4,
               "Imported tree name={} node={}",

eden/fs/store/hg/HgDatapackStore.h

@@ -49,7 +49,7 @@ class HgDatapackStore {
       const std::vector<Hash>& ids,
       const std::vector<HgProxyHash>& hashes,
       LocalStore::WriteBatch* writeBatch,
-      std::vector<folly::Promise<std::unique_ptr<Tree>>*>* promises);
+      std::vector<folly::Promise<std::unique_ptr<Tree>>>* promises);
 
   std::unique_ptr<Tree> getTree(
       const RelativePath& path,

eden/fs/store/hg/HgQueuedBackingStore.cpp

@@ -169,11 +169,13 @@ void HgQueuedBackingStore::processTreeImportRequests(
   proxyHashes.reserve(requests.size());
   promises.reserve(requests.size());
 
+  bool prefetchMetadata = false;
   for (auto& request : requests) {
     auto* treeImport = request->getRequest<HgImportRequest::TreeImport>();
     auto& hash = treeImport->hash;
     auto* promise =
         request->getPromise<HgImportRequest::TreeImport::Response>();
+    prefetchMetadata |= treeImport->prefetchMetadata;
 
     traceBus_->publish(HgImportTraceEvent::start(
         request->getUnique(), HgImportTraceEvent::TREE, treeImport->proxyHash));
@@ -188,7 +190,7 @@ void HgQueuedBackingStore::processTreeImportRequests(
     promises.emplace_back(promise);
   }
 
-  backingStore_->getTreeBatch(hashes, proxyHashes, promises);
+  backingStore_->getTreeBatch(hashes, proxyHashes, promises, prefetchMetadata);
 
   {
     auto request = requests.begin();

eden/fs/store/hg/test/HgBackingStoreTest.cpp

@@ -175,7 +175,5 @@ TEST_F(HgBackingStoreTest, skipMetadataPrefetch) {
   metadataImporter->getTreeMetadataCalled = false;
   backingStore->getTree(tree->getHash(), ObjectFetchContext::getNullContext())
       .get(kTestTimeout);
-  // Temporarily comment this off since the current diff doesn't implement batch
-  // tree metadata fetching. The next diff uncomments this line.
-  // EXPECT_TRUE(metadataImporter->getTreeMetadataCalled);
+  EXPECT_TRUE(metadataImporter->getTreeMetadataCalled);
 }