getting the object store working

src/objstore.cc

@@ -0,0 +1,59 @@
+#include "objstore.h"
+
+const size_t ObjStoreClient::CHUNK_SIZE = 8 * 1024;
+
+Status ObjStoreClient::upload_data_to(slice data, ObjRef objref, ObjStore::Stub& stub) {
+  ObjChunk chunk;
+  ClientContext context;
+  AckReply reply;
+  std::unique_ptr<ClientWriter<ObjChunk> > writer(stub.StreamObj(&context, &reply));
+  const char* head = data.data;
+  for (size_t i = 0; i < data.len; i += CHUNK_SIZE) {
+    chunk.set_objref(objref);
+    chunk.set_totalsize(data.len);
+    chunk.set_data(head + i, std::min(CHUNK_SIZE, data.len - i));
+    if (!writer->Write(chunk)) {
+      std::cout << "write failed" << std::endl;
+      // throw std::runtime_error("write failed");
+    }
+  }
+  writer->WritesDone();
+  return writer->Finish();
+}
+
+void ObjStoreServiceImpl::allocate_memory(ObjRef objref, size_t size) {
+  std::ostringstream stream;
+  stream << "obj-" << memory_names_.size();
+  std::string name = stream.str();
+  // Make sure that the name is not taken yet
+  shared_memory_object::remove(name.c_str());
+  memory_names_.push_back(name);
+  // Make room for boost::interprocess metadata
+  size_t new_size = (size / page_size + 2) * page_size;
+  shared_object& object = memory_[objref];
+  object.name = name;
+  object.memory = std::make_shared<managed_shared_memory>(create_only, name.c_str(), new_size);
+  object.ptr.data = static_cast<char*>(memory_[objref].memory->allocate(size));
+  object.ptr.len = size;
+}
+
+void start_objstore(const char* objstore_address) {
+  ObjStoreServiceImpl service;
+  ServerBuilder builder;
+
+  builder.AddListeningPort(std::string(objstore_address), grpc::InsecureServerCredentials());
+  builder.RegisterService(&service);
+  std::unique_ptr<Server> server(builder.BuildAndStart());
+
+  server->Wait();
+}
+
+int main(int argc, char** argv) {
+  if (argc != 2) {
+    return 1;
+  }
+
+  start_objstore(argv[1]);
+
+  return 0;
+}

src/objstore.h

