From 5349e63c2b01aaaa3f5e08ac6d72d05b8c8dfac6 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Sun, 15 Aug 2021 11:15:14 +0300
Subject: [PATCH] Add Timeout::get_timeout.

---
 tdactor/td/actor/Timeout.h             | 3 +++
 tdactor/td/actor/impl/Actor-decl.h     | 1 +
 tdactor/td/actor/impl/Actor.h          | 3 +++
 tdactor/td/actor/impl/Scheduler-decl.h | 2 ++
 tdactor/td/actor/impl/Scheduler.cpp    | 6 ++++++
 tdactor/td/actor/impl/Scheduler.h      | 3 +++
 tdutils/td/utils/Heap.h                | 6 ++++++
 7 files changed, 24 insertions(+)

diff --git a/tdactor/td/actor/Timeout.h b/tdactor/td/actor/Timeout.h
index de4216546..63f6cfdf2 100644
--- a/tdactor/td/actor/Timeout.h
+++ b/tdactor/td/actor/Timeout.h
@@ -35,6 +35,9 @@ class Timeout final : public Actor {
   bool has_timeout() const {
     return Actor::has_timeout();
   }
+  double get_timeout() const {
+    return Actor::get_timeout();
+  }
   void set_timeout_in(double timeout) {
     Actor::set_timeout_in(timeout);
   }
diff --git a/tdactor/td/actor/impl/Actor-decl.h b/tdactor/td/actor/impl/Actor-decl.h
index d6eabd2be..8a107375f 100644
--- a/tdactor/td/actor/impl/Actor-decl.h
+++ b/tdactor/td/actor/impl/Actor-decl.h
@@ -67,6 +67,7 @@ class Actor : public ObserverBase {
   void stop();
   void do_stop();
   bool has_timeout() const;
+  double get_timeout() const;
   void set_timeout_in(double timeout_in);
   void set_timeout_at(double timeout_at);
   void cancel_timeout();
diff --git a/tdactor/td/actor/impl/Actor.h b/tdactor/td/actor/impl/Actor.h
index 94fd9c065..3577093ce 100644
--- a/tdactor/td/actor/impl/Actor.h
+++ b/tdactor/td/actor/impl/Actor.h
@@ -54,6 +54,9 @@ inline void Actor::do_stop() {
 inline bool Actor::has_timeout() const {
   return Scheduler::instance()->has_actor_timeout(this);
 }
+inline double Actor::get_timeout() const {
+  return Scheduler::instance()->get_actor_timeout(this);
+}
 inline void Actor::set_timeout_in(double timeout_in) {
   Scheduler::instance()->set_actor_timeout_in(this, timeout_in);
 }
diff --git a/tdactor/td/actor/impl/Scheduler-decl.h b/tdactor/td/actor/impl/Scheduler-decl.h
index add1d702d..2b17d90ed 100644
--- a/tdactor/td/actor/impl/Scheduler-decl.h
+++ b/tdactor/td/actor/impl/Scheduler-decl.h
@@ -126,6 +126,7 @@ class Scheduler {
   void finish_migrate_actor(Actor *actor);
 
   bool has_actor_timeout(const Actor *actor) const;
+  double get_actor_timeout(const Actor *actor) const;
   void set_actor_timeout_in(Actor *actor, double timeout);
   void set_actor_timeout_at(Actor *actor, double timeout_at);
   void cancel_actor_timeout(Actor *actor);
@@ -176,6 +177,7 @@ class Scheduler {
   void start_migrate_actor(ActorInfo *actor_info, int32 dest_sched_id);
 
   bool has_actor_timeout(const ActorInfo *actor_info) const;
+  double get_actor_timeout(const ActorInfo *actor_info) const;
   void set_actor_timeout_in(ActorInfo *actor_info, double timeout);
   void set_actor_timeout_at(ActorInfo *actor_info, double timeout_at);
   void cancel_actor_timeout(ActorInfo *actor_info);
diff --git a/tdactor/td/actor/impl/Scheduler.cpp b/tdactor/td/actor/impl/Scheduler.cpp
index c5c5c88a5..3ae5a2468 100644
--- a/tdactor/td/actor/impl/Scheduler.cpp
+++ b/tdactor/td/actor/impl/Scheduler.cpp
@@ -392,6 +392,7 @@ void Scheduler::do_migrate_actor(ActorInfo *actor_info, int32 dest_sched_id) {
 void Scheduler::start_migrate_actor(Actor *actor, int32 dest_sched_id) {
   start_migrate_actor(actor->get_info(), dest_sched_id);
 }
+
 void Scheduler::start_migrate_actor(ActorInfo *actor_info, int32 dest_sched_id) {
   VLOG(actor) << "Start migrate actor: " << tag("name", actor_info) << tag("ptr", actor_info)
               << tag("actor_count", actor_count_);
@@ -406,6 +407,11 @@ void Scheduler::start_migrate_actor(ActorInfo *actor_info, int32 dest_sched_id)
   cancel_actor_timeout(actor_info);
 }
 
+double Scheduler::get_actor_timeout(const ActorInfo *actor_info) const {
+  const HeapNode *heap_node = actor_info->get_heap_node();
+  return heap_node->in_heap() ? timeout_queue_.get_key(heap_node) - Time::now() : 0.0;
+}
+
 void Scheduler::set_actor_timeout_in(ActorInfo *actor_info, double timeout) {
   if (timeout > 1e10) {
     timeout = 1e10;
diff --git a/tdactor/td/actor/impl/Scheduler.h b/tdactor/td/actor/impl/Scheduler.h
index 0d432ca2f..1212509c0 100644
--- a/tdactor/td/actor/impl/Scheduler.h
+++ b/tdactor/td/actor/impl/Scheduler.h
@@ -302,6 +302,9 @@ inline void Scheduler::finish_migrate_actor(Actor *actor) {
 inline bool Scheduler::has_actor_timeout(const Actor *actor) const {
   return has_actor_timeout(actor->get_info());
 }
+inline double Scheduler::get_actor_timeout(const Actor *actor) const {
+  return get_actor_timeout(actor->get_info());
+}
 inline void Scheduler::set_actor_timeout_in(Actor *actor, double timeout) {
   set_actor_timeout_in(actor->get_info(), timeout);
 }
diff --git a/tdutils/td/utils/Heap.h b/tdutils/td/utils/Heap.h
index 697b8b5f8..c9bbbd901 100644
--- a/tdutils/td/utils/Heap.h
+++ b/tdutils/td/utils/Heap.h
@@ -37,6 +37,12 @@ class KHeap {
     return array_[0].key_;
   }
 
+  KeyT get_key(const HeapNode *node) const {
+    size_t pos = static_cast<size_t>(node->pos_);
+    CHECK(pos < array_.size());
+    return array_[pos].key_;
+  }
+
   const HeapNode *top() const {
     return array_[0].node_;
   }