Merge pull request #1221 from tli2/master

Clean up Brain framework code

Author: pervazea

src/include/brain/brain.h

@@ -15,52 +15,88 @@
 #include <string>
 #include <cstdio>
 #include <utility>
+#include <functional>
 #include "capnp/ez-rpc.h"
 #include "peloton/capnp/peloton_service.capnp.h"
 #include "common/notifiable_task.h"
 
 namespace peloton {
 namespace brain {
 
+/**
+ * Provides an access point to the various resources available to the jobs in
+ * the brain, such as RPC and Catalog.
+ */
 class BrainEnvironment {
-  // TODO(tianyu): provide interface for accessing various resources of the brain,
-  // such as network connection to a peloton engine.
+  // TODO(tianyu): fill in as needed
 };
 
+/**
+ * Interface that represents a piece of task to be run on Brain. To use this
+ * abstract class, extend it with a concrete class and fill in the method
+ * OnJobInvocation.
+ */
 class BrainJob {
  public:
-  // TODO(tianyu): Extend this interface for richer interaction
-  virtual void RunJob(BrainEnvironment *) = 0;
+  explicit BrainJob(BrainEnvironment *env) : env_(env) {}
+
+  virtual ~BrainJob() = default;
+
+  // This is separate from the user-defined OnJobInvocation to allow for better
+  // interfacing with the libevent API.
+  /**
+   * Invokes this job to be run. The brain framework will call this method.
+   *
+   */
+  inline void Invoke() { OnJobInvocation(env_); }
+  // TODO(tianyu): Extend this interface for richer behavior
+  /**
+   * Executed as the main body of the job, filled in by the user. Use the
+   * provided BrainEnvironment for interaction with Brain's resources.
+   */
+  virtual void OnJobInvocation(BrainEnvironment *) = 0;
+ private:
+  BrainEnvironment *env_;
 };
 
-class ExampleBrainJob: public BrainJob {
+/**
+ * Simple implementation of a BrainJob.
+ */
+class SimpleBrainJob : public BrainJob {
  public:
-   void RunJob(BrainEnvironment *) override {
-     // TODO(tianyu): Replace with real address
-     capnp::EzRpcClient client("localhost:15445");
-     PelotonService::Client peloton_service = client.getMain<PelotonService>();
-     auto request = peloton_service.createIndexRequest();
-     request.getRequest().setIndexKeys(42);
-     auto response = request.send().wait(client.getWaitScope());
-   }
+  explicit SimpleBrainJob(BrainEnvironment *env,
+                          std::function<void(BrainEnvironment *)> task)
+      : BrainJob(env), task_(std::move(task)) {}
+  inline void OnJobInvocation(BrainEnvironment *env) override { task_(env); }
+ private:
+  std::function<void(BrainEnvironment *)> task_;
 };
 
+/**
+ * Main running component of the brain. Events can be registered on this event
+ * loop and once Run is called, it will invoke handlers every specified time
+ * interval
+ */
 class Brain {
  public:
   // TODO(tianyu): Add necessary parameters to initialize the brain's resources
   Brain() : scheduler_(0) {}
 
+  ~Brain() {
+    for (auto entry : jobs_)
+      delete entry.second;
+  }
+
+  template <typename BrainJob, typename... Args>
   inline void RegisterJob(const struct timeval *period,
-                          std::string name,
-                          BrainJob &job) {
-    auto callback = [](int, short, void *pair) {
-      auto *job_env_pair = reinterpret_cast<std::pair<BrainJob *, BrainEnvironment *> *>(pair);
-      job_env_pair->first->RunJob(job_env_pair->second);
+                          std::string name, Args... args) {
+    auto *job = new BrainJob(&env_, args...);
+    jobs_[name] = job;
+    auto callback = [](int, short, void *arg) {
+      reinterpret_cast<BrainJob *>(arg)->Invoke();
     };
-    // TODO(tianyu) Deal with this memory
-    auto *pair = new std::pair<BrainJob *, BrainEnvironment *>(&job, &env_);
     job_handles_[name] =
-        scheduler_.RegisterPeriodicEvent(period, callback, pair);
+        scheduler_.RegisterPeriodicEvent(period, callback, job);
   }
 
   inline void Run() {
@@ -73,8 +109,9 @@ class Brain {
 
  private:
   NotifiableTask scheduler_;
+  std::unordered_map<std::string, BrainJob *> jobs_;
   std::unordered_map<std::string, struct event *> job_handles_;
   BrainEnvironment env_;
 };
-}
-}
+} // namespace brain
+} // namespace peloton

src/include/common/event_util.h

@@ -21,6 +21,7 @@
 #include "common/logger.h"
 
 namespace peloton {
+
 /**
  * Static utility class with wrappers for libevent functions.
  *

src/include/common/notifiable_task.h

@@ -62,6 +62,7 @@ class NotifiableTask {
    */
   inline int Id() const { return task_id_; }
 
+
   /**
    * @brief Register an event with the event base associated with this
    * notifiable task.

src/main/peloton/peloton.cpp

@@ -22,7 +22,7 @@
 // For GFlag's built-in help message flag
 DECLARE_bool(help);
 
-void RunPelotonEngine() {
+void RunPelotonServer() {
   try {
     // Setup
     peloton::PelotonInit::Initialize();
@@ -49,8 +49,18 @@ void RunPelotonBrain() {
   // TODO(tianyu): register jobs here
   struct timeval one_second;
   one_second.tv_sec = 1;
-  peloton::brain::ExampleBrainJob job;
-  brain.RegisterJob(&one_second, "test", job);
+  one_second.tv_usec = 0;
+
+  auto example_task = [](peloton::brain::BrainEnvironment *) {
+    // TODO(tianyu): Replace with real address
+    capnp::EzRpcClient client("localhost:15445");
+    PelotonService::Client peloton_service = client.getMain<PelotonService>();
+    auto request = peloton_service.createIndexRequest();
+    request.getRequest().setIndexKeys(42);
+    auto response = request.send().wait(client.getWaitScope());
+  };
+
+  brain.RegisterJob<peloton::brain::SimpleBrainJob>(&one_second, "test", example_task);
   brain.Run();
 }
 
@@ -76,6 +86,6 @@ int main(int argc, char *argv[]) {
       peloton::settings::SettingId::brain))
     RunPelotonBrain();
   else
-    RunPelotonEngine();
+    RunPelotonServer();
   return 0;
 }

src/network/peloton_server.cpp

@@ -270,8 +270,8 @@ void PelotonServer::ServerLoop() {
   if (settings::SettingsManager::GetBool(settings::SettingId::rpc_enabled)) {
     int rpc_port =
         settings::SettingsManager::GetInt(settings::SettingId::rpc_port);
-    auto rpc_task = std::make_shared<PelotonRpcHandlerTask>(("127.0.0.1:"
-        + std::to_string(rpc_port)).c_str());
+    std::string address = "127.0.0.1:" + std::to_string(rpc_port);
+    auto rpc_task = std::make_shared<PelotonRpcHandlerTask>(address.c_str());
     DedicatedThreadRegistry::GetInstance()
         .RegisterDedicatedThread<PelotonRpcHandlerTask>(this, rpc_task);
   }