a

Signed-off-by: Calvin Neo <[email protected]>
pingcap · Aug 28, 2024 · 9aa9d0e · 9aa9d0e
1 parent f2e5fb8
commit 9aa9d0e
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 0 deletions.
diff --git a/proxy_components/proxy_server/src/status_server/mod.rs b/proxy_components/proxy_server/src/status_server/mod.rs
@@ -206,6 +206,20 @@ where
             None => 60,
         };
 
+        let enable_period = match query_pairs.get("period") {
+            Some(val) => match val.parse() {
+                Ok(val) => val,
+                Err(err) => return Ok(make_response(StatusCode::BAD_REQUEST, err.to_string())),
+            },
+            None => 0,
+        };
+
+        if enable_period == 0 {
+            let msg = "set prof.active = true";
+            set_prof_active(true);
+            return Ok(make_response(StatusCode::OK, msg));
+        }
+
         let interval = Duration::from_secs(interval);
         let period = GLOBAL_TIMER_HANDLE
             .interval(Instant::now() + interval, interval)
@@ -226,11 +240,24 @@ where
     }
 
     fn deactivate_heap_prof(_req: Request<Body>) -> hyper::Result<Response<Body>> {
+        let disable_period = match query_pairs.get("period") {
+            Some(val) => match val.parse() {
+                Ok(val) => val,
+                Err(err) => return Ok(make_response(StatusCode::BAD_REQUEST, err.to_string())),
+            },
+            None => 1,
+        };
+
         let body = if deactivate_heap_profile() {
             "deactivate heap profile success"
         } else {
             "no heap profile is running"
         };
+
+        if disable_period == 1 {
+            set_prof_active(false);
+        }
+
         Ok(make_response(StatusCode::OK, body))
     }
 

diff --git a/proxy_components/proxy_server/src/status_server/profile.rs b/proxy_components/proxy_server/src/status_server/profile.rs
@@ -127,6 +127,19 @@ where
     ProfileGuard::new(on_start, on_end, end.boxed())?.await
 }
 
+pub fn set_prof_active(val: bool) -> Result<(), String> {
+    let activate = has_activate_prof();
+    if activate == val {
+        return Ok(());
+    }
+    if val {
+        activate_prof().map_err(|e| format!("activate_prof: {}", e))?;
+    } else {
+        deactivate_prof().map_err(|e| format!("deactivate_prof: {}", e))?;
+    }
+    Ok(())
+}
+
 /// Activate heap profile and call `callback` if successfully.
 /// `deactivate_heap_profile` can only be called after it's notified from
 /// `callback`.