make ExecutorWithPriority's priority changeable

Summary:
We need the ability to dynamically change the priority of an executor when it's executing work. Specifically, we want to de-prioritize a query's work units while it is being executed if it has taken up too much resources.
Pass a callback function to ExecutorWithPriority which sets the priority upon every add.

Reviewed By: andriigrynenko

Differential Revision: D18344806

fbshipit-source-id: c0c59d3982637703e6e7c3793930b9d9fadce6fe

folly/executors/ExecutorWithPriority-inl.h

+/*
+ * 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 <glog/logging.h>
+
+namespace folly {
+namespace detail {
+template <typename Callback>
+class ExecutorWithPriorityImpl : public virtual Executor {
+ public:
+  static Executor::KeepAlive<ExecutorWithPriorityImpl<std::decay_t<Callback>>>
+  create(Executor::KeepAlive<Executor> executor, Callback&& callback) {
+    return makeKeepAlive(new ExecutorWithPriorityImpl<std::decay_t<Callback>>(
+        executor, std::move(callback)));
+  }
+  ExecutorWithPriorityImpl(ExecutorWithPriorityImpl const&) = delete;
+  ExecutorWithPriorityImpl& operator=(ExecutorWithPriorityImpl const&) = delete;
+  ExecutorWithPriorityImpl(ExecutorWithPriorityImpl&&) = delete;
+  ExecutorWithPriorityImpl& operator=(ExecutorWithPriorityImpl&&) = delete;
+
+  void add(Func func) override {
+    int8_t priority = callback_();
+    executor_->addWithPriority(std::move(func), priority);
+  }
+
+ protected:
+  bool keepAliveAcquire() override {
+    auto keepAliveCounter =
+        keepAliveCounter_.fetch_add(1, std::memory_order_relaxed);
+    DCHECK(keepAliveCounter > 0);
+    return true;
+  }
+
+  void keepAliveRelease() override {
+    auto keepAliveCounter =
+        keepAliveCounter_.fetch_sub(1, std::memory_order_acq_rel);
+    DCHECK(keepAliveCounter > 0);
+    if (keepAliveCounter == 1) {
+      delete this;
+    }
+  }
+
+ private:
+  ExecutorWithPriorityImpl(
+      Executor::KeepAlive<Executor> executor,
+      Callback&& callback)
+      : executor_(std::move(executor)), callback_(std::move(callback)) {}
+  std::atomic<ssize_t> keepAliveCounter_{1};
+  Executor::KeepAlive<Executor> executor_;
+  Callback callback_;
+};
+} // namespace detail
+
+template <typename Callback>
+Executor::KeepAlive<> ExecutorWithPriority::createDynamic(
+    Executor::KeepAlive<Executor> executor,
+    Callback&& callback) {
+  return detail::ExecutorWithPriorityImpl<std::decay_t<Callback>>::create(
+      executor, std::move(callback));
+}
+} // namespace folly

folly/executors/ExecutorWithPriority.cpp

@@ -15,33 +15,12 @@
  */
 
 #include <folly/executors/ExecutorWithPriority.h>
-#include <glog/logging.h>
 
 namespace folly {
-Executor::KeepAlive<ExecutorWithPriority> ExecutorWithPriority::create(
-    KeepAlive<Executor> executor,
+Executor::KeepAlive<> ExecutorWithPriority::create(
+    Executor::KeepAlive<Executor> executor,
     int8_t priority) {
-  return makeKeepAlive<ExecutorWithPriority>(
-      new ExecutorWithPriority(std::move(executor), priority));
-}
-
-void ExecutorWithPriority::add(Func func) {
-  executor_->addWithPriority(std::move(func), priority_);
-}
-
-bool ExecutorWithPriority::keepAliveAcquire() {
-  auto keepAliveCounter =
-      keepAliveCounter_.fetch_add(1, std::memory_order_relaxed);
-  DCHECK(keepAliveCounter > 0);
-  return true;
-}
-
-void ExecutorWithPriority::keepAliveRelease() {
-  auto keepAliveCounter =
-      keepAliveCounter_.fetch_sub(1, std::memory_order_acq_rel);
-  DCHECK(keepAliveCounter > 0);
-  if (keepAliveCounter == 1) {
-    delete this;
-  }
+  return ExecutorWithPriority::createDynamic(
+      executor, [priority]() { return priority; });
 }
 } // namespace folly

folly/executors/ExecutorWithPriority.h

@@ -20,30 +20,17 @@
 #include <atomic>
 
 namespace folly {
-
-class ExecutorWithPriority : public virtual Executor {
+class ExecutorWithPriority {
  public:
-  ExecutorWithPriority(ExecutorWithPriority const&) = delete;
-  ExecutorWithPriority& operator=(ExecutorWithPriority const&) = delete;
-  ExecutorWithPriority(ExecutorWithPriority&&) = delete;
-  ExecutorWithPriority& operator=(ExecutorWithPriority&&) = delete;
+  template <typename Callback>
+  static Executor::KeepAlive<> createDynamic(
+      Executor::KeepAlive<Executor> executor,
+      Callback&& callback);
 
-  static Executor::KeepAlive<ExecutorWithPriority> create(
-      KeepAlive<Executor> executor,
+  static Executor::KeepAlive<> create(
+      Executor::KeepAlive<Executor> executor,
       int8_t priority);
-
-  void add(Func func) override;
-
- protected:
-  bool keepAliveAcquire() override;
-  void keepAliveRelease() override;
-
- private:
-  ExecutorWithPriority(KeepAlive<Executor> executor, int8_t priority)
-      : executor_(std::move(executor)), priority_(priority) {}
-
-  std::atomic<ssize_t> keepAliveCounter_{1};
-  KeepAlive<Executor> executor_;
-  int8_t priority_;
 };
 } // namespace folly
+
+#include <folly/executors/ExecutorWithPriority-inl.h>

folly/executors/test/ExecutorWithPriorityTest.cpp

@@ -16,6 +16,7 @@
 
 #include <folly/executors/ExecutorWithPriority.h>
 #include <folly/executors/CPUThreadPoolExecutor.h>
+#include <folly/futures/Future.h>
 #include <folly/portability/GTest.h>
 
 using namespace folly;
@@ -49,3 +50,27 @@ TEST(ExecutorWithPriorityTest, addWithCorrectPriorityTest) {
   pool->join();
   EXPECT_EQ(100, completed);
 }
+
+TEST(ExecutorWithPriorityTest, updatePriorityUsingCallback) {
+  std::atomic<int8_t> curPri = -25;
+  auto completed = 0;
+  auto callback = [&curPri]() { return curPri += 1; };
+  // priorities range from -25 to 25
+  auto pool = std::make_shared<CPUThreadPoolExecutor>(
+      0 /*numThreads*/, 51 /*numPriorities*/);
+  {
+    auto executorWithPriority = ExecutorWithPriority::createDynamic(
+        getKeepAliveToken(pool.get()), callback);
+    for (int i = 0; i < 50; i++) {
+      auto task = [&completed, i]() {
+        completed++;
+        EXPECT_EQ(completed, 50 - i);
+      };
+      executorWithPriority->add(std::move(task));
+    }
+    pool->setNumThreads(1);
+  }
+  pool->join();
+  EXPECT_EQ(50, completed);
+  EXPECT_EQ(25, curPri);
+}