From 6b846d2bc08e6e42894b6bdbcf479a0d33764ac9 Mon Sep 17 00:00:00 2001 From: Robert Nishihara Date: Sun, 24 Apr 2016 19:06:14 -0700 Subject: [PATCH 1/2] support multiple object stores, part one --- lib/orchpy/orchpy/worker.py | 20 +++++++++---------- protos/orchestra.proto | 24 +++++++++++------------ src/ipc.cc | 11 ++++++++--- src/objstore.h | 38 ++++++++++++++++++++++++------------ src/scheduler.cc | 31 ++++++++++++++++++----------- src/serialize.cc | 2 +- test/arrays_test.py | 2 +- test/runtest.py | 39 ++++++++++++++++++++++++++++++++----- 8 files changed, 111 insertions(+), 56 deletions(-) diff --git a/lib/orchpy/orchpy/worker.py b/lib/orchpy/orchpy/worker.py index 79d0be25d..09709158b 100644 --- a/lib/orchpy/orchpy/worker.py +++ b/lib/orchpy/orchpy/worker.py @@ -14,11 +14,11 @@ class Worker(object): def put_object(self, objref, value): """Put `value` in the local object store with objref `objref`. This assumes that the value for `objref` has not yet been placed in the local object store.""" - if type(value) == np.ndarray: - orchpy.lib.put_arrow(self.handle, objref, value) - else: - object_capsule, contained_objrefs = serialization.serialize(self.handle, value) # contained_objrefs is a list of the objrefs contained in object_capsule - orchpy.lib.put_object(self.handle, objref, object_capsule, contained_objrefs) + # if type(value) == np.ndarray: + # orchpy.lib.put_arrow(self.handle, objref, value) + # else: + object_capsule, contained_objrefs = serialization.serialize(self.handle, value) # contained_objrefs is a list of the objrefs contained in object_capsule + orchpy.lib.put_object(self.handle, objref, object_capsule, contained_objrefs) def get_object(self, objref): """ @@ -27,11 +27,11 @@ class Worker(object): WARNING: get_object can only be called on a canonical objref. """ - if orchpy.lib.is_arrow(self.handle, objref): - return orchpy.lib.get_arrow(self.handle, objref) - else: - object_capsule = orchpy.lib.get_object(self.handle, objref) - return serialization.deserialize(self.handle, object_capsule) + # if orchpy.lib.is_arrow(self.handle, objref): + # return orchpy.lib.get_arrow(self.handle, objref) + # else: + object_capsule = orchpy.lib.get_object(self.handle, objref) + return serialization.deserialize(self.handle, object_capsule) def alias_objrefs(self, alias_objref, target_objref): """Make `alias_objref` refer to the same object that `target_objref` refers to.""" diff --git a/protos/orchestra.proto b/protos/orchestra.proto index 88b5c9f67..9c63f8ad1 100644 --- a/protos/orchestra.proto +++ b/protos/orchestra.proto @@ -149,10 +149,10 @@ message SchedulerInfoReply { // Object stores service ObjStore { - // Request to deliver an object to another object store (called by the scheduler) - rpc DeliverObj(DeliverObjRequest) returns (AckReply); + // Tell the object store to begin pulling an object from another object store (called by the scheduler) + rpc StartDelivery(StartDeliveryRequest) returns (AckReply); // Accept incoming data from another object store, as a stream of object chunks - rpc StreamObj(stream ObjChunk) returns (AckReply); + rpc StreamObjTo(StreamObjToRequest) returns (stream ObjChunk); // Notify the object store about objref aliasing. This is called by the scheduler rpc NotifyAlias(NotifyAliasRequest) returns (AckReply); // Tell the object store to deallocate an object held by the object store. This is called by the scheduler. @@ -161,8 +161,8 @@ service ObjStore { rpc ObjStoreInfo(ObjStoreInfoRequest) returns (ObjStoreInfoReply); } -message DeliverObjRequest { - string objstore_address = 1; // Object store to deliver the object to +message StartDeliveryRequest { + string objstore_address = 1; // Object store to pull the object from uint64 objref = 2; // Reference of object that gets delivered } @@ -174,9 +174,13 @@ message RegisterObjReply { uint64 handle = 1; // Handle to memory segment where object is stored } -message ObjChunk { +message StreamObjToRequest { uint64 objref = 1; // Object reference of the object being streamed - uint64 totalsize = 2; // Total size of the object +} + +message ObjChunk { + uint64 total_size = 1; // Total size of the object + uint64 metadata_offset = 2; // Offset of the arrow metadata bytes data = 3; // Data for this chunk of the object } @@ -193,12 +197,6 @@ message GetObjRequest { uint64 objref = 1; // Object reference of the object being requested by the worker } -message GetObjReply { - string bucket = 1; // Name of the shared memory segment where the object is stored - uint64 handle = 2; // Shared memory pointer to the object - uint64 size = 3; // Total size of the object in bytes -} - // These messages are for getting information about the object store state message ObjStoreInfoRequest { diff --git a/src/ipc.cc b/src/ipc.cc index bb4cac307..0a7a2bb2d 100644 --- a/src/ipc.cc +++ b/src/ipc.cc @@ -31,9 +31,9 @@ MemorySegmentPool::MemorySegmentPool(ObjStoreId objstoreid, bool create) : objst // 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 void MemorySegmentPool::open_segment(SegmentId segmentid, size_t size) { - ORCH_LOG(ORCH_DEBUG, "opening segmentid " << segmentid); + ORCH_LOG(ORCH_DEBUG, "Opening segmentid " << segmentid << " on object store " << objstoreid_ << " with create_mode_ = " << create_mode_); if (segmentid != segments_.size() && create_mode_) { - ORCH_LOG(ORCH_FATAL, "Attempting to open segmentid " << segmentid << " on the object store, but segments_.size() = " << segments_.size()); + ORCH_LOG(ORCH_FATAL, "Object store " << objstoreid_ << " is attempting to open segmentid " << segmentid << " on the object store, but segments_.size() = " << segments_.size()); } if (segmentid >= segments_.size()) { // resize and initialize segments_ int current_size = segments_.size(); @@ -90,7 +90,12 @@ void MemorySegmentPool::deallocate(ObjHandle pointer) { // returns address of the object refered to by the handle, needs to be called on // the process that will use the address uint8_t* MemorySegmentPool::get_address(ObjHandle pointer) { - open_segment(pointer.segmentid()); + if (create_mode_ && segments_[pointer.segmentid()].second != SegmentStatusType::OPENED) { + ORCH_LOG(ORCH_FATAL, "Object store " << objstoreid_ << " is attempting to call get_address on segmentid " << pointer.segmentid() << ", which has not been opened yet."); + } + if (!create_mode_) { + open_segment(pointer.segmentid()); + } managed_shared_memory* segment = segments_[pointer.segmentid()].first.get(); return static_cast(segment->get_address_from_handle(pointer.ipcpointer())); } diff --git a/src/objstore.h b/src/objstore.h index 750d7aa6b..36cda7df9 100644 --- a/src/objstore.h +++ b/src/objstore.h @@ -17,40 +17,50 @@ using grpc::ServerBuilder; using grpc::ServerReader; using grpc::ServerContext; using grpc::ClientContext; -using grpc::ClientWriter; +using grpc::ServerWriter; +using grpc::ClientReader; 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); -}; - -enum MemoryStatusType {READY = 0, NOT_READY = 1, DEALLOCATED = 2, NOT_PRESENT = 3}; +// READY: This is used to indicate that the object has been copied from a +// worker and is ready to be used. +// NOT_READY: This is used to indicate that memory has been allocated for the +// object, but the object hasn't been copied from a worker yet. +// DEALLOCATED: This is used to indicate that the object has been deallocated. +// NOT_PRESENT: This is used to indicate that space has not been allocated for +// this object in this object store. +// PRE_ALLOCED: This is used to indicate that the memory has not yet been +// alloced, but it will be alloced soon. This is set when we call +// StartDelivery. +enum MemoryStatusType {READY = 0, NOT_READY = 1, DEALLOCATED = 2, NOT_PRESENT = 3, PRE_ALLOCED = 4}; class ObjStoreService final : public ObjStore::Service { public: ObjStoreService(const std::string& objstore_address, std::shared_ptr scheduler_channel); - // Status DeliverObj(ServerContext* context, const DeliverObjRequest* request, AckReply* reply) override; - // Status StreamObj(ServerContext* context, ServerReader* reader, AckReply* reply) override; + Status StartDelivery(ServerContext* context, const StartDeliveryRequest* request, AckReply* reply) override; + Status StreamObjTo(ServerContext* context, const StreamObjToRequest* request, ServerWriter* writer) override; Status NotifyAlias(ServerContext* context, const NotifyAliasRequest* request, AckReply* reply) override; Status DeallocateObject(ServerContext* context, const DeallocateObjectRequest* request, AckReply* reply) override; Status ObjStoreInfo(ServerContext* context, const ObjStoreInfoRequest* request, ObjStoreInfoReply* reply) override; void start_objstore_service(); private: + void pull_data_from(ObjRef objref, ObjStore::Stub& stub); // 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); void process_worker_request(const ObjRequest request); void process_objstore_request(const ObjRequest request); void process_requests(); void process_pulls_for_objref(ObjRef objref); + ObjHandle alloc(ObjRef objref, size_t size); + void object_ready(ObjRef objref, size_t metadata_offset); + static const size_t CHUNK_SIZE; std::string objstore_address_; ObjStoreId objstoreid_; // id of this objectstore in the scheduler object store table std::shared_ptr segmentpool_; + std::mutex segmentpool_lock_; std::vector > memory_; // object reference -> (memory address, memory status) std::mutex memory_lock_; std::unordered_map> objstores_; @@ -58,9 +68,13 @@ private: std::unique_ptr scheduler_stub_; std::vector > pull_queue_; std::mutex pull_queue_lock_; - MessageQueue recv_queue_; - std::vector > send_queues_; // workerid -> queue + MessageQueue recv_queue_; // This queue is used by workers to send tasks to the object store. + std::vector > send_queues_; // This maps workerid -> queue. The object store uses these queues to send replies to the relevant workers. std::thread communicator_thread_; + + std::vector > delivery_threads_; // TODO(rkn): document + // TODO(rkn): possibly add lock, and properly remove these threads from the delivery_threads_ when the deliveries are done + }; #endif diff --git a/src/scheduler.cc b/src/scheduler.cc index 4f55497c9..8f0bdd0f1 100644 --- a/src/scheduler.cc +++ b/src/scheduler.cc @@ -191,8 +191,14 @@ Status SchedulerService::SchedulerInfo(ServerContext* context, const SchedulerIn return Status::OK; } +// TODO(rkn): This could execute multiple times with the same arguments before +// the delivery finishes, but we only want it to happen once. Currently, the +// redundancy is handled by the object store, which will only execute the +// delivery once. However, we may want to handle it in the scheduler in the +// future. +// +// deliver_object assumes that the aliasing for objref has already been completed. That is, has_canonical_objref(objref) == true void SchedulerService::deliver_object(ObjRef objref, ObjStoreId from, ObjStoreId to) { - // deliver_object assumes that the aliasing for objref has already been completed. That is, has_canonical_objref(objref) == true if (from == to) { ORCH_LOG(ORCH_FATAL, "attempting to deliver objref " << objref << " from objstore " << from << " to itself."); } @@ -201,12 +207,12 @@ void SchedulerService::deliver_object(ObjRef objref, ObjStoreId from, ObjStoreId } ClientContext context; AckReply reply; - DeliverObjRequest request; + StartDeliveryRequest request; ObjRef canonical_objref = get_canonical_objref(objref); request.set_objref(canonical_objref); std::lock_guard lock(objstores_lock_); - request.set_objstore_address(objstores_[to].address); - objstores_[from].objstore_stub->DeliverObj(&context, request, &reply); + request.set_objstore_address(objstores_[from].address); + objstores_[to].objstore_stub->StartDelivery(&context, request, &reply); } void SchedulerService::schedule() { @@ -403,14 +409,16 @@ void SchedulerService::get_info(const SchedulerInfoRequest& request, SchedulerIn } +// pick_objstore assumes that objtable_lock_ has been acquired +// pick_objstore must be called with a canonical_objref ObjStoreId SchedulerService::pick_objstore(ObjRef canonical_objref) { - // pick_objstore must be called with a canonical_objref std::mt19937 rng; if (!is_canonical(canonical_objref)) { ORCH_LOG(ORCH_FATAL, "Attempting to call pick_objstore with a non-canonical objref, (objref " << canonical_objref << ")"); } std::uniform_int_distribution uni(0, objtable_[canonical_objref].size() - 1); - return uni(rng); + ObjStoreId objstoreid = objtable_[canonical_objref][uni(rng)]; + return objstoreid; } bool SchedulerService::is_canonical(ObjRef objref) { @@ -433,7 +441,7 @@ void SchedulerService::perform_pulls() { continue; } ObjRef canonical_objref = get_canonical_objref(objref); - ORCH_LOG(ORCH_DEBUG, "attempting to pull objref " << pull.second << " with canonical objref " << canonical_objref); + ORCH_LOG(ORCH_DEBUG, "attempting to pull objref " << pull.second << " with canonical objref " << canonical_objref << " to objstore " << get_store(workerid)); objtable_lock_.lock(); int num_stores = objtable_[canonical_objref].size(); @@ -570,16 +578,17 @@ void SchedulerService::deallocate_object(ObjRef canonical_objref) { // so the lock must before outside of these methods (it is acquired in // DecrementRefCount). ORCH_LOG(ORCH_REFCOUNT, "Deallocating canonical_objref " << canonical_objref << "."); - ClientContext context; - AckReply reply; - DeallocateObjectRequest request; - request.set_canonical_objref(canonical_objref); { std::lock_guard objtable_lock(objtable_lock_); auto &objstores = objtable_[canonical_objref]; std::lock_guard objstores_lock(objstores_lock_); // TODO(rkn): Should this be inside the for loop instead? for (int i = 0; i < objstores.size(); ++i) { + ClientContext context; + AckReply reply; + DeallocateObjectRequest request; + request.set_canonical_objref(canonical_objref); ObjStoreId objstoreid = objstores[i]; + ORCH_LOG(ORCH_REFCOUNT, "Attempting to deallocate canonical_objref " << canonical_objref << " from objstore " << objstoreid); objstores_[objstoreid].objstore_stub->DeallocateObject(&context, request, &reply); } } diff --git a/src/serialize.cc b/src/serialize.cc index c564eae6f..3fbf83fdf 100644 --- a/src/serialize.cc +++ b/src/serialize.cc @@ -156,7 +156,7 @@ std::shared_ptr read_flat_array(BufferMemorySource* source, int64_ std::shared_ptr reader; Status s = ipc::RowBatchReader::Open(source, metadata_offset, &reader); if (!s.ok()) { - ORCH_LOG(ORCH_FATAL, s.ToString()); + ORCH_LOG(ORCH_FATAL, "Error in read_flat_array: " << s.ToString()); } auto field = std::make_shared("data", npy_traits::primitive_type); std::shared_ptr schema(new arrow::Schema({field})); diff --git a/test/arrays_test.py b/test/arrays_test.py index 4a2644694..685474b1f 100644 --- a/test/arrays_test.py +++ b/test/arrays_test.py @@ -81,7 +81,7 @@ class ArraysDistTest(unittest.TestCase): def testMethods(self): test_dir = os.path.dirname(os.path.abspath(__file__)) test_path = os.path.join(test_dir, "testrecv.py") - services.start_cluster(return_drivers=False, num_workers_per_objstore=8, worker_path=test_path) + services.start_cluster(return_drivers=False, num_objstores=2, num_workers_per_objstore=8, worker_path=test_path) x = dist.zeros([9, 25, 51], "float") y = dist.assemble(x) diff --git a/test/runtest.py b/test/runtest.py index 858fd5e28..349c00167 100644 --- a/test/runtest.py +++ b/test/runtest.py @@ -78,11 +78,40 @@ class ObjStoreTest(unittest.TestCase): self.assertEqual(result, data) # pushing an object, shipping it to another worker, and pulling it shouldn't change it - # for data in ["h", "h" * 10000, 0, 0.0]: - # objref = worker.push(data, worker1) - # response = objstore1_stub.DeliverObj(orchestra_pb2.DeliverObjRequest(objref=objref.val, objstore_address=address(IP_ADDRESS, objstore2_port)), TIMEOUT_SECONDS) - # result = worker.pull(objref, worker2) - # self.assertEqual(result, data) + for data in ["h", "h" * 10000, 0, 0.0, [1, 2, 3, "a", (1, 2)], ("a", ("b", 3))]: + objref = worker.push(data, w1) + result = worker.pull(objref, w2) + self.assertEqual(result, data) + + # pushing an array, shipping it to another worker, and pulling it shouldn't change it + for data in [np.zeros([10, 20]), np.random.normal(size=[45, 25])]: + objref = worker.push(data, w1) + result = worker.pull(objref, w2) + self.assertTrue(np.alltrue(result == data)) + + """ + # pulling multiple times shouldn't matter + for data in [np.zeros([10, 20]), np.random.normal(size=[45, 25]), np.zeros([10, 20], dtype=np.dtype("float64")), np.zeros([10, 20], dtype=np.dtype("float32")), np.zeros([10, 20], dtype=np.dtype("int64")), np.zeros([10, 20], dtype=np.dtype("int32"))]: + objref = worker.push(data, w1) + result = worker.pull(objref, w2) + result = worker.pull(objref, w2) + result = worker.pull(objref, w2) + self.assertTrue(np.alltrue(result == data)) + """ + + # shipping a numpy array inside something else should be fine + data = ("a", np.random.normal(size=[10, 10])) + objref = worker.push(data, w1) + result = worker.pull(objref, w2) + self.assertTrue(data[0] == result[0]) + self.assertTrue(np.alltrue(data[1] == result[1])) + + # shipping a numpy array inside something else should be fine + data = ["a", np.random.normal(size=[10, 10])] + objref = worker.push(data, w1) + result = worker.pull(objref, w2) + self.assertTrue(data[0] == result[0]) + self.assertTrue(np.alltrue(data[1] == result[1])) services.cleanup() From 4e65dfc2aeae72f9c6d9f0d91e693d026f570c43 Mon Sep 17 00:00:00 2001 From: Philipp Moritz Date: Sun, 24 Apr 2016 19:09:02 -0700 Subject: [PATCH 2/2] support multiple object stores, part two --- src/objstore.cc | 242 ++++++++++++++++++++++++++++-------------------- 1 file changed, 141 insertions(+), 101 deletions(-) diff --git a/src/objstore.cc b/src/objstore.cc index 5f637aa64..b3e1e0b95 100644 --- a/src/objstore.cc +++ b/src/objstore.cc @@ -1,25 +1,44 @@ #include "objstore.h" #include -const size_t ObjStoreClient::CHUNK_SIZE = 8 * 1024; +const size_t ObjStoreService::CHUNK_SIZE = 8 * 1024; // this method needs to be protected by a objstore_lock_ -Status ObjStoreClient::upload_data_to(slice data, ObjRef objref, ObjStore::Stub& stub) { +// TODO(rkn): Make sure that we do not in fact need the objstore_lock_. We want multiple deliveries to be able to happen simultaneously. +void ObjStoreService::pull_data_from(ObjRef objref, ObjStore::Stub& stub) { + ORCH_LOG(ORCH_DEBUG, "Objstore " << objstoreid_ << " is beginning to pull objref " << objref); ObjChunk chunk; ClientContext context; - AckReply reply; - std::unique_ptr > writer(stub.StreamObj(&context, &reply)); - const uint8_t* 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)) { - ORCH_LOG(ORCH_FATAL, "stream connection prematurely closed") - } + StreamObjToRequest stream_request; + stream_request.set_objref(objref); + std::unique_ptr > reader(stub.StreamObjTo(&context, stream_request)); + + size_t total_size = 0; + ObjHandle handle; + if (reader->Read(&chunk)) { + total_size = chunk.total_size(); + handle = alloc(objref, total_size); } - writer->WritesDone(); - return writer->Finish(); + size_t num_bytes = 0; + segmentpool_lock_.lock(); + uint8_t* data = segmentpool_->get_address(handle); + segmentpool_lock_.unlock(); + do { + if (num_bytes + chunk.data().size() > total_size) { + ORCH_LOG(ORCH_FATAL, "The reader attempted to stream too many bytes."); + } + std::memcpy(data, chunk.data().c_str(), chunk.data().size()); + data += chunk.data().size(); + num_bytes += chunk.data().size(); + } while (reader->Read(&chunk)); + Status status = reader->Finish(); // Right now we don't use the status. + + // finalize object + if (num_bytes != total_size) { + ORCH_LOG(ORCH_FATAL, "Streamed objref " << objref << ", but num_bytes != total_size"); + } + object_ready(objref, chunk.metadata_offset()); + ORCH_LOG(ORCH_DEBUG, "finished streaming data, objref was " << objref << " and size was " << num_bytes); } ObjStoreService::ObjStoreService(const std::string& objstore_address, std::shared_ptr scheduler_channel) @@ -44,15 +63,35 @@ ObjStore::Stub& ObjStoreService::get_objstore_stub(const std::string& objstore_a return *objstores_[objstore_address]; } -/* -Status ObjStoreService::DeliverObj(ServerContext* context, const DeliverObjRequest* request, AckReply* reply) { - std::lock_guard objstores_lock(objstores_lock_); - ObjStore::Stub& stub = get_objstore_stub(request->objstore_address()); +Status ObjStoreService::StartDelivery(ServerContext* context, const StartDeliveryRequest* request, AckReply* reply) { + // TODO(rkn): We're pushing the delivery task onto a new thread so that this method can return immediately. This matters + // because the scheduler holds a lock while DeliverObj is being called. The correct solution is to make DeliverObj + // an asynchronous call (and similarly with the rest of the object store service methods). + std::string address = request->objstore_address(); ObjRef objref = request->objref(); - Status status = ObjStoreClient::upload_data_to(memory_[objref].ptr, objref, stub); - return status; + { + std::lock_guard memory_lock(memory_lock_); + if (objref >= memory_.size()) { + memory_.resize(objref + 1, std::make_pair(ObjHandle(), MemoryStatusType::NOT_PRESENT)); + } + if (memory_[objref].second == MemoryStatusType::NOT_PRESENT) { + } + else if (memory_[objref].second == MemoryStatusType::DEALLOCATED) { + ORCH_LOG(ORCH_FATAL, "Objstore " << objstoreid_ << " is attempting to get objref " << objref << ", but memory_[objref] == DEALLOCATED."); + } + else { + ORCH_LOG(ORCH_DEBUG, "Objstore " << objstoreid_ << " already has objref " << objref << " or it is already being shipped, so no need to pull it again."); + return Status::OK; + } + memory_[objref].second = MemoryStatusType::PRE_ALLOCED; + } + delivery_threads_.push_back(std::make_shared([this, address, objref]() { + std::lock_guard objstores_lock(objstores_lock_); + ObjStore::Stub& stub = get_objstore_stub(address); + pull_data_from(objref, stub); + })); + return Status::OK; } -*/ Status ObjStoreService::ObjStoreInfo(ServerContext* context, const ObjStoreInfoRequest* request, ObjStoreInfoReply* reply) { std::lock_guard memory_lock(memory_lock_); @@ -73,70 +112,59 @@ Status ObjStoreService::ObjStoreInfo(ServerContext* context, const ObjStoreInfoR return Status::OK; } -/* -Status ObjStoreService::StreamObj(ServerContext* context, ServerReader* reader, AckReply* reply) { - ORCH_LOG(ORCH_VERBOSE, "begin to stream data to object store " << objstoreid_); - memory_lock_.lock(); +Status ObjStoreService::StreamObjTo(ServerContext* context, const StreamObjToRequest* request, ServerWriter* writer) { + ORCH_LOG(ORCH_DEBUG, "begin to stream data from object store " << objstoreid_); ObjChunk chunk; - ObjRef objref = 0; - size_t totalsize = 0; - if (reader->Read(&chunk)) { - objref = chunk.objref(); - totalsize = chunk.totalsize(); - allocate_memory(objref, totalsize); + ObjRef objref = request->objref(); + memory_lock_.lock(); + if (objref >= memory_.size()) { + ORCH_LOG(ORCH_FATAL, "Objstore " << objstoreid_ << " is attempting to use objref " << objref << " in StreamObjTo, but this objref is not present in the object store."); } - size_t num_bytes = 0; - char* data = memory_[objref].ptr.data; - - do { - if (num_bytes + chunk.data().size() > totalsize) { - memory_lock_.unlock(); - return Status::CANCELLED; + if (memory_[objref].second != MemoryStatusType::READY) { + ORCH_LOG(ORCH_FATAL, "Objstore " << objstoreid_ << " is attempting to stream objref " << objref << ", but memory_[objref].second != MemoryStatusType::READY."); + } + ObjHandle handle = memory_[objref].first; + memory_lock_.unlock(); // TODO(rkn): Make sure we don't still need to hold on to this lock. + segmentpool_lock_.lock(); + const uint8_t* head = segmentpool_->get_address(handle); + segmentpool_lock_.unlock(); + size_t size = handle.size(); + for (size_t i = 0; i < size; i += CHUNK_SIZE) { + chunk.set_metadata_offset(handle.metadata_offset()); + chunk.set_total_size(size); + chunk.set_data(head + i, std::min(CHUNK_SIZE, size - i)); + if (!writer->Write(chunk)) { + ORCH_LOG(ORCH_FATAL, "stream connection prematurely closed") } - std::memcpy(data, chunk.data().c_str(), chunk.data().size()); - data += chunk.data().size(); - num_bytes += chunk.data().size(); - } while (reader->Read(&chunk)); - - ORCH_LOG(ORCH_VERBOSE, "finished streaming data, objref was " << objref << " and size was " << num_bytes); - - memory_lock_.unlock(); - - ClientContext objready_context; - ObjReadyRequest objready_request; - objready_request.set_objref(objref); - objready_request.set_objstoreid(objstoreid_); - AckReply objready_reply; - scheduler_stub_->ObjReady(&objready_context, objready_request, &objready_reply); - + } return Status::OK; } -*/ Status ObjStoreService::NotifyAlias(ServerContext* context, const NotifyAliasRequest* request, AckReply* reply) { // NotifyAlias assumes that the objstore already holds canonical_objref ObjRef alias_objref = request->alias_objref(); ObjRef canonical_objref = request->canonical_objref(); ORCH_LOG(ORCH_DEBUG, "Aliasing objref " << alias_objref << " with objref " << canonical_objref); - std::lock_guard memory_lock(memory_lock_); - if (canonical_objref >= memory_.size()) { - ORCH_LOG(ORCH_FATAL, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " is not in the objstore.") + { + std::lock_guard memory_lock(memory_lock_); + if (canonical_objref >= memory_.size()) { + ORCH_LOG(ORCH_FATAL, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " is not in the objstore.") + } + if (memory_[canonical_objref].second == MemoryStatusType::NOT_READY) { + ORCH_LOG(ORCH_FATAL, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " is not ready yet in the objstore.") + } + if (memory_[canonical_objref].second == MemoryStatusType::NOT_PRESENT) { + ORCH_LOG(ORCH_FATAL, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " is not present in the objstore.") + } + if (memory_[canonical_objref].second == MemoryStatusType::DEALLOCATED) { + ORCH_LOG(ORCH_FATAL, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " has already been deallocated.") + } + if (alias_objref >= memory_.size()) { + memory_.resize(alias_objref + 1, std::make_pair(ObjHandle(), MemoryStatusType::NOT_PRESENT)); + } + memory_[alias_objref].first = memory_[canonical_objref].first; + memory_[alias_objref].second = MemoryStatusType::READY; } - if (memory_[canonical_objref].second == MemoryStatusType::NOT_READY) { - ORCH_LOG(ORCH_FATAL, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " is not ready yet in the objstore.") - } - if (memory_[canonical_objref].second == MemoryStatusType::NOT_PRESENT) { - ORCH_LOG(ORCH_FATAL, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " is not present in the objstore.") - } - if (memory_[canonical_objref].second == MemoryStatusType::DEALLOCATED) { - ORCH_LOG(ORCH_FATAL, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " has already been deallocated.") - } - if (alias_objref >= memory_.size()) { - memory_.resize(alias_objref + 1, std::make_pair(ObjHandle(), MemoryStatusType::NOT_PRESENT)); - } - memory_[alias_objref].first = memory_[canonical_objref].first; - memory_[alias_objref].second = MemoryStatusType::READY; - ObjRequest done_request; done_request.type = ObjRequestType::ALIAS_DONE; done_request.objref = alias_objref; @@ -154,7 +182,9 @@ 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_lock_.lock(); segmentpool_->deallocate(memory_[canonical_objref].first); + segmentpool_lock_.unlock(); memory_[canonical_objref].second = MemoryStatusType::DEALLOCATED; return Status::OK; } @@ -197,15 +227,8 @@ 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); - send_queues_[request.workerid].send(&reply); - std::lock_guard memory_lock(memory_lock_); - if (memory_[request.objref].second != MemoryStatusType::NOT_PRESENT) { - ORCH_LOG(ORCH_FATAL, "Attempting to allocate space for objref " << request.objref << ", but memory_[objref].second != MemoryStatusType::NOT_PRESENT, it equals " << memory_[request.objref].second); - } - memory_[request.objref].first = reply; - memory_[request.objref].second = MemoryStatusType::NOT_READY; + ObjHandle handle = alloc(request.objref, request.size); // This method acquires memory_lock_ + send_queues_[request.workerid].send(&handle); } break; case ObjRequestType::GET: { @@ -214,7 +237,7 @@ void ObjStoreService::process_worker_request(const ObjRequest request) { if (item.second == MemoryStatusType::READY) { ORCH_LOG(ORCH_DEBUG, "Responding to GET request: returning objref " << request.objref); send_queues_[request.workerid].send(&item.first); - } else if (item.second == MemoryStatusType::NOT_READY || item.second == MemoryStatusType::NOT_PRESENT) { + } else if (item.second == MemoryStatusType::NOT_READY || item.second == MemoryStatusType::NOT_PRESENT || item.second == MemoryStatusType::PRE_ALLOCED) { std::lock_guard lock(pull_queue_lock_); pull_queue_.push_back(std::make_pair(request.workerid, request.objref)); } else { @@ -223,24 +246,7 @@ void ObjStoreService::process_worker_request(const ObjRequest request) { } break; case ObjRequestType::WORKER_DONE: { - { - std::lock_guard memory_lock(memory_lock_); - std::pair& item = memory_[request.objref]; - if (item.second != MemoryStatusType::NOT_READY) { - ORCH_LOG(ORCH_FATAL, "A worker notified the object store that objref " << request.objref << " has been written to the object store, but memory_[objref].second != NOT_READY."); - } - item.first.set_metadata_offset(request.metadata_offset); - item.second = MemoryStatusType::READY; - } - process_pulls_for_objref(request.objref); - // Tell the scheduler that the object arrived - // TODO(pcm): put this in a separate thread so we don't have to pay the latency here - ClientContext objready_context; - ObjReadyRequest objready_request; - objready_request.set_objref(request.objref); - objready_request.set_objstoreid(objstoreid_); - AckReply objready_reply; - scheduler_stub_->ObjReady(&objready_context, objready_request, &objready_reply); + object_ready(request.objref, request.metadata_offset); // This method acquires memory_lock_ } break; default: { @@ -293,6 +299,40 @@ void ObjStoreService::process_pulls_for_objref(ObjRef objref) { } } +ObjHandle ObjStoreService::alloc(ObjRef objref, size_t size) { + segmentpool_lock_.lock(); + ObjHandle handle = segmentpool_->allocate(size); + segmentpool_lock_.unlock(); + std::lock_guard memory_lock(memory_lock_); + if (memory_[objref].second != MemoryStatusType::NOT_PRESENT && memory_[objref].second != MemoryStatusType::PRE_ALLOCED) { + ORCH_LOG(ORCH_FATAL, "Attempting to allocate space for objref " << objref << ", but memory_[objref].second = " << memory_[objref].second); + } + memory_[objref].first = handle; + memory_[objref].second = MemoryStatusType::NOT_READY; + return handle; +} + +void ObjStoreService::object_ready(ObjRef objref, size_t metadata_offset) { + { + std::lock_guard memory_lock(memory_lock_); + std::pair& item = memory_[objref]; + if (item.second != MemoryStatusType::NOT_READY) { + ORCH_LOG(ORCH_FATAL, "A worker notified the object store that objref " << objref << " has been written to the object store, but memory_[objref].second != NOT_READY."); + } + item.first.set_metadata_offset(metadata_offset); + item.second = MemoryStatusType::READY; + } + process_pulls_for_objref(objref); + // Tell the scheduler that the object arrived + // TODO(pcm): put this in a separate thread so we don't have to pay the latency here + ClientContext objready_context; + ObjReadyRequest objready_request; + objready_request.set_objref(objref); + objready_request.set_objstoreid(objstoreid_); + AckReply objready_reply; + scheduler_stub_->ObjReady(&objready_context, objready_request, &objready_reply); +} + void ObjStoreService::start_objstore_service() { communicator_thread_ = std::thread([this]() { ORCH_LOG(ORCH_INFO, "started object store communicator server");