@@ -0,0 +1,133 @@
+#ifndef ORCHESTRA_OBJSTORE_SERVER_H
+#define ORCHESTRA_OBJSTORE_SERVER_H
+
+#include <unordered_map>
+#include <memory>
+#include <iostream>
+#include <boost/interprocess/managed_shared_memory.hpp>
+#include <grpc++/grpc++.h>
+
+using namespace boost::interprocess;
+
+#include "orchestra/orchestra.h"
+#include "orchestra.grpc.pb.h"
+#include "types.pb.h"
+
+#include "orchlib.h"
+
+using grpc::Server;
+using grpc::ServerBuilder;
+using grpc::ServerReader;
+using grpc::ServerContext;
+using grpc::ClientContext;
+using grpc::ClientWriter;
+using grpc::Status;
+
+using grpc::Channel;
+
+class ObjStoreClient {
+public:
+  static const size_t CHUNK_SIZE;
+  static Status upload_data_to(slice data, ObjRef objref, ObjStore::Stub& stub);
+};
+
+struct shared_object {
+  std::string name;
+  std::shared_ptr<managed_shared_memory> memory;
+  slice ptr;
+};
+
+class ObjStoreServiceImpl final : public ObjStore::Service {
+  std::vector<std::string> memory_names_;
+  std::unordered_map<ObjRef, shared_object> memory_;
+  std::mutex memory_lock_;
+  size_t page_size = mapped_region::get_page_size();
+  std::unordered_map<std::string, std::unique_ptr<ObjStore::Stub>> objstores_;
+
+  void allocate_memory(ObjRef objref, size_t size);
+
+  // check if we already connected to the other objstore, if yes, return reference to connection, otherwise connect
+  ObjStore::Stub& get_objstore_stub(const std::string& objstore_address) {
+    auto iter = objstores_.find(objstore_address);
+    if (iter != objstores_.end())
+      return *(iter->second);
+    auto channel = grpc::CreateChannel(objstore_address, grpc::InsecureChannelCredentials());
+    objstores_.emplace(objstore_address, ObjStore::NewStub(channel));
+    return *objstores_[objstore_address];
+  }
+
+public:
+  ObjStoreServiceImpl() {}
+
+  ~ObjStoreServiceImpl() {
+    for (const auto& segment_name : memory_names_) {
+      shared_memory_object::remove(segment_name.c_str());
+    }
+  }
+
+  Status DeliverObj(ServerContext* context, const DeliverObjRequest* request, AckReply* reply) override {
+    ObjStore::Stub& stub = get_objstore_stub(request->objstore_address());
+    ObjRef objref = request->objref();
+
+    // TODO: Have to introduce wait condition
+
+    return ObjStoreClient::upload_data_to(memory_[objref].ptr, objref, stub);
+  }
+
+  Status DebugInfo(ServerContext* context, const DebugInfoRequest* request, DebugInfoReply* reply) override {
+    for (const auto& entry : memory_) {
+      reply->add_objref(entry.first);
+    }
+    return Status::OK;
+  }
+
+  Status GetObj(ServerContext* context, const GetObjRequest* request, GetObjReply* reply) override {
+    ObjRef objref = request->objref();
+    std::cout << "getobj lock";
+    memory_lock_.lock();
+    shared_object& object = memory_[objref];
+    reply->set_bucket(object.name);
+    auto handle = object.memory->get_handle_from_address(object.ptr.data);
+    reply->set_handle(handle);
+    reply->set_size(object.ptr.len);
+    memory_lock_.unlock();
+    std::cout << "getobj unlock";
+    return Status::OK;
+  }
+
+  Status StreamObj(ServerContext* context, ServerReader<ObjChunk>* reader, AckReply* reply) override {
+    std::cout << "stream obj lock" << std::endl;
+    memory_lock_.lock();
+    ObjChunk chunk;
+    ObjRef objref = 0;
+    size_t totalsize = 0;
+    if (reader->Read(&chunk)) {
+      objref = chunk.objref();
+      totalsize = chunk.totalsize();
+      allocate_memory(objref, totalsize);
+    }
+    size_t num_bytes = 0;
+    char* data = memory_[objref].ptr.data;
+
+    std::cout << "before loop " << totalsize << std::endl;
+
+    do {
+      if (num_bytes + chunk.data().size() > totalsize) {
+        std::cout << "cancelled" << std::endl;
+        memory_lock_.unlock();
+        return Status::CANCELLED;
+      }
+      std::memcpy(data, chunk.data().c_str(), chunk.data().size());
+      data += chunk.data().size();
+      num_bytes += chunk.data().size();
+      std::cout << "looping " << num_bytes << std::endl;
+    } while (reader->Read(&chunk));
+
+    std::cout << "finished" << std::endl;
+    memory_lock_.unlock();
+    std::cout << "stream obj unlock" << std::endl;
+    return Status::OK;
+  }
+};
+
+#endif

src/orchlib.cc

@@ -0,0 +1,29 @@
+#include "worker.h"
+
+Worker* orch_create_context(const char* server_addr, const char* worker_addr, const char* objstore_addr) {
+  auto server_channel = grpc::CreateChannel(server_addr, grpc::InsecureChannelCredentials());
+  auto objstore_channel = grpc::CreateChannel(objstore_addr, grpc::InsecureChannelCredentials());
+  Worker* worker = new Worker(server_channel, objstore_channel);
+  worker->register_worker(std::string(worker_addr), std::string(objstore_addr));
+  return worker;
+}
+
+size_t orch_remote_call(Worker* worker, RemoteCallRequest* request) {
+  return worker->RemoteCall(request);
+}
+
+void orch_main_loop(Worker* worker) {
+  worker->MainLoop();
+}
+
+size_t orch_push(Worker* worker, Obj* obj) {
+  return worker->PushObj(obj);
+}
+
+slice orch_get_serialized_obj(Worker* worker, ObjRef objref) {
+  return worker->GetSerializedObj(objref);
+}
+
+void orch_register_function(Worker* worker, const char* name, size_t num_return_vals) {
+  // worker->register_function(std::string(name), num_return_vals);
+}

src/orchlib.h

