From da9e05e8a7d6667424bdffcd10086cc9a85ff1a9 Mon Sep 17 00:00:00 2001
From: Yunqiao Zhang <yunqiao@fb.com>
Date: Fri, 26 Aug 2016 12:53:24 -0700
Subject: [PATCH] collectOne

Summary:
The resultant future of collectOne will be fulfilled when the first one of
the future in the list completes without exception. If all input futures throws
exception, the resultant future will get the last exception that was thrown.

Reviewed By: andriigrynenko

Differential Revision: D3764760

fbshipit-source-id: 76484254e35182eddc8266865853d65c28170f82
---
 folly/futures/Future-inl.h         | 31 ++++++++++++++
 folly/futures/helpers.h            | 17 ++++++++
 folly/futures/test/CollectTest.cpp | 67 ++++++++++++++++++++++++++++++
 3 files changed, 115 insertions(+)

diff --git a/folly/futures/Future-inl.h b/folly/futures/Future-inl.h
index 0d171006..1cd4035e 100644
--- a/folly/futures/Future-inl.h
+++ b/folly/futures/Future-inl.h
@@ -695,6 +695,37 @@ collectAny(InputIterator first, InputIterator last) {
   return ctx->p.getFuture();
 }
 
+// collectAnyWithoutException (iterator)
+
+template <class InputIterator>
+Future<std::pair<
+    size_t,
+    typename std::iterator_traits<InputIterator>::value_type::value_type>>
+collectAnyWithoutException(InputIterator first, InputIterator last) {
+  typedef
+      typename std::iterator_traits<InputIterator>::value_type::value_type T;
+
+  struct CollectAnyWithoutExceptionContext {
+    CollectAnyWithoutExceptionContext(){};
+    Promise<std::pair<size_t, T>> p;
+    std::atomic<bool> done{false};
+    std::atomic<size_t> nFulfilled{0};
+    size_t nTotal;
+  };
+
+  auto ctx = std::make_shared<CollectAnyWithoutExceptionContext>();
+  ctx->nTotal = std::distance(first, last);
+
+  mapSetCallback<T>(first, last, [ctx](size_t i, Try<T>&& t) {
+    if (!t.hasException() && !ctx->done.exchange(true)) {
+      ctx->p.setValue(std::make_pair(i, std::move(t.value())));
+    } else if (++ctx->nFulfilled == ctx->nTotal) {
+      ctx->p.setException(t.exception());
+    }
+  });
+  return ctx->p.getFuture();
+}
+
 // collectN (iterator)
 
 template <class InputIterator>
diff --git a/folly/futures/helpers.h b/folly/futures/helpers.h
index 75f92537..7f074a06 100644
--- a/folly/futures/helpers.h
+++ b/folly/futures/helpers.h
@@ -226,6 +226,23 @@ auto collectAny(Collection&& c) -> decltype(collectAny(c.begin(), c.end())) {
   return collectAny(c.begin(), c.end());
 }
 
+/** Similar to collectAny, collectAnyWithoutException return the first Future to
+ * complete without exceptions. If none of the future complete without
+ * excpetions, the last exception will be returned as a result.
+  */
+template <class InputIterator>
+Future<std::pair<
+    size_t,
+    typename std::iterator_traits<InputIterator>::value_type::value_type>>
+collectAnyWithoutException(InputIterator first, InputIterator last);
+
+/// Sugar for the most common case
+template <class Collection>
+auto collectAnyWithoutException(Collection&& c)
+    -> decltype(collectAnyWithoutException(c.begin(), c.end())) {
+  return collectAnyWithoutException(c.begin(), c.end());
+}
+
 /** when n Futures have completed, the Future completes with a vector of
   the index and Try of those n Futures (the indices refer to the original
   order, but the result vector will be in an arbitrary order)
diff --git a/folly/futures/test/CollectTest.cpp b/folly/futures/test/CollectTest.cpp
index 5f05a0b6..673161aa 100644
--- a/folly/futures/test/CollectTest.cpp
+++ b/folly/futures/test/CollectTest.cpp
@@ -333,6 +333,73 @@ TEST(Collect, collectAny) {
   }
 }
 
+TEST(Collect, collectAnyWithoutException) {
+  {
+    std::vector<Promise<int>> promises(10);
+    std::vector<Future<int>> futures;
+
+    for (auto& p : promises) {
+      futures.push_back(p.getFuture());
+    }
+
+    auto onef = collectAnyWithoutException(futures);
+
+    /* futures were moved in, so these are invalid now */
+    EXPECT_FALSE(onef.isReady());
+
+    promises[7].setValue(42);
+    EXPECT_TRUE(onef.isReady());
+    auto& idx_fut = onef.value();
+    EXPECT_EQ(7, idx_fut.first);
+    EXPECT_EQ(42, idx_fut.second);
+  }
+
+  // some exception before ready
+  {
+    std::vector<Promise<int>> promises(10);
+    std::vector<Future<int>> futures;
+
+    for (auto& p : promises) {
+      futures.push_back(p.getFuture());
+    }
+
+    auto onef = collectAnyWithoutException(futures);
+
+    EXPECT_FALSE(onef.isReady());
+
+    promises[3].setException(eggs);
+    EXPECT_FALSE(onef.isReady());
+    promises[4].setException(eggs);
+    EXPECT_FALSE(onef.isReady());
+    promises[0].setValue(99);
+    EXPECT_TRUE(onef.isReady());
+    auto& idx_fut = onef.value();
+    EXPECT_EQ(0, idx_fut.first);
+    EXPECT_EQ(99, idx_fut.second);
+  }
+
+  // all exceptions
+  {
+    std::vector<Promise<int>> promises(10);
+    std::vector<Future<int>> futures;
+
+    for (auto& p : promises) {
+      futures.push_back(p.getFuture());
+    }
+
+    auto onef = collectAnyWithoutException(futures);
+
+    EXPECT_FALSE(onef.isReady());
+    for (int i = 0; i < 9; ++i) {
+      promises[i].setException(eggs);
+    }
+    EXPECT_FALSE(onef.isReady());
+
+    promises[9].setException(eggs);
+    EXPECT_TRUE(onef.isReady());
+    EXPECT_TRUE(onef.hasException());
+  }
+}
 
 TEST(Collect, alreadyCompleted) {
   {
-- 
2.34.1