changing segment names to include objstoreid (#50)

protos/orchestra.proto

@@ -59,6 +59,7 @@ message RegisterWorkerRequest {
 
 message RegisterWorkerReply {
   uint64 workerid = 1; // Worker ID assigned by the scheduler
+  uint64 objstoreid = 2; // The Object store ID of the worker's local object store
 }
 
 message RegisterObjStoreRequest {

src/ipc.cc

@@ -26,7 +26,7 @@ int64_t BufferMemorySource::Size() const {
   return size_;
 }
 
-MemorySegmentPool::MemorySegmentPool(bool create) : create_mode_(create) { }
+MemorySegmentPool::MemorySegmentPool(ObjStoreId objstoreid, bool create) : objstoreid_(objstoreid), create_mode_(create) { }
 
 // creates a memory segment if it is not already there; if the pool is in create mode,
 // space is allocated, if it is in open mode, the shared memory is mapped into the process
@@ -49,7 +49,7 @@ void MemorySegmentPool::open_segment(SegmentId segmentid, size_t size) {
   if (segments_[segmentid].second == SegmentStatusType::CLOSED) {
     ORCH_LOG(ORCH_FATAL, "Attempting to open segmentid " << segmentid << ", but segments_[segmentid].second == SegmentStatusType::CLOSED.");
   }
-  std::string segment_name = std::string("segment:") + std::to_string(segmentid);
+  std::string segment_name = get_segment_name(segmentid);
   if (create_mode_) {
     assert(size > 0);
     shared_memory_object::remove(segment_name.c_str()); // remove segment if it has not been properly removed from last run
@@ -62,7 +62,7 @@ void MemorySegmentPool::open_segment(SegmentId segmentid, size_t size) {
 
 void MemorySegmentPool::close_segment(SegmentId segmentid) {
   ORCH_LOG(ORCH_DEBUG, "CLOSING segmentid " << segmentid);
-  std::string segment_name = std::string("segment:") + std::to_string(segmentid);
+  std::string segment_name = get_segment_name(segmentid);
   shared_memory_object::remove(segment_name.c_str());
   segments_[segmentid].first.reset();
   segments_[segmentid].second = SegmentStatusType::CLOSED;
@@ -95,9 +95,14 @@ uint8_t* MemorySegmentPool::get_address(ObjHandle pointer) {
   return static_cast<uint8_t*>(segment->get_address_from_handle(pointer.ipcpointer()));
 }
 
+// returns the name of the segment
+std::string MemorySegmentPool::get_segment_name(SegmentId segmentid) {
+  return std::string("objstore:") + std::to_string(objstoreid_) + std::string(":segment:") + std::to_string(segmentid);
+}
+
 MemorySegmentPool::~MemorySegmentPool() {
   for (size_t segmentid = 0; segmentid < segments_.size(); ++segmentid) {
-    std::string segment_name = std::string("segment:") + std::to_string(segmentid);
+    std::string segment_name = get_segment_name(segmentid);
     segments_[segmentid].first.reset();
     shared_memory_object::remove(segment_name.c_str());
   }

src/ipc.h

@@ -2,6 +2,7 @@
 #define ORCHESTRA_IPC_H
 
 #include <iostream>
+#include <limits>
 
 #include <boost/interprocess/managed_shared_memory.hpp>
 #include <boost/interprocess/ipc/message_queue.hpp>
@@ -145,15 +146,17 @@ enum SegmentStatusType {UNOPENED = 0, OPENED = 1, CLOSED = 2};
 
 class MemorySegmentPool {
 public:
-  MemorySegmentPool(bool create = false); // can be used in two modes: create mode and open mode (see above)
+  MemorySegmentPool(ObjStoreId objstoreid, bool create); // can be used in two modes: create mode and open mode (see above)
   ~MemorySegmentPool();
   ObjHandle allocate(size_t nbytes); // allocate memory, potentially creating a new segment (only run on object store)
   void deallocate(ObjHandle pointer); // deallocate object, potentially deallocating a new segment (only run on object store)
   uint8_t* get_address(ObjHandle pointer); // get address of shared object
+  std::string get_segment_name(SegmentId segmentid); // get the name of a segment
 private:
   void open_segment(SegmentId segmentid, size_t size = 0); // create a segment or map an existing one into memory
   void close_segment(SegmentId segmentid); // close a segment
   bool create_mode_; // true in the object stores, false on the workers
+  ObjStoreId objstoreid_; // the identity of the associated object store
   size_t page_size_ = mapped_region::get_page_size();
   std::vector<std::pair<std::unique_ptr<managed_shared_memory>, SegmentStatusType> > segments_;
 };

src/objstore.cc

@@ -23,7 +23,7 @@ Status ObjStoreClient::upload_data_to(slice data, ObjRef objref, ObjStore::Stub&
 }
 
 ObjStoreService::ObjStoreService(const std::string& objstore_address, std::shared_ptr<Channel> scheduler_channel)
-  : scheduler_stub_(Scheduler::NewStub(scheduler_channel)), segmentpool_(true), objstore_address_(objstore_address) {
+  : scheduler_stub_(Scheduler::NewStub(scheduler_channel)), objstore_address_(objstore_address) {
   recv_queue_.connect(std::string("queue:") + objstore_address + std::string(":obj"), true);
   ClientContext context;
   RegisterObjStoreRequest request;
@@ -31,6 +31,7 @@ ObjStoreService::ObjStoreService(const std::string& objstore_address, std::share
   RegisterObjStoreReply reply;
   scheduler_stub_->RegisterObjStore(&context, request, &reply);
   objstoreid_ = reply.objstoreid();
+  segmentpool_ = std::make_shared<MemorySegmentPool>(objstoreid_, true);
 }
 
 // this method needs to be protected by a objstores_lock_
@@ -153,7 +154,7 @@ Status ObjStoreService::DeallocateObject(ServerContext* context, const Deallocat
   if (canonical_objref >= memory_.size()) {
     ORCH_LOG(ORCH_FATAL, "Attempting to deallocate canonical_objref " << canonical_objref << ", but it is not in the objstore.");
   }
-  segmentpool_.deallocate(memory_[canonical_objref].first);
+  segmentpool_->deallocate(memory_[canonical_objref].first);
   memory_[canonical_objref].second = MemoryStatusType::DEALLOCATED;
   return Status::OK;
 }
@@ -197,7 +198,7 @@ void ObjStoreService::process_worker_request(const ObjRequest request) {
   switch (request.type) {
     case ObjRequestType::ALLOC: {
         // TODO(rkn): Does segmentpool_ need a lock around it?
-        ObjHandle reply = segmentpool_.allocate(request.size);
+        ObjHandle reply = segmentpool_->allocate(request.size);
         send_queues_[request.workerid].send(&reply);
         std::lock_guard<std::mutex> memory_lock(memory_lock_);
         if (memory_[request.objref].second != MemoryStatusType::NOT_PRESENT) {

src/objstore.h

@@ -50,7 +50,7 @@ private:
 
   std::string objstore_address_;
   ObjStoreId objstoreid_; // id of this objectstore in the scheduler object store table
-  MemorySegmentPool segmentpool_;
+  std::shared_ptr<MemorySegmentPool> segmentpool_;
   std::vector<std::pair<ObjHandle, MemoryStatusType> > memory_; // object reference -> (memory address, memory status)
   std::mutex memory_lock_;
   std::unordered_map<std::string, std::unique_ptr<ObjStore::Stub>> objstores_;

src/scheduler.cc

@@ -106,9 +106,12 @@ Status SchedulerService::RegisterObjStore(ServerContext* context, const Register
 }
 
 Status SchedulerService::RegisterWorker(ServerContext* context, const RegisterWorkerRequest* request, RegisterWorkerReply* reply) {
-  WorkerId workerid = register_worker(request->worker_address(), request->objstore_address());
+  std::pair<WorkerId, ObjStoreId> info = register_worker(request->worker_address(), request->objstore_address());
+  WorkerId workerid = info.first;
+  ObjStoreId objstoreid = info.second;
   ORCH_LOG(ORCH_INFO, "registered worker with workerid " << workerid);
   reply->set_workerid(workerid);
+  reply->set_objstoreid(objstoreid);
   schedule();
   return Status::OK;
 }
@@ -260,7 +263,7 @@ bool SchedulerService::can_run(const Call& task) {
   return true;
 }
 
-WorkerId SchedulerService::register_worker(const std::string& worker_address, const std::string& objstore_address) {
+std::pair<WorkerId, ObjStoreId> SchedulerService::register_worker(const std::string& worker_address, const std::string& objstore_address) {
   ORCH_LOG(ORCH_INFO, "registering worker " << worker_address << " connected to object store " << objstore_address);
   ObjStoreId objstoreid = std::numeric_limits<size_t>::max();
   for (int num_attempts = 0; num_attempts < 5; ++num_attempts) {
@@ -288,7 +291,7 @@ WorkerId SchedulerService::register_worker(const std::string& worker_address, co
   avail_workers_lock_.lock();
   avail_workers_.push_back(workerid);
   avail_workers_lock_.unlock();
-  return workerid;
+  return std::make_pair(workerid, objstoreid);
 }
 
 ObjRef SchedulerService::register_new_object() {

src/scheduler.h

@@ -66,7 +66,7 @@ public:
   // checks if the dependencies of the task are met
   bool can_run(const Call& task);
   // register a worker and its object store (if it has not been registered yet)
-  WorkerId register_worker(const std::string& worker_address, const std::string& objstore_address);
+  std::pair<WorkerId, ObjStoreId> register_worker(const std::string& worker_address, const std::string& objstore_address);
   // register a new object with the scheduler and return its object reference
   ObjRef register_new_object();
   // register the location of the object reference in the object table

src/worker.cc

@@ -34,6 +34,8 @@ void Worker::register_worker(const std::string& worker_address, const std::strin
   ClientContext context;
   Status status = scheduler_stub_->RegisterWorker(&context, request, &reply);
   workerid_ = reply.workerid();
+  objstoreid_ = reply.objstoreid();
+  segmentpool_ = std::make_shared<MemorySegmentPool>(objstoreid_, false);
   request_obj_queue_.connect(std::string("queue:") + objstore_address + std::string(":obj"), false);
   std::string queue_name = std::string("queue:") + objstore_address + std::string(":worker:") + std::to_string(workerid_) + std::string(":obj");
   receive_obj_queue_.connect(queue_name, true);
@@ -78,7 +80,7 @@ slice Worker::get_object(ObjRef objref) {
   ObjHandle result;
   receive_obj_queue_.receive(&result);
   slice slice;
-  slice.data = segmentpool_.get_address(result);
+  slice.data = segmentpool_->get_address(result);
   slice.len = result.size();
   return slice;
 }
@@ -106,7 +108,7 @@ void Worker::put_object(ObjRef objref, const Obj* obj, std::vector<ObjRef> &cont
   }
   ObjHandle result;
   receive_obj_queue_.receive(&result);
-  uint8_t* target = segmentpool_.get_address(result);
+  uint8_t* target = segmentpool_->get_address(result);
   std::memcpy(target, &data[0], data.size());
   request.type = ObjRequestType::WORKER_DONE;
   request.metadata_offset = 0;
@@ -136,7 +138,7 @@ void Worker::put_arrow(ObjRef objref, PyArrayObject* array) {
   request_obj_queue_.send(&request);
   ObjHandle result;
   receive_obj_queue_.receive(&result);
-  store_arrow(array, result, &segmentpool_);
+  store_arrow(array, result, segmentpool_.get());
   request.type = ObjRequestType::WORKER_DONE;
   request.metadata_offset = result.metadata_offset();
   request_obj_queue_.send(&request);
@@ -153,7 +155,7 @@ PyArrayObject* Worker::get_arrow(ObjRef objref) {
   request_obj_queue_.send(&request);
   ObjHandle result;
   receive_obj_queue_.receive(&result);
-  return (PyArrayObject*)deserialize_array(result, &segmentpool_);
+  return (PyArrayObject*)deserialize_array(result, segmentpool_.get());
 }
 
 bool Worker::is_arrow(ObjRef objref) {

src/worker.h

@@ -89,10 +89,11 @@ class Worker {
   MessageQueue<Call*> receive_queue_;
   managed_shared_memory segment_;
   WorkerId workerid_;
+  ObjStoreId objstoreid_;
   std::string worker_address_;
   MessageQueue<ObjRequest> request_obj_queue_;
   MessageQueue<ObjHandle> receive_obj_queue_;
-  MemorySegmentPool segmentpool_;
+  std::shared_ptr<MemorySegmentPool> segmentpool_;
 };
 
 #endif