ray/streaming/src/queue/queue_handler.h

#ifndef _QUEUE_SERVICE_H_
#define _QUEUE_SERVICE_H_

#include <boost/asio.hpp>
#include <boost/bind.hpp>
#include <boost/thread.hpp>
#include <thread>

#include "queue.h"
#include "util/streaming_logging.h"

namespace ray {
namespace streaming {

/// Base class of UpstreamQueueMessageHandler and DownstreamQueueMessageHandler.
/// A queue service manages a group of queues, upstream queues or downstream queues of
/// the current actor. Each queue service holds a boost.asio io_service, to handle
/// messages asynchronously. When a message received by Writer/Reader in ray call thread,
/// the message was delivered to
/// UpstreamQueueMessageHandler/DownstreamQueueMessageHandler, then the ray call thread
/// returns immediately. The queue service parses meta infomation from the message,
/// including queue_id actor_id, etc, and dispatchs message to queue according to
/// queue_id.
class QueueMessageHandler {
 public:
  /// Construct a QueueMessageHandler instance.
  /// \param[in] core_worker CoreWorker C++ pointer of current actor, used to call Core
  /// Worker's api.
  ///            For Python worker, the pointer can be obtained from
  ///            ray.worker.global_worker.core_worker; For Java worker, obtained from
  ///            RayNativeRuntime object through java reflection.
  /// \param[in] actor_id actor id of current actor.
  QueueMessageHandler(CoreWorker *core_worker, const ActorID &actor_id)
      : core_worker_(core_worker),
        actor_id_(actor_id),
        queue_dummy_work_(queue_service_) {
    Start();
  }

  virtual ~QueueMessageHandler() { Stop(); }

  /// Dispatch message buffer to asio service.
  /// \param[in] buffer serialized message received from peer actor.
  void DispatchMessageAsync(std::shared_ptr<LocalMemoryBuffer> buffer);

  /// Dispatch message buffer to asio service synchronously, and wait for handle result.
  /// \param[in] buffer serialized message received from peer actor.
  /// \return handle result.
  std::shared_ptr<LocalMemoryBuffer> DispatchMessageSync(
      std::shared_ptr<LocalMemoryBuffer> buffer);

  /// Get transport to a peer actor specified by actor_id.
  /// \param[in] actor_id actor id of peer actor
  /// \return transport
  std::shared_ptr<Transport> GetOutTransport(const ObjectID &actor_id);

  /// The actual function where message being dispatched, called by DispatchMessageAsync
  /// and DispatchMessageSync.
  /// \param[in] buffer serialized message received from peer actor.
  /// \param[in] callback the callback function used by DispatchMessageSync, called
  ///            after message processed complete. The std::shared_ptr<LocalMemoryBuffer>
  ///            parameter is the return value.
  virtual void DispatchMessageInternal(
      std::shared_ptr<LocalMemoryBuffer> buffer,
      std::function<void(std::shared_ptr<LocalMemoryBuffer>)> callback) = 0;

  /// Save actor_id of the peer actor specified by queue_id. For a upstream queue, the
  /// peer actor refer specifically to the actor in current ray cluster who has a
  /// downstream queue with same queue_id, and vice versa.
  /// \param[in] queue_id queue id of current queue.
  /// \param[in] actor_id actor_id actor id of corresponded peer actor.
  void SetPeerActorID(const ObjectID &queue_id, const ActorID &actor_id);

  /// Obtain the actor id of the peer actor specified by queue_id.
  /// \return actor id
  ActorID GetPeerActorID(const ObjectID &queue_id);

  /// Release all queues in current queue service.
  void Release();

 private:
  /// Start asio service
  void Start();
  /// Stop asio service
  void Stop();
  /// The callback function of internal thread.
  void QueueThreadCallback() { queue_service_.run(); }

 protected:
  /// CoreWorker C++ pointer of current actor
  CoreWorker *core_worker_;
  /// actor_id actor id of current actor
  ActorID actor_id_;
  /// Helper function, parse message buffer to Message object.
  std::shared_ptr<Message> ParseMessage(std::shared_ptr<LocalMemoryBuffer> buffer);

