From e602963840f91a673443aff98d14d339194882d1 Mon Sep 17 00:00:00 2001
From: Vladislav Isenbaev <winger@fb.com>
Date: Thu, 30 Apr 2015 15:39:27 -0700
Subject: [PATCH] Fix race condition in collect(..)

Summary: This is a temporary fix (until D2015320 is checked in) for race condition(s) in collect(..) method.

Test Plan:
Run unit tests
Run buffalo_aggregator canary

Reviewed By: jsedgwick@fb.com

Subscribers: folly-diffs@, jsedgwick, yfeldblum, chalfant

FB internal diff: D2037406

Tasks: 6894157

Signature: t1:2037406:1430435227:ed9612d016cdbd708e2deba02dc4fe0b59632f5a
---
 folly/futures/Future-inl.h | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/folly/futures/Future-inl.h b/folly/futures/Future-inl.h
index 927a3142..00cfba22 100644
--- a/folly/futures/Future-inl.h
+++ b/folly/futures/Future-inl.h
@@ -620,13 +620,13 @@ struct CollectContext {
     Optional<T>
    >::type VecT;
 
-  explicit CollectContext(int n) : count(0), threw(false) {
+  explicit CollectContext(int n) : count(0), success_count(0), threw(false) {
     results.resize(n);
   }
 
   Promise<std::vector<T>> p;
   std::vector<VecT> results;
-  std::atomic<size_t> count;
+  std::atomic<size_t> count, success_count;
   std::atomic_bool threw;
 
   typedef std::vector<T> result_type;
@@ -647,10 +647,10 @@ struct CollectContext {
 template <>
 struct CollectContext<void> {
 
-  explicit CollectContext(int n) : count(0), threw(false) {}
+  explicit CollectContext(int n) : count(0), success_count(0), threw(false) {}
 
   Promise<void> p;
-  std::atomic<size_t> count;
+  std::atomic<size_t> count, success_count;
   std::atomic_bool threw;
 
   typedef void result_type;
@@ -690,7 +690,6 @@ collect(InputIterator first, InputIterator last) {
      assert(i < n);
      auto& f = *first;
      f.setCallback_([ctx, i, n](Try<T> t) {
-       auto c = ++ctx->count;
 
        if (t.hasException()) {
          if (!ctx->threw.exchange(true)) {
@@ -698,12 +697,12 @@ collect(InputIterator first, InputIterator last) {
          }
        } else if (!ctx->threw) {
          ctx->addResult(i, t);
-         if (c == n) {
+         if (++ctx->success_count == n) {
            ctx->setValue();
          }
        }
 
-       if (c == n) {
+       if (++ctx->count == n) {
          delete ctx;
        }
      });
-- 
2.34.1