Fix ASAN failure in FutureDAG test

[folly.git] / folly / experimental / FutureDAG.h

diff --git a/folly/experimental/FutureDAG.h b/folly/experimental/FutureDAG.h
index 42391336eccfef95a189df69b25cb206c8236c59..ee3dcffbf649f6f74d6953f50f9492640ca39a19 100644 (file)
--- a/folly/experimental/FutureDAG.h
+++ b/folly/experimental/FutureDAG.h
@@ -1,5 +1,5 @@
 /*
- * Copyright 2016 Facebook, Inc.
+ * Copyright 2017 Facebook, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -35,7 +35,11 @@ class FutureDAG : public std::enable_shared_from_this<FutureDAG> {
   }
 
   void remove(Handle a) {
-    if (nodes.size() > a && nodes[a].hasDependents) {
+    if (a >= nodes.size()) {
+      return;
+    }
+
+    if (nodes[a].hasDependents) {
       for (auto& node : nodes) {
         auto& deps = node.dependencies;
         deps.erase(
@@ -47,6 +51,7 @@ class FutureDAG : public std::enable_shared_from_this<FutureDAG> {
         }
       }
     }
+
     nodes.erase(nodes.begin() + a);
   }
 
@@ -129,10 +134,9 @@ class FutureDAG : public std::enable_shared_from_this<FutureDAG> {
     }
 
     nodes[sourceHandle].promise.setValue();
-    auto that = shared_from_this();
-    return nodes[sinkHandle].promise.getFuture().ensure([that] {}).then(
-        [this, sourceHandle, sinkHandle]() {
-          clean_state(sourceHandle, sinkHandle);
+    return nodes[sinkHandle].promise.getFuture().then(
+        [that = shared_from_this(), sourceHandle, sinkHandle]() {
+          that->clean_state(sourceHandle, sinkHandle);
         });
   }
 
@@ -205,19 +209,19 @@ class FutureDAGFunctor {
   std::vector<T> dep_states;
   T result() {
     return state;
-  };
+  }
   // execReset() runs DAG & clears all nodes except for source
   void execReset() {
     this->dag->go().get();
     this->dag->reset();
-  };
+  }
   void exec() {
     this->dag->go().get();
-  };
-  virtual void operator()(){};
+  }
+  virtual void operator()(){}
   explicit FutureDAGFunctor(T init_val) : state(init_val) {}
   FutureDAGFunctor() : state() {}
-  virtual ~FutureDAGFunctor(){};
+  virtual ~FutureDAGFunctor(){}
 };
 
-} // folly
+} // namespace folly