Add AsyncSource template for speculative execution (#903)

Summary:
AsyncSource encapsulates a background computation. These can be
scheduled on a background executor. The difference between this and a
future is that when the user requires the result, AsyncSource will
perform the async peration on the caller's thread and will turn the
background operation into a no-op.

Pull Request resolved: #903

Reviewed By: pedroerp

Differential Revision: D33687621

Pulled By: oerling

fbshipit-source-id: 46f3e7613bd67bc67e7b43336fc64e0d74c9c8ba

Author: Orri Erling

velox/common/base/AsyncSource.h

@@ -0,0 +1,110 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <functional>
+#include <memory>
+
+#include <folly/executors/QueuedImmediateExecutor.h>
+#include <folly/futures/Future.h>
+
+namespace facebook::velox {
+
+// A future-like object that prefabricates Items on an executor and
+// allows consumer threads to pick items as they are ready. If the
+// consumer needs the item before the executor started making it,
+// the consumer will make it instead. If multiple consumers request
+// the same item, exactly one gets it.
+template <typename Item>
+class AsyncSource {
+ public:
+  explicit AsyncSource(std::function<std::unique_ptr<Item>()> make)
+      : make_(make) {}
+
+  // Makes an item if it is not already made. To be called on a background
+  // executor.
+  void prepare() {
+    std::function<std::unique_ptr<Item>()> make = nullptr;
+    {
+      std::lock_guard<std::mutex> l(mutex_);
+      if (!make_) {
+        return;
+      }
+      making_ = true;
+      std::swap(make, make_);
+    }
+    item_ = make();
+    {
+      std::lock_guard<std::mutex> l(mutex_);
+      making_ = false;
+      if (promise_) {
+        promise_->setValue(true);
+        promise_ = nullptr;
+      }
+    }
+  }
+
+  // Returns the item to the first caller and nullptr to subsequent callers. If
+  // the item is preparing on the executor, waits for the item and otherwise
+  // makes it on the caller thread.
+  std::unique_ptr<Item> move() {
+    std::function<std::unique_ptr<Item>()> make = nullptr;
+    folly::SemiFuture<bool> wait(false);
+    {
+      std::lock_guard<std::mutex> l(mutex_);
+      if (item_) {
+        return std::move(item_);
+      }
+      if (promise_) {
+        // Somebody else is now waiting for the item to be made.
+        return nullptr;
+      }
+      if (making_) {
+        promise_ = std::make_unique<folly::Promise<bool>>();
+        wait = promise_->getSemiFuture();
+      } else {
+        if (!make_) {
+          return nullptr;
+        }
+        std::swap(make, make_);
+      }
+    }
+    // Outside of mutex_.
+    if (make) {
+      return make();
+    }
+    auto& exec = folly::QueuedImmediateExecutor::instance();
+    std::move(wait).via(&exec).wait();
+    std::lock_guard<std::mutex> l(mutex_);
+    return std::move(item_);
+  }
+
+  // If true, move() will not block. But there is no guarantee that somebody
+  // else will not get the item first.
+  bool hasValue() const {
+    return item_ != nullptr;
+  }
+
+ private:
+  std::mutex mutex_;
+  // True if 'prepare() is making the item.
+  bool making_{false};
+  std::unique_ptr<folly::Promise<bool>> promise_;
+  std::unique_ptr<Item> item_;
+  std::function<std::unique_ptr<Item>()> make_;
+};
+} // namespace facebook::velox

velox/common/tests/AsyncSourceTest.cpp

@@ -0,0 +1,99 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "velox/common/base/AsyncSource.h"
+#include <fmt/format.h>
+#include <folly/Random.h>
+#include <folly/Synchronized.h>
+#include <gtest/gtest.h>
+#include <thread>
+
+using namespace facebook::velox;
+
+// A sample class to be constructed via AsyncSource.
+struct Gizmo {
+  explicit Gizmo(int32_t _id) : id(_id) {}
+
+  const int32_t id;
+};
+
+TEST(AsyncSourceTest, basic) {
+  AsyncSource<Gizmo> gizmo([]() { return std::make_unique<Gizmo>(11); });
+  EXPECT_FALSE(gizmo.hasValue());
+  gizmo.prepare();
+  EXPECT_TRUE(gizmo.hasValue());
+  auto value = gizmo.move();
+  EXPECT_FALSE(gizmo.hasValue());
+  EXPECT_EQ(11, value->id);
+}
+
+TEST(AsyncSourceTest, threads) {
+  constexpr int32_t kNumThreads = 10;
+  constexpr int32_t kNumGizmos = 2000;
+  folly::Synchronized<std::unordered_set<int32_t>> results;
+  std::vector<std::shared_ptr<AsyncSource<Gizmo>>> gizmos;
+  for (auto i = 0; i < kNumGizmos; ++i) {
+    gizmos.push_back(std::make_shared<AsyncSource<Gizmo>>([i]() {
+      std::this_thread::sleep_for(std::chrono::milliseconds(1)); // NOLINT
+      return std::make_unique<Gizmo>(i);
+    }));
+  }
+
+  std::vector<std::thread> threads;
+  threads.reserve(kNumThreads);
+  for (int32_t threadIndex = 0; threadIndex < kNumThreads; ++threadIndex) {
+    threads.push_back(std::thread([threadIndex, &gizmos, &results]() {
+      if (threadIndex < kNumThreads / 2) {
+        // The first half of the threads prepare Gizmos in the background.
+        for (auto i = 0; i < kNumGizmos; ++i) {
+          gizmos[i]->prepare();
+        }
+      } else {
+        // The rest of the threads first get random Gizmos and then do a pass
+        // over all the Gizmos to make sure all get collected. We assert that
+        // each Gizmo is obtained once.
+        folly::Random::DefaultGenerator rng;
+        for (auto i = 0; i < kNumGizmos / 3; ++i) {
+          auto gizmo =
+              gizmos[folly::Random::rand32(rng) % gizmos.size()]->move();
+          if (gizmo) {
+            results.withWLock([&](auto& set) {
+              EXPECT_TRUE(set.find(gizmo->id) == set.end());
+              set.insert(gizmo->id);
+            });
+          }
+        }
+        for (auto i = 0; i < gizmos.size(); ++i) {
+          auto gizmo = gizmos[i]->move();
+          if (gizmo) {
+            results.withWLock([&](auto& set) {
+              EXPECT_TRUE(set.find(gizmo->id) == set.end());
+              set.insert(gizmo->id);
+            });
+          }
+        }
+      }
+    }));
+  }
+  for (auto& thread : threads) {
+    thread.join();
+  }
+  results.withRLock([&](auto& set) {
+    for (auto i = 0; i < kNumGizmos; ++i) {
+      EXPECT_TRUE(set.find(i) != set.end());
+    }
+  });
+}

velox/common/tests/CMakeLists.txt

@@ -14,6 +14,7 @@
 
 add_executable(
   velox_common_test
+  AsyncSourceTest.cpp
   ExceptionTests.cpp
   RangeTest.cpp
   BitUtilTest.cpp