From 64dc6f21a9dd3d314eba5bc3fac85f5f1b49ccab Mon Sep 17 00:00:00 2001
From: Chenyang Wu <hpghjfs@yeah.net>
Date: Sun, 20 Mar 2022 16:00:32 +0800
Subject: [PATCH] fix `timeout` and add `timer`

---
 src/AdaOPS.jl  | 9 +++++++--
 src/planner.jl | 8 ++++----
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/AdaOPS.jl b/src/AdaOPS.jl
index c8dfabb..2ddd83e 100644
--- a/src/AdaOPS.jl
+++ b/src/AdaOPS.jl
@@ -81,7 +81,7 @@ parameters match the definitions in the paper exactly.
 Further information can be found in the field docstrings (e.g.
 `?AdaOPSSolver.xi`)
 """
-@with_kw struct AdaOPSSolver{N, R<:AbstractRNG} <: Solver
+@with_kw struct AdaOPSSolver{N, R<:AbstractRNG, T} <: Solver
     "The target gap between the upper and the lower bound at the root of the AdaOPS tree."
     epsilon_0::Float64                      = 1e-3
 
@@ -140,7 +140,12 @@ Further information can be found in the field docstrings (e.g.
     tree_in_info::Bool                      = false
 
     "Issue an warning when the planning time surpass the time limit by `timeout_warning_threshold` times"
-    timeout_warning_threshold::Float64     = T_max * 2.0
+    timeout_warning_threshold::Float64      = T_max * 2.0
+
+    "Search iterations end when `timer() - start_time ≥ T_max`.
+    Default to `() -> 1e-9*time_ns()`, which gets the current system time.
+    Switch to `() -> 1e-6*CPUtime_us()` when the CPU time is concerned."
+    timer::T                                = () -> 1e-9*time_ns()
 
     "Number of pre-allocated belief nodes"
     num_b::Int                              = 50_000
diff --git a/src/planner.jl b/src/planner.jl
index 56f680b..0eb4cdd 100644
--- a/src/planner.jl
+++ b/src/planner.jl
@@ -2,20 +2,20 @@ function build_tree(p::AdaOPSPlanner, b0)
     D = AdaOPSTree(p, b0)
     b = 1
     trial = 1
-    start = CPUtime_us()
 
     Depth = sizehint!(Int[], 2000)
     sol = solver(p)
+    start = sol.timer()
     while D.u[1]-D.l[1] > sol.epsilon_0 &&
-          CPUtime_us()-start < sol.T_max*1e6 &&
+          sol.timer()-start < sol.T_max &&
           trial <= sol.max_trials
         push!(Depth, explore!(D, 1, p))
         trial += 1
     end
-    if (CPUtime_us()-start)*1e-6 > sol.T_max*sol.timeout_warning_threshold
+    if sol.timer()-start > sol.timeout_warning_threshold
         @warn(@sprintf("Surpass the time limit. The actual runtime is %3.1fs.
         Hyperparameters: delta=%4.2f, m_min=%3d, m_max=%3d, zeta=%4.2f, grid=%s, bounds=%s",
-        (CPUtime_us()-start)*1e-6, sol.delta, sol.m_min, sol.m_max, sol.zeta, typeof(sol.grid), typeof(sol.bounds)))
+        sol.timer()-start, sol.delta, sol.m_min, sol.m_max, sol.zeta, typeof(sol.grid), typeof(sol.bounds)))
         info_analysis(Dict(:tree=>D, :depth=>Depth))
     end
     return D, Depth