@@ -0,0 +1,20 @@
+
+
+extern "C" {
+
+struct slice {
+  char* data;
+  size_t len;
+};
+
+struct Worker;
+struct RemoteCallRequest;
+struct Value;
+
+Worker* orch_create_context(const char* server_addr, const char* worker_addr, const char* objstore_addr);
+size_t orch_remote_call(Worker* context, RemoteCallRequest* request);
+size_t orch_push(Worker* context, Obj* value);
+void orch_main_loop(Worker* worker);
+slice orch_get_serialized_obj(Worker* worker, size_t objref);
+void orch_register_function(Worker* worker, const char* name, size_t num_return_vals);
+}

src/scheduler.h

@@ -0,0 +1,145 @@
+#ifndef ORCHESTRA_SCHEDULER_H
+#define ORCHESTRA_SCHEDULER_H
+
+#include <deque>
+
+#include <grpc++/grpc++.h>
+
+#include "orchestra/orchestra.h"
+#include "orchestra.grpc.pb.h"
+#include "types.pb.h"
+
+using grpc::Server;
+using grpc::ServerBuilder;
+using grpc::ServerReader;
+using grpc::ServerContext;
+using grpc::Status;
+
+using grpc::Channel;
+
+struct WorkerHandle {
+  std::shared_ptr<Channel> channel;
+  ObjStoreId objstoreid;
+};
+
+struct ObjStoreHandle {
+  std::shared_ptr<Channel> channel;
+  std::string address;
+};
+
+class Scheduler {
+  // Vector of all workers registered in the system. Their index in this vector
+  // is the workerid.
+  std::vector<WorkerHandle> workers_;
+  std::mutex workers_lock_;
+  // Vector of all workers that are currently idle.
+  std::vector<WorkerId> available_workers_;
+  // Vector of all object stores registered in the system. Their index in this
+  // vector is the objstoreid.
+  std::vector<ObjStoreHandle> objstores_;
+  grpc::mutex objstores_lock_;
+  // Mapping from objref to list of object stores where the object is stored.
+  ObjTable objtable_;
+  std::mutex objtable_lock_;
+  // Hash map from function names to workers where the function is registered.
+  FnTable fntable_;
+  std::mutex fntable_lock_;
+  // List of pending tasks.
+  std::deque<std::unique_ptr<Call> > tasks_;
+  std::mutex tasks_lock_;
+public:
+  // returns number of return values of task
+  size_t add_task(const Call& task) {
+    fntable_lock_.lock();
+    size_t num_return_vals = 2; // fn_table_[task.name()].num_return_vals();
+    fntable_lock_.unlock();
+    // std::unique_ptr<Call> task_ptr(new Call(task)); // TODO: perform copy outside
+    tasks_lock_.lock();
+    // tasks_.push_back(task_ptr);
+    tasks_lock_.unlock();
+    return num_return_vals;
+  }
+  WorkerId register_worker(const std::string& worker_address, const std::string& objstore_address) {
+    ObjStoreId objstoreid = std::numeric_limits<size_t>::max();
+    objstores_lock_.lock();
+    for (size_t i = 0; i < objstores_.size(); ++i) {
+      std::cout << "adress: " << objstores_[i].address << std::endl;
+      std::cout << "my adress: " << objstore_address << std::endl;
+      if (objstores_[i].address == objstore_address) {
+        objstoreid = i;
+      }
+    }
+    if (objstoreid == std::numeric_limits<size_t>::max()) {
+      // throw objstore_not_registered_error("objectstore not registered");
+      std::cout << "bad bad bad" << std::endl;
+    }
+    objstores_lock_.unlock();
+    workers_lock_.lock();
+    WorkerId result = workers_.size();
+    workers_.push_back(WorkerHandle());
+    workers_[result].channel = grpc::CreateChannel(worker_address, grpc::InsecureChannelCredentials());
+    workers_[result].objstoreid = objstoreid;
+    workers_lock_.unlock();
+    return result;
+  }
+  ObjStoreId register_objstore(const std::string& objstore_address) {
+    // auto handle = ObjStoreHandle(objstore_address);
+    // auto handlecopy = handle;
+    // auto handle = ObjStoreHandle("0.0.0.0:22222");
+    objstores_lock_.lock();
+    std::cout << "capacity" << objstores_.capacity() << std::endl;
+    ObjStoreId result = objstores_.size();
+    // auto handle = ObjStoreHandle(objstore_address);
+    // objstores_.emplace_back(objstore_address);
+    objstores_.push_back(ObjStoreHandle());
+
+    objstores_[result].channel = grpc::CreateChannel(objstore_address, grpc::InsecureChannelCredentials());
+    objstores_[result].address = std::string(objstore_address);
+
+    // auto handlecopy = handle;
+    // auto handle = grpc::CreateChannel(objstore_address, grpc::InsecureChannelCredentials());
+    // auto handlecopy = grpc::CreateChannel(objstore_address, grpc::InsecureChannelCredentials());
+    objstores_lock_.unlock();
+    return result;
+  }
+  ObjRef register_new_object() {
+    objtable_lock_.lock();
+    ObjRef result = objtable_.size();
+    objtable_.push_back(std::vector<ObjStoreId>());
+    objtable_lock_.unlock();
+    return result;
+  }
+  void add_objstore_to_obj(ObjRef objref, ObjStoreId objstoreid) {
+    objtable_lock_.lock();
+    // do a binary search
+    auto pos = std::lower_bound(objtable_[objref].begin(), objtable_[objref].end(), objstoreid);
+    if (pos == objtable_[objref].end() || objstoreid < *pos) {
+      objtable_[objref].insert(pos, objstoreid);
+    }
+    objtable_lock_.unlock();
+  }
+  ObjStoreId get_store(WorkerId workerid) {
+    workers_lock_.lock();
+    ObjStoreId result = workers_[workerid].objstoreid;
+    workers_lock_.unlock();
+    return result;
+  }
+  void register_function(const std::string& name, WorkerId workerid, size_t num_return_vals) {
+    fntable_lock_.lock();
+    FnInfo& info = fntable_[name];
+    info.set_num_return_vals(num_return_vals);
+    info.add_worker(workerid);
+    fntable_lock_.unlock();
+  }
+  /*
+  void debug_info(DebugInfoReply* debug_info) {
+    fntable_lock_.lock();
+    for (const auto& entry : fntable_) {
+      debug_info->
+    }
+    fntable_lock_.lock();
+  }
+  */
+};
+
+#endif

