Integrate metric items into raylet

src/ray/gcs/redis_context.cc

@@ -4,6 +4,9 @@
 
 #include <sstream>
 
+#include "ray/stats/stats.h"
+#include "ray/util/util.h"
+
 extern "C" {
 #include "ray/thirdparty/hiredis/adapters/ae.h"
 #include "ray/thirdparty/hiredis/async.h"
@@ -18,13 +21,21 @@ namespace {
 /// A helper function to call the callback and delete it from the callback
 /// manager if necessary.
 void ProcessCallback(int64_t callback_index, const std::string &data) {
-  if (callback_index >= 0) {
-    bool delete_callback =
-        ray::gcs::RedisCallbackManager::instance().get(callback_index)(data);
-    // Delete the callback if necessary.
-    if (delete_callback) {
-      ray::gcs::RedisCallbackManager::instance().remove(callback_index);
-    }
+  RAY_CHECK(callback_index >= 0) << "The callback index must be greater than 0, "
+                                 << "but it actually is " << callback_index;
+  auto callback_item = ray::gcs::RedisCallbackManager::instance().get(callback_index);
+  if (!callback_item.is_subscription) {
+    // Record the redis latency for non-subscription redis operations.
+    auto end_time = current_sys_time_us();
+    ray::stats::RedisLatency().Record(end_time - callback_item.start_time);
+  }
+  // Invoke the callback.
+  if (callback_item.callback != nullptr) {
+    callback_item.callback(data);
+  }
+  if (!callback_item.is_subscription) {
+    // Delete the callback if it's not a subscription callback.
+    ray::gcs::RedisCallbackManager::instance().remove(callback_index);
   }
 }
 
@@ -104,18 +115,20 @@ void SubscribeRedisCallback(void *c, void *r, void *privdata) {
   ProcessCallback(callback_index, data);
 }
 
-int64_t RedisCallbackManager::add(const RedisCallback &function) {
-  callbacks_.emplace(num_callbacks_, function);
+int64_t RedisCallbackManager::add(const RedisCallback &function, bool is_subscription) {
+  auto start_time = current_sys_time_us();
+  callback_items_.emplace(num_callbacks_,
+                          CallbackItem(function, is_subscription, start_time));
   return num_callbacks_++;
 }
 
-RedisCallback &RedisCallbackManager::get(int64_t callback_index) {
-  RAY_CHECK(callbacks_.find(callback_index) != callbacks_.end());
-  return callbacks_[callback_index];
+RedisCallbackManager::CallbackItem &RedisCallbackManager::get(int64_t callback_index) {
+  RAY_CHECK(callback_items_.find(callback_index) != callback_items_.end());
+  return callback_items_[callback_index];
 }
 
 void RedisCallbackManager::remove(int64_t callback_index) {
-  callbacks_.erase(callback_index);
+  callback_items_.erase(callback_index);
 }
 
 #define REDIS_CHECK_ERROR(CONTEXT, REPLY)                     \
@@ -217,8 +230,7 @@ Status RedisContext::RunAsync(const std::string &command, const UniqueID &id,
                               const uint8_t *data, int64_t length,
                               const TablePrefix prefix, const TablePubsub pubsub_channel,
                               RedisCallback redisCallback, int log_length) {
-  int64_t callback_index =
-      redisCallback != nullptr ? RedisCallbackManager::instance().add(redisCallback) : -1;
+  int64_t callback_index = RedisCallbackManager::instance().add(redisCallback, false);
   if (length > 0) {
     if (log_length >= 0) {
       std::string redis_command = command + " %d %d %b %b %d";
@@ -278,7 +290,7 @@ Status RedisContext::SubscribeAsync(const ClientID &client_id,
   RAY_CHECK(pubsub_channel != TablePubsub::NO_PUBLISH)
       << "Client requested subscribe on a table that does not support pubsub";
 
-  int64_t callback_index = RedisCallbackManager::instance().add(redisCallback);
+  int64_t callback_index = RedisCallbackManager::instance().add(redisCallback, true);
   RAY_CHECK(out_callback_index != nullptr);
   *out_callback_index = callback_index;
   int status = 0;

src/ray/gcs/redis_context.h

@@ -19,9 +19,8 @@ namespace ray {
 
 namespace gcs {
 /// Every callback should take in a vector of the results from the Redis
-/// operation and return a bool indicating whether the callback should be
-/// deleted once called.
-using RedisCallback = std::function<bool(const std::string &)>;
+/// operation.
+using RedisCallback = std::function<void(const std::string &)>;
 
 class RedisCallbackManager {
  public:
@@ -30,9 +29,24 @@ class RedisCallbackManager {
     return instance;
   }
 
-  int64_t add(const RedisCallback &function);
+  struct CallbackItem {
+    CallbackItem() = default;
 
-  RedisCallback &get(int64_t callback_index);
+    CallbackItem(const RedisCallback &callback, bool is_subscription,
+                 int64_t start_time) {
+      this->callback = callback;
+      this->is_subscription = is_subscription;
+      this->start_time = start_time;
+    }
+
+    RedisCallback callback;
+    bool is_subscription;
+    int64_t start_time;
+  };
+
+  int64_t add(const RedisCallback &function, bool is_subscription);
+
+  CallbackItem &get(int64_t callback_index);
 
   /// Remove a callback.
   void remove(int64_t callback_index);
@@ -43,7 +57,7 @@ class RedisCallbackManager {
   ~RedisCallbackManager() {}
 
   int64_t num_callbacks_ = 0;
-  std::unordered_map<int64_t, RedisCallback> callbacks_;
+  std::unordered_map<int64_t, CallbackItem> callback_items_;
 };
 
 class RedisContext {

src/ray/gcs/tables.cc

@@ -48,7 +48,6 @@ Status Log<ID, Data>::Append(const JobID &job_id, const ID &id,
     if (done != nullptr) {
       (done)(client_, id, *dataT);
     }
-    return true;
   };
   flatbuffers::FlatBufferBuilder fbb;
   fbb.ForceDefaults(true);
@@ -73,7 +72,6 @@ Status Log<ID, Data>::AppendAt(const JobID &job_id, const ID &id,
         (failure)(client_, id, *dataT);
       }
     }
-    return true;
   };
   flatbuffers::FlatBufferBuilder fbb;
   fbb.ForceDefaults(true);
@@ -101,7 +99,6 @@ Status Log<ID, Data>::Lookup(const JobID &job_id, const ID &id, const Callback &
       }
       lookup(client_, id, results);
     }
-    return true;
   };
   std::vector<uint8_t> nil;
   return GetRedisContext(id)->RunAsync("RAY.TABLE_LOOKUP", id, nil.data(), nil.size(),
@@ -153,10 +150,8 @@ Status Log<ID, Data>::Subscribe(const JobID &job_id, const ClientID &client_id,
         subscribe(client_, id, root->notification_mode(), results);
       }
     }
-    // We do not delete the callback after calling it since there may be
-    // more subscription messages.
-    return false;
   };
+
   subscribe_callback_index_ = 1;
   for (auto &context : shard_contexts_) {
     RAY_RETURN_NOT_OK(context->SubscribeAsync(client_id, pubsub_channel_, callback,
@@ -229,7 +224,6 @@ Status Table<ID, Data>::Add(const JobID &job_id, const ID &id,
     if (done != nullptr) {
       (done)(client_, id, *dataT);
     }
-    return true;
   };
   flatbuffers::FlatBufferBuilder fbb;
   fbb.ForceDefaults(true);
@@ -295,7 +289,6 @@ Status Set<ID, Data>::Add(const JobID &job_id, const ID &id,
     if (done != nullptr) {
       (done)(client_, id, *dataT);
     }
-    return true;
   };
   flatbuffers::FlatBufferBuilder fbb;
   fbb.ForceDefaults(true);
@@ -313,7 +306,6 @@ Status Set<ID, Data>::Remove(const JobID &job_id, const ID &id,
     if (done != nullptr) {
       (done)(client_, id, *dataT);
     }
-    return true;
   };
   flatbuffers::FlatBufferBuilder fbb;
   fbb.ForceDefaults(true);

src/ray/object_manager/connection_pool.cc

@@ -147,4 +147,25 @@ std::string ConnectionPool::DebugString() const {
   return result.str();
 }
 
+void ConnectionPool::RecordMetrics() const {
+  stats::ConnectionPoolStats().Record(
+      message_send_connections_.size(),
+      {{stats::ValueTypeKey, "num_message_send_connections"}});
+  stats::ConnectionPoolStats().Record(
+      transfer_send_connections_.size(),
+      {{stats::ValueTypeKey, "num_transfer_send_connections"}});
+  stats::ConnectionPoolStats().Record(
+      available_transfer_send_connections_.size(),
+      {{stats::ValueTypeKey, "num_avail_message_send_connections"}});
+  stats::ConnectionPoolStats().Record(
+      available_transfer_send_connections_.size(),
+      {{stats::ValueTypeKey, "num_avail_transfer_send_connections"}});
+  stats::ConnectionPoolStats().Record(
+      message_receive_connections_.size(),
+      {{stats::ValueTypeKey, "num_message_receive_connections"}});
+  stats::ConnectionPoolStats().Record(
+      transfer_receive_connections_.size(),
+      {{stats::ValueTypeKey, "num_transfer_receive_connections"}});
+}
+
 }  // namespace ray

src/ray/object_manager/connection_pool.h

@@ -19,6 +19,7 @@
 #include "ray/object_manager/format/object_manager_generated.h"
 #include "ray/object_manager/object_directory.h"
 #include "ray/object_manager/object_manager_client_connection.h"
+#include "ray/stats/stats.h"
 
 namespace asio = boost::asio;
 
@@ -95,6 +96,9 @@ class ConnectionPool {
   /// \return string.
   std::string DebugString() const;
 
+  /// Record metrics.
+  void RecordMetrics() const;
+
   /// This object cannot be copied for thread-safety.
   RAY_DISALLOW_COPY_AND_ASSIGN(ConnectionPool);
 

src/ray/object_manager/object_manager.cc

@@ -1,5 +1,6 @@
 #include "ray/object_manager/object_manager.h"
 #include "ray/common/common_protocol.h"
+#include "ray/stats/stats.h"
 #include "ray/util/util.h"
 
 namespace asio = boost::asio;
@@ -966,4 +967,19 @@ std::string ObjectManager::DebugString() const {
   return result.str();
 }
 
+void ObjectManager::RecordMetrics() const {
+  stats::ObjectManagerStats().Record(local_objects_.size(),
+                                     {{stats::ValueTypeKey, "num_local_objects"}});
+  stats::ObjectManagerStats().Record(active_wait_requests_.size(),
+                                     {{stats::ValueTypeKey, "num_active_wait_requests"}});
+  stats::ObjectManagerStats().Record(
+      unfulfilled_push_requests_.size(),
+      {{stats::ValueTypeKey, "num_unfulfilled_push_requests"}});
+  stats::ObjectManagerStats().Record(pull_requests_.size(),
+                                     {{stats::ValueTypeKey, "num_pull_requests"}});
+  stats::ObjectManagerStats().Record(profile_events_.size(),
+                                     {{stats::ValueTypeKey, "num_profile_events"}});
+  connection_pool_.RecordMetrics();
+}
+
 }  // namespace ray

src/ray/object_manager/object_manager.h

@@ -187,6 +187,9 @@ class ObjectManager : public ObjectManagerInterface {
   /// \return string.
   std::string DebugString() const;
 
+  /// Record metrics.
+  void RecordMetrics() const;
+
  private:
   friend class TestObjectManager;
 

src/ray/raylet/lineage_cache.cc

@@ -1,4 +1,5 @@
 #include "lineage_cache.h"
+#include "ray/stats/stats.h"
 
 #include <sstream>
 
@@ -471,6 +472,17 @@ std::string LineageCache::DebugString() const {
   return result.str();
 }
 
+void LineageCache::RecordMetrics() const {
+  stats::LineageCacheStats().Record(committed_tasks_.size(),
+                                    {{stats::ValueTypeKey, "num_committed_tasks"}});
+  stats::LineageCacheStats().Record(lineage_.GetChildrenSize(),
+                                    {{stats::ValueTypeKey, "num_children"}});
+  stats::LineageCacheStats().Record(subscribed_tasks_.size(),
+                                    {{stats::ValueTypeKey, "num_subscribed_tasks"}});
+  stats::LineageCacheStats().Record(lineage_.GetEntries().size(),
+                                    {{stats::ValueTypeKey, "num_lineages"}});
+}
+
 }  // namespace raylet
 
 }  // namespace ray

src/ray/raylet/lineage_cache.h

@@ -299,6 +299,9 @@ class LineageCache {
   /// \return string.
   std::string DebugString() const;
 
+  /// Record metrics.
+  void RecordMetrics() const;
+
  private:
   FRIEND_TEST(LineageCacheTest, BarReturnsZeroOnNull);
   /// Flush a task that is in UNCOMMITTED_READY state.