syntax = "proto3";

package dgs;

enum WorkerType {
  Sampling   = 0;
  Serving    = 1;
};

// Common message types

message UpStreamInfoPb {
  repeated string sub_kafka_servers = 1;
  string sub_kafka_topic = 2;
  uint32 sub_kafka_partition_num = 3;
  repeated uint32 sub_kafka_pids = 4;
}

message DownStreamKafkaInfoPb {
  repeated string pub_kafka_servers = 1;
  string pub_kafka_topic = 2;
  uint32 pub_kafka_partition_num = 3;
}

// When sampling worker publishing sample updates, they will be partitioned
// by worker-wise strategy according to the total worker number of serving workers.
message DownStreamPartitionInfoPb {
  string worker_partition_strategy = 1;
  uint32 worker_partition_num = 2;
  // vector router: kafka partition id -> serving worker id
  repeated uint32 kafka_to_wid = 3;
}

message ManagedStorePids {
  uint32 worker_id = 1;
  repeated uint32 pids = 2;
}

message StorePartitionInfoPb {
  string partition_strategy = 1;
  uint32 partition_num = 2;
  repeated ManagedStorePids managed_pids_group = 3;
}

message KafkaPartitionOffsetInfoPb {
  uint32 pid = 1;
  int64  ready_offset = 2;
}

message StorePartitionBackupInfoPb {
  uint32  pid = 1;
  bool    valid = 2;
  uint32  vertex_bid = 3;
  uint32  edge_bid = 4;
}

message SubsPartitionBackupInfoPb {
  uint32  pid = 1;
  bool    valid = 2;
  uint32  bid = 3;
}

message CheckpointInfoPb {
  repeated KafkaPartitionOffsetInfoPb sub_kafka_offsets = 1;
  repeated StorePartitionBackupInfoPb sample_store_backups = 2;
  repeated SubsPartitionBackupInfoPb  subs_table_backups = 3;
}

message SamplingInitInfoPb {
  string query_plan = 1;
  uint32 num_local_shards = 2;
  repeated string ipaddrs = 3;
  StorePartitionInfoPb store_partition_info = 4;
  UpStreamInfoPb upstream_info = 5;
  DownStreamKafkaInfoPb ds_kafka_info = 6;
  DownStreamPartitionInfoPb ds_partition_info = 7;
  CheckpointInfoPb checkpoint_info = 8;
}

message ServingInitInfoPb {
  string query_plan = 1;
  uint32 num_local_shards = 2;
  StorePartitionInfoPb store_partition_info = 3;
  UpStreamInfoPb   upstream_info = 4;
  CheckpointInfoPb checkpoint_info = 5;
}

/// rpc message types

message RegisterWorkerRequestPb {
  WorkerType worker_type = 1;
  uint32     worker_id = 2;
  string     worker_ip = 3;
}

message RegisterWorkerResponsePb {
  bool   suc = 1;
  uint32 num_workers = 2;
}

message GetInitInfoRequestPb {
  WorkerType worker_type = 1;
  uint32     worker_id = 2;
}

message GetInitInfoResponsePb {
  SamplingInitInfoPb   sampling_info = 1;
  ServingInitInfoPb    serving_info = 2;
  bool                 terminate_service = 3;
}

message ReportStartedRequestPb {
  WorkerType worker_type = 1;
  uint32     worker_id = 2;
}

message ReportStartedResponsePb {
  bool terminate_service = 1;
}

message ReportStatisticsRequestPb {
  WorkerType worker_type = 1;
  uint32     worker_id = 2;
  repeated KafkaPartitionOffsetInfoPb ready_kafka_offsets = 3;
}

message ReportStatisticsResponsePb {
  bool terminate_service = 1;
  bool do_backup = 2;
}

message ReportBackupFinishedRequestPb {
  WorkerType worker_type = 1;
  uint32     worker_id = 2;
  repeated StorePartitionBackupInfoPb sample_store_backups = 3;
  repeated SubsPartitionBackupInfoPb  subs_table_backups = 4;
}

message ReportBackupFinishedResponsePb {
  bool is_valid = 1;
}

service Coordinator {
  // From DataLoader, SamplingWorkers and ServingWorkers, register worker and get init info.
  rpc RegisterWorker (RegisterWorkerRequestPb) returns (RegisterWorkerResponsePb) {}

  rpc GetInitInfo (GetInitInfoRequestPb) returns (GetInitInfoResponsePb) {}

  rpc ReportStarted (ReportStartedRequestPb) returns (ReportStartedResponsePb) {}

  rpc ReportStatistics (ReportStatisticsRequestPb) returns (ReportStatisticsResponsePb) {}

  rpc ReportBackupFinished (ReportBackupFinishedRequestPb) returns (ReportBackupFinishedResponsePb) {}
}