 private:
  /// Map from queue id to a actor id of the queue's peer actor.
  std::unordered_map<ObjectID, ActorID> actors_;
  /// Map from queue id to a transport of the queue's peer actor.
  std::unordered_map<ObjectID, std::shared_ptr<Transport>> out_transports_;
  /// The internal thread which asio service run with.
  std::thread queue_thread_;
  /// The internal asio service.
  boost::asio::io_service queue_service_;
  /// The asio work which keeps queue_service_ alive.
  boost::asio::io_service::work queue_dummy_work_;
};

/// UpstreamQueueMessageHandler holds and manages all upstream queues of current actor.
class UpstreamQueueMessageHandler : public QueueMessageHandler {
 public:
  /// Construct a UpstreamQueueMessageHandler instance.
  UpstreamQueueMessageHandler(CoreWorker *core_worker, const ActorID &actor_id)
      : QueueMessageHandler(core_worker, actor_id) {}
  /// Create a upstream queue.
  /// \param[in] queue_id queue id of the queue to be created.
  /// \param[in] peer_actor_id actor id of peer actor.
  /// \param[in] size the max memory size of the queue.
  std::shared_ptr<WriterQueue> CreateUpstreamQueue(const ObjectID &queue_id,
                                                   const ActorID &peer_actor_id,
                                                   uint64_t size);
  /// Check whether the upstream queue specified by queue_id exists or not.
  bool UpstreamQueueExists(const ObjectID &queue_id);
  /// Wait all queues in queue_ids vector ready, until timeout.
  /// \param[in] queue_ids a group of queues.
  /// \param[in] timeout_ms max timeout time interval for wait all queues.
  /// \param[out] failed_queues a group of queues which are not ready when timeout.
  void WaitQueues(const std::vector<ObjectID> &queue_ids, int64_t timeout_ms,
                  std::vector<ObjectID> &failed_queues);
  /// Handle notify message from corresponded downstream queue.
  void OnNotify(std::shared_ptr<NotificationMessage> notify_msg);
  /// Obtain upstream queue specified by queue_id.
  std::shared_ptr<streaming::WriterQueue> GetUpQueue(const ObjectID &queue_id);
  /// Release all upstream queues
  void ReleaseAllUpQueues();

  virtual void DispatchMessageInternal(
      std::shared_ptr<LocalMemoryBuffer> buffer,
      std::function<void(std::shared_ptr<LocalMemoryBuffer>)> callback) override;

  static std::shared_ptr<UpstreamQueueMessageHandler> CreateService(
      CoreWorker *core_worker, const ActorID &actor_id);
  static std::shared_ptr<UpstreamQueueMessageHandler> GetService();

  static RayFunction peer_sync_function_;
  static RayFunction peer_async_function_;

 private:
  bool CheckQueueSync(const ObjectID &queue_ids);

 private:
  std::unordered_map<ObjectID, std::shared_ptr<streaming::WriterQueue>> upstream_queues_;
  static std::shared_ptr<UpstreamQueueMessageHandler> upstream_handler_;
};

/// UpstreamQueueMessageHandler holds and manages all downstream queues of current actor.
class DownstreamQueueMessageHandler : public QueueMessageHandler {
 public:
  DownstreamQueueMessageHandler(CoreWorker *core_worker, const ActorID &actor_id)
      : QueueMessageHandler(core_worker, actor_id) {}
  std::shared_ptr<ReaderQueue> CreateDownstreamQueue(const ObjectID &queue_id,
                                                     const ActorID &peer_actor_id);
  bool DownstreamQueueExists(const ObjectID &queue_id);

  void UpdateDownActor(const ObjectID &queue_id, const ActorID &actor_id);

  std::shared_ptr<LocalMemoryBuffer> OnCheckQueue(
      std::shared_ptr<CheckMessage> check_msg);

  std::shared_ptr<streaming::ReaderQueue> GetDownQueue(const ObjectID &queue_id);

  void ReleaseAllDownQueues();

  void OnData(std::shared_ptr<DataMessage> msg);
  virtual void DispatchMessageInternal(
      std::shared_ptr<LocalMemoryBuffer> buffer,
      std::function<void(std::shared_ptr<LocalMemoryBuffer>)> callback);

  static std::shared_ptr<DownstreamQueueMessageHandler> CreateService(
      CoreWorker *core_worker, const ActorID &actor_id);
  static std::shared_ptr<DownstreamQueueMessageHandler> GetService();
  static RayFunction peer_sync_function_;
  static RayFunction peer_async_function_;

 private:
  std::unordered_map<ObjectID, std::shared_ptr<streaming::ReaderQueue>>
      downstream_queues_;
  static std::shared_ptr<DownstreamQueueMessageHandler> downstream_handler_;
};

}  // namespace streaming
}  // namespace ray
#endif