[Streaming] Streaming data transfer and python integration (#6185)

streaming/src/queue/queue.cc

@@ -0,0 +1,211 @@
+#include "queue.h"
+#include <chrono>
+#include <thread>
+#include "queue_handler.h"
+#include "util/streaming_util.h"
+
+namespace ray {
+namespace streaming {
+
+bool Queue::Push(QueueItem item) {
+  std::unique_lock<std::mutex> lock(mutex_);
+  if (max_data_size_ < item.DataSize() + data_size_) return false;
+
+  buffer_queue_.push_back(item);
+  data_size_ += item.DataSize();
+  readable_cv_.notify_one();
+  return true;
+}
+
+QueueItem Queue::FrontProcessed() {
+  std::unique_lock<std::mutex> lock(mutex_);
+  STREAMING_CHECK(buffer_queue_.size() != 0) << "WriterQueue Pop fail";
+
+  if (watershed_iter_ == buffer_queue_.begin()) {
+    return InvalidQueueItem();
+  }
+
+  QueueItem item = buffer_queue_.front();
+  return item;
+}
+
+QueueItem Queue::PopProcessed() {
+  std::unique_lock<std::mutex> lock(mutex_);
+  STREAMING_CHECK(buffer_queue_.size() != 0) << "WriterQueue Pop fail";
+
+  if (watershed_iter_ == buffer_queue_.begin()) {
+    return InvalidQueueItem();
+  }
+
+  QueueItem item = buffer_queue_.front();
+  buffer_queue_.pop_front();
+  data_size_ -= item.DataSize();
+  data_size_sent_ -= item.DataSize();
+  return item;
+}
+
+QueueItem Queue::PopPending() {
+  std::unique_lock<std::mutex> lock(mutex_);
+  auto it = std::next(watershed_iter_);
+  QueueItem item = *it;
+  data_size_sent_ += it->DataSize();
+  buffer_queue_.splice(watershed_iter_, buffer_queue_, it, std::next(it));
+  return item;
+}
+
+QueueItem Queue::PopPendingBlockTimeout(uint64_t timeout_us) {
+  std::unique_lock<std::mutex> lock(mutex_);
+  std::chrono::system_clock::time_point point =
+      std::chrono::system_clock::now() + std::chrono::microseconds(timeout_us);
+  if (readable_cv_.wait_until(lock, point, [this] {
+        return std::next(watershed_iter_) != buffer_queue_.end();
+      })) {
+    auto it = std::next(watershed_iter_);
+    QueueItem item = *it;
+    data_size_sent_ += it->DataSize();
+    buffer_queue_.splice(watershed_iter_, buffer_queue_, it, std::next(it));
+    return item;
+
+  } else {
+    uint8_t data[1];
+    return QueueItem(QUEUE_INVALID_SEQ_ID, data, 1, 0, true);
+  }
+}
+
+QueueItem Queue::BackPending() {
+  std::unique_lock<std::mutex> lock(mutex_);
+  if (std::next(watershed_iter_) == buffer_queue_.end()) {
+    uint8_t data[1];
+    return QueueItem(QUEUE_INVALID_SEQ_ID, data, 1, 0, true);
+  }
+  return buffer_queue_.back();
+}
+
+bool Queue::IsPendingEmpty() {
+  std::unique_lock<std::mutex> lock(mutex_);
+  return std::next(watershed_iter_) == buffer_queue_.end();
+}
+
+bool Queue::IsPendingFull(uint64_t data_size) {
+  std::unique_lock<std::mutex> lock(mutex_);
+  return max_data_size_ < data_size + data_size_;
+}
+
+size_t Queue::ProcessedCount() {
+  std::unique_lock<std::mutex> lock(mutex_);
+  if (watershed_iter_ == buffer_queue_.begin()) return 0;
+
+  auto begin = buffer_queue_.begin();
+  auto end = std::prev(watershed_iter_);
+
+  return end->SeqId() + 1 - begin->SeqId();
+}
+
+size_t Queue::PendingCount() {
+  std::unique_lock<std::mutex> lock(mutex_);
+  if (std::next(watershed_iter_) == buffer_queue_.end()) return 0;
+
+  auto begin = std::next(watershed_iter_);
+  auto end = std::prev(buffer_queue_.end());
+
+  return begin->SeqId() - end->SeqId() + 1;
+}
+
+Status WriterQueue::Push(uint64_t seq_id, uint8_t *data, uint32_t data_size,
+                         uint64_t timestamp, bool raw) {
+  if (IsPendingFull(data_size)) {
+    return Status::OutOfMemory("Queue Push OutOfMemory");
+  }
+
+  while (is_pulling_) {
+    STREAMING_LOG(INFO) << "This queue is sending pull data, wait.";
+    std::this_thread::sleep_for(std::chrono::milliseconds(10));
+  }
+
+  QueueItem item(seq_id, data, data_size, timestamp, raw);
+  Queue::Push(item);
+  return Status::OK();
+}
+
+void WriterQueue::Send() {
+  while (!IsPendingEmpty()) {
+    // FIXME: front -> send -> pop
+    QueueItem item = PopPending();
+    DataMessage msg(actor_id_, peer_actor_id_, queue_id_, item.SeqId(), item.Buffer(),
+                    item.IsRaw());
+    std::unique_ptr<LocalMemoryBuffer> buffer = msg.ToBytes();
+    STREAMING_CHECK(transport_ != nullptr);
+    transport_->Send(std::move(buffer),
+                     DownstreamQueueMessageHandler::peer_async_function_);
+  }
+}
+
+Status WriterQueue::TryEvictItems() {
+  STREAMING_LOG(INFO) << "TryEvictItems";
+  QueueItem item = FrontProcessed();
+  uint64_t first_seq_id = item.SeqId();
+  STREAMING_LOG(INFO) << "TryEvictItems first_seq_id: " << first_seq_id
+                      << " min_consumed_id_: " << min_consumed_id_
+                      << " eviction_limit_: " << eviction_limit_;
+  if (min_consumed_id_ == QUEUE_INVALID_SEQ_ID || first_seq_id > min_consumed_id_) {
+    return Status::OutOfMemory("The queue is full and some reader doesn't consume");
+  }
+
+  if (eviction_limit_ == QUEUE_INVALID_SEQ_ID || first_seq_id > eviction_limit_) {
+    return Status::OutOfMemory("The queue is full and eviction limit block evict");
+  }
+
+  uint64_t evict_target_seq_id = std::min(min_consumed_id_, eviction_limit_);
+
+  while (item.SeqId() <= evict_target_seq_id) {
+    PopProcessed();
+    STREAMING_LOG(INFO) << "TryEvictItems directly " << item.SeqId();
+    item = FrontProcessed();
+  }
+  return Status::OK();
+}
+
+void WriterQueue::OnNotify(std::shared_ptr<NotificationMessage> notify_msg) {
+  STREAMING_LOG(INFO) << "OnNotify target seq_id: " << notify_msg->SeqId();
+  min_consumed_id_ = notify_msg->SeqId();
+}
+
+void ReaderQueue::OnConsumed(uint64_t seq_id) {
+  STREAMING_LOG(INFO) << "OnConsumed: " << seq_id;
+  QueueItem item = FrontProcessed();
+  while (item.SeqId() <= seq_id) {
+    PopProcessed();
+    item = FrontProcessed();
+  }
+  Notify(seq_id);
+}
+
+void ReaderQueue::Notify(uint64_t seq_id) {
+  std::vector<TaskArg> task_args;
+  CreateNotifyTask(seq_id, task_args);
+  // SubmitActorTask
+
+  NotificationMessage msg(actor_id_, peer_actor_id_, queue_id_, seq_id);
+  std::unique_ptr<LocalMemoryBuffer> buffer = msg.ToBytes();
+
+  transport_->Send(std::move(buffer), UpstreamQueueMessageHandler::peer_async_function_);
+}
+
+void ReaderQueue::CreateNotifyTask(uint64_t seq_id, std::vector<TaskArg> &task_args) {}
+
+void ReaderQueue::OnData(QueueItem &item) {
+  if (item.SeqId() != expect_seq_id_) {
+    STREAMING_LOG(WARNING) << "OnData ignore seq_id: " << item.SeqId()
+                           << " expect_seq_id_: " << expect_seq_id_;
+    return;
+  }
+
+  last_recv_seq_id_ = item.SeqId();
+  STREAMING_LOG(DEBUG) << "ReaderQueue::OnData seq_id: " << last_recv_seq_id_;
+
+  Push(item);
+  expect_seq_id_++;
+}
+
+}  // namespace streaming
+}  // namespace ray