mirror of
https://github.com/wassname/ray.git
synced 2026-07-07 06:55:53 +08:00
Robert Nishihara
67 lines
2.5 KiB
C++
67 lines
2.5 KiB
C++
#ifndef ORCHESTRA_OBJSTORE_H
|
|
#define ORCHESTRA_OBJSTORE_H
|
|
|
|
#include <unordered_map>
|
|
#include <memory>
|
|
#include <thread>
|
|
#include <iostream>
|
|
#include <grpc++/grpc++.h>
|
|
|
|
#include "orchestra/orchestra.h"
|
|
#include "orchestra.grpc.pb.h"
|
|
#include "types.pb.h"
|
|
#include "ipc.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);
|
|
};
|
|
|
|
enum MemoryStatusType {READY = 0, NOT_READY = 1, DEALLOCATED = 2, NOT_PRESENT = 3};
|
|
|
|
class ObjStoreService final : public ObjStore::Service {
|
|
public:
|
|
ObjStoreService(const std::string& objstore_address, std::shared_ptr<Channel> scheduler_channel);
|
|
|
|
// Status DeliverObj(ServerContext* context, const DeliverObjRequest* request, AckReply* reply) override;
|
|
// Status StreamObj(ServerContext* context, ServerReader<ObjChunk>* reader, AckReply* reply) 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:
|
|
// 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);
|
|
|
|
std::string objstore_address_;
|
|
ObjStoreId objstoreid_; // id of this objectstore in the scheduler object store table
|
|
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_;
|
|
std::mutex objstores_lock_;
|
|
std::unique_ptr<Scheduler::Stub> scheduler_stub_;
|
|
std::vector<std::pair<WorkerId, ObjRef> > pull_queue_;
|
|
std::mutex pull_queue_lock_;
|
|
MessageQueue<ObjRequest> recv_queue_;
|
|
std::vector<MessageQueue<ObjHandle> > send_queues_; // workerid -> queue
|
|
std::thread communicator_thread_;
|
|
};
|
|
|
|
#endif
|