// This file defines the GRPC interface between scheduler, object stores and
// workers. These are used for communication over the network.

// Terminology:
// Worker: A cluster consists of multiple worker processes (typically one
//   per core) which execute tasks that can access objects from object stores.
// Object store: Typically there is one object store per node which holds the
//   objects locally stored on that node.
// Scheduler: The scheduler process keeps track of a mapping from object
//   references to object stores, orchestrates data transfer between object
//   stores and assigns tasks to workers.

syntax = "proto3";

import "types.proto";

// Scheduler

service Scheduler {
  // Register a new worker with the scheduler
  rpc RegisterWorker(RegisterWorkerRequest) returns (RegisterWorkerReply);
  // Register an object store with the scheduler
  rpc RegisterObjStore(RegisterObjStoreRequest) returns (RegisterObjStoreReply);
  // Tell the scheduler that a worker can execute a certain function
  rpc RegisterFunction(RegisterFunctionRequest) returns (AckReply);
  // Asks the scheduler to execute a task, immediately returns an object reference to the result
  rpc SubmitTask(SubmitTaskRequest) returns (SubmitTaskReply);
  // Increment the count of the object reference
  rpc IncrementCount(ChangeCountRequest) returns (AckReply);
  // Decrement the count of the object reference
  rpc DecrementCount(ChangeCountRequest) returns (AckReply);
  // Request an object reference for an object that will be put in an object store
  rpc PutObj(PutObjRequest) returns (PutObjReply);
  // Request delivery of an object from an object store that holds the object to the local object store
  rpc RequestObj(RequestObjRequest) returns (AckReply);
  // Used by the worker to tell the scheduler that two objrefs should refer to the same object
  rpc AliasObjRefs(AliasObjRefsRequest) returns (AckReply);
  // Used by an object store to tell the scheduler that an object is ready (i.e. has been finalized and can be shared)
  rpc ObjReady(ObjReadyRequest) returns (AckReply);
  // Increments the reference count of a particular object reference
  rpc IncrementRefCount(IncrementRefCountRequest) returns (AckReply);
  // Decrements the reference count of a particular object reference
  rpc DecrementRefCount(DecrementRefCountRequest) returns (AckReply);
  // Used by the worker to notify the scheduler about which objrefs a particular object contains
  rpc AddContainedObjRefs(AddContainedObjRefsRequest) returns (AckReply);
  // Used by the worker to ask for work, this also returns the status of the previous task if there was one
  rpc ReadyForNewTask(ReadyForNewTaskRequest) returns (AckReply);
  // Get information about the scheduler state
  rpc SchedulerInfo(SchedulerInfoRequest) returns (SchedulerInfoReply);
  // Get information about tasks
  rpc TaskInfo(TaskInfoRequest) returns (TaskInfoReply);
}

message AckReply {
}

message RegisterWorkerRequest {
  string worker_address = 1; // IP address of the worker being registered
  string objstore_address = 2; // IP address of the object store the worker is connected to
}

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 {
  string objstore_address = 1; // IP address of the object store being registered
}

message RegisterObjStoreReply {
  uint64 objstoreid = 1; // Object store ID assigned by the scheduler
}

message RegisterFunctionRequest {
  uint64 workerid = 1; // Worker that can execute the function
  string fnname = 2; // Name of the function that is registered
  uint64 num_return_vals = 3; // Number of return values of the function
}

message SubmitTaskRequest {
  uint64 workerid = 1; // The ID of the worker submitting the task
  Task task = 2; // Contains name of the function to be executed and arguments
}

message SubmitTaskReply {
  repeated uint64 result = 1; // Object references of the function return values
  bool function_registered = 2; // True if the function was registered; false otherwise
}

message RequestObjRequest {
  uint64 workerid = 1; // Worker that tries to request the object
  uint64 objref = 2; // Object reference of the object being requested
}

message PutObjRequest {
  uint64 workerid = 1; // Worker that tries to put an object
}

message PutObjReply {
  uint64 objref = 1; // Object reference assigned by the scheduler to the object
}