src/scheduler_server.cc

@@ -0,0 +1,50 @@
+#include "scheduler_server.h"
+
+Status SchedulerServerServiceImpl::RemoteCall(ServerContext* context, const RemoteCallRequest* request, RemoteCallReply* reply) {
+  size_t num_return_vals = scheduler_->add_task(request->call());
+  for (size_t i = 0; i < num_return_vals; ++i) {
+    ObjRef result = scheduler_->register_new_object();
+    reply->add_result(result);
+  }
+  return Status::OK;
+}
+
+Status SchedulerServerServiceImpl::PushObj(ServerContext* context, const PushObjRequest* request, PushObjReply* reply) {
+  ObjRef objref = scheduler_->register_new_object();
+  ObjStoreId objstoreid = scheduler_->get_store(request->workerid());
+  scheduler_->add_objstore_to_obj(objref, objstoreid);
+  reply->set_objref(objref);
+  return Status::OK;
+}
+
+  /*
+  Status PushObj(ServerContext* context, ServerReader<ObjChunk> *reader, AckReply* reply) override {
+    ObjChunk chunk;
+    while (reader->Read(&chunk)) {
+
+    }
+    std::cout << "got chunks" << std::endl;
+    return Status::OK;
+  }
+  */
+
+void start_scheduler_server(const char* server_address) {
+  SchedulerServerServiceImpl service;
+  ServerBuilder builder;
+
+  builder.AddListeningPort(std::string(server_address), grpc::InsecureServerCredentials());
+  builder.RegisterService(&service);
+  std::unique_ptr<Server> server(builder.BuildAndStart());
+
+  server->Wait();
+}
+
+int main(int argc, char** argv) {
+  if (argc != 2) {
+    return 1;
+  }
+
+  start_scheduler_server(argv[1]);
+
+  return 0;
+}

src/scheduler_server.h

@@ -0,0 +1,45 @@
+#ifndef ORCHESTRA_SCHEDULER_SERVER_H
+#define ORCHESTRA_SCHEDULER_SERVER_H
+
+#include <iostream>
+#include <memory>
+#include <string>
+#include <mutex>
+
+#include "scheduler.h"
+
+
+class SchedulerServerServiceImpl final : public SchedulerServer::Service {
+  ObjTable objtable_;
+  std::unique_ptr<Scheduler> scheduler_;
+public:
+  SchedulerServerServiceImpl() : scheduler_(new Scheduler()) {
+  }
+  Status RemoteCall(ServerContext* context, const RemoteCallRequest* request, RemoteCallReply* reply) override;
+  Status PushObj(ServerContext* context, const PushObjRequest* request, PushObjReply* reply) override;
+  Status PullObj(ServerContext* context, const PullObjRequest* request, AckReply* reply) override {
+    return Status::OK;
+  }
+  Status RegisterWorker(ServerContext* context, const RegisterWorkerRequest* request, RegisterWorkerReply* reply) override {
+    WorkerId workerid = scheduler_->register_worker(request->worker_address(), request->objstore_address());
+    reply->set_workerid(workerid);
+    return Status::OK;
+  }
+  Status RegisterObjStore(ServerContext* context, const RegisterObjStoreRequest* request, RegisterObjStoreReply* reply) override {
+    try {
+      reply->set_objstoreid(scheduler_->register_objstore(request->address()));
+    } catch (...) {
+      std::cout << "caught exception" << std::endl;
+    }
+    return Status::OK;
+  }
+  Status RegisterFunction(ServerContext* context, const RegisterFunctionRequest* request, AckReply* reply) override {
+    scheduler_->register_function(request->fnname(), request->workerid(), request->num_return_vals());
+    return Status::OK;
+  }
+  Status GetDebugInfo(ServerContext* context, const GetDebugInfoRequest* request, GetDebugInfoReply* reply) override {
+    return Status::OK;
+  }
+};
+
+#endif

src/server.cc

@@ -1,88 +0,0 @@
-#include <iostream>
-#include <memory>
-#include <string>
-#include <mutex>
-
-#include <grpc++/grpc++.h>
-
-#include "orchestra.grpc.pb.h"
-
-using grpc::Server;
-using grpc::ServerBuilder;
-using grpc::ServerContext;
-using grpc::Status;
-
-typedef size_t ObjRef;
-typedef size_t WorkerId;
-typedef std::vector<std::vector<WorkerId> > ObjTable;
-
-class OrchestraScheduler {
-
-};
-
-class OrchestraServer {
-  ObjTable objtable;
-  std::mutex mutex;
-public:
-  ObjRef register_new_object() {
-    mutex.lock();
-    ObjRef result = objtable.size();
-    // std::cout << "size " << result << std::endl;
-    objtable.push_back(std::vector<WorkerId>());
-    mutex.unlock();
-    return result;
-  }
-  void register_object(ObjRef objref, WorkerId workerid) {
-    mutex.lock();
-    objtable[objref].push_back(workerid);
-    mutex.unlock();
-  }
-};
-
-// Logic and data behind the server's behavior.
-class OrchestraServiceImpl final : public Orchestra::Service {
-  ObjTable objtable;
-  std::unique_ptr<OrchestraServer> server;
-public:
-  OrchestraServiceImpl() : server(new OrchestraServer()) {
-  }
-  Status RemoteCall(ServerContext* context, const RemoteCallRequest* request,
-                  RemoteCallReply* reply) override {
-    // std::cout << "called" << std::endl;
-    ObjRef objref = server->register_new_object();
-    reply->set_result(objref);
-    // std::string prefix("Hello ");
-    // reply->set_message(prefix + request->name());
-    return Status::OK;
-  }
-  Status RegisterWorker(ServerContext* context, const RegisterWorkerRequest* request,
-    RegisterWorkerReply* reply) override {
-      std::cout << "register worker" << std::endl;
-      return Status::OK;
-  }
-};
-
-void RunServer() {
-  std::string server_address("0.0.0.0:50052");
-  OrchestraServiceImpl service;
-
-  ServerBuilder builder;
-  // Listen on the given address without any authentication mechanism.
-  builder.AddListeningPort(server_address, grpc::InsecureServerCredentials());
-  // Register "service" as the instance through which we'll communicate with
-  // clients. In this case it corresponds to an *synchronous* service.
-  builder.RegisterService(&service);
-  // Finally assemble the server.
-  std::unique_ptr<Server> server(builder.BuildAndStart());
-  std::cout << "Server listening on " << server_address << std::endl;
-
-  // Wait for the server to shutdown. Note that some other thread must be
-  // responsible for shutting down the server for this call to ever return.
-  server->Wait();
-}
-
-int main(int argc, char** argv) {
-  RunServer();
-
-  return 0;
-}