message AliasObjRefsRequest {
  uint64 alias_objref = 1; // ObjRef which will be aliased
  uint64 target_objref = 2; // The target ObjRef
}

message ObjReadyRequest {
  uint64 objref = 1; // Object reference of the object that has been finalized
  uint64 objstoreid = 2; // ID of the object store the object lives on
}

message IncrementRefCountRequest {
  repeated uint64 objref = 1; // Object references whose reference count should be incremented. Duplicates will be incremented multiple times.
}

message AddContainedObjRefsRequest {
  uint64 objref = 1; // The objref of the object in question
  repeated uint64 contained_objref = 2; // Object references contained in the object
}

message DecrementRefCountRequest {
  repeated uint64 objref = 1; // Object references whose reference count should be decremented. Duplicates will be decremented multiple times.
}

message ReadyForNewTaskRequest {
  uint64 workerid = 1; // ID of the worker which executed the task
  message PreviousTaskInfo {
    bool task_succeeded = 1; // True if the task succeeded, false if it threw an exception
    string error_message = 2; // The contents of the exception, if the task threw an exception
  }
  PreviousTaskInfo previous_task_info = 2; // Information about the previous task, this is only present if there was a previous task
}

message ChangeCountRequest {
  uint64 objref = 1; // Object reference of the object whose reference count is increased or decreased
}

// The following messages are used to get information about the scheduler state

message SchedulerInfoRequest {
}

message FnTableEntry {
  repeated uint64 workerid = 1; // ID of the worker that can execute the function
  uint64 num_return_vals = 2; // Number of return values of the function
}

message SchedulerInfoReply {
  repeated uint64 operationid = 1; // OperationIds of the tasks on the task queue
  repeated uint64 avail_worker = 3; // List of workers waiting to get a task assigned
  map<string, FnTableEntry> function_table = 2; // Table of all available remote function
  repeated uint64 target_objref = 4; // The target_objrefs_ data structure
  repeated uint64 reference_count = 5; // The reference_counts_ data structure
}

// Object stores

service ObjStore {
  // Tell the object store to begin getting 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 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.
  rpc DeallocateObject(DeallocateObjectRequest) returns (AckReply);
  // Get info about the object store state
  rpc ObjStoreInfo(ObjStoreInfoRequest) returns (ObjStoreInfoReply);
}

message StartDeliveryRequest {
  string objstore_address = 1; // Object store to get the object from
  uint64 objref = 2; // Reference of object that gets delivered
}

message RegisterObjRequest {
  uint64 objref = 1; // Reference of object that gets registered
}

message RegisterObjReply {
  uint64 handle = 1; // Handle to memory segment where object is stored
}

message StreamObjToRequest {
  uint64 objref = 1; // Object reference of the object being streamed
}

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
}

message NotifyAliasRequest {
  uint64 alias_objref = 1; // The objref being aliased
  uint64 canonical_objref = 2; // The canonical objref that points to the actual object
}

message DeallocateObjectRequest {
  uint64 canonical_objref = 1; // The canonical objref of the object to deallocate
}

message GetObjRequest {
  uint64 objref = 1; // Object reference of the object being requested by the worker
}

message TaskInfoRequest {
}

message TaskInfoReply {
  repeated TaskStatus failed_task = 1;
  repeated TaskStatus running_task = 2;
  uint64 num_succeeded = 3;
  // TODO(mehrdadn): We'll want to return information from computation_graph since it's important for visualizing tasks that have been completed etc.
}

// These messages are for getting information about the object store state

message ObjStoreInfoRequest {
  repeated uint64 objref = 1; // Object references we want to retrieve from the store for inspection
}

message ObjStoreInfoReply {
  repeated uint64 objref = 1; // List of object references in the store
  repeated Obj obj = 2; // Protocol buffer objects that were requested
}

// Workers

service WorkerService {
  rpc ExecuteTask(ExecuteTaskRequest) returns (ExecuteTaskReply); // Scheduler calls a function from the worker
}

message ExecuteTaskRequest {
  Task task = 1; // Contains name of the function to be executed and arguments
}

message ExecuteTaskReply {

}