chore(BOUN): Add debugging endpoints to certificate-orchestrator (#2501)

This change adds two canister methods: - list_registrations - list_tasks which allow to examine the queue contents and the existing registrations.
dfinity · Nov 13, 2024 · 1c9f049 · 1c9f049
1 parent b477d59
commit 1c9f049
Show file tree

Hide file tree

Showing 5 changed files with 334 additions and 4 deletions.
diff --git a/rs/boundary_node/certificate_issuance/certificate_orchestrator/interface.did b/rs/boundary_node/certificate_issuance/certificate_orchestrator/interface.did
@@ -88,6 +88,16 @@ type RemoveRegistrationResponse = variant {
     Err: RemoveRegistrationError;
 };
 
+type ListRegistrationsError = variant {
+    Unauthorized;
+    UnexpectedError: text;
+};
+
+type ListRegistrationsResponse = variant {
+    Ok: vec record { text; Registration };
+    Err: ListRegistrationsError;
+};
+
 type GetCertificateError = variant {
     NotFound;
     Unauthorized;
@@ -155,6 +165,16 @@ type PeekTaskResponse = variant {
     Err: PeekTaskError;
 };
 
+type ListTasksError = variant {
+    Unauthorized;
+    UnexpectedError: text;
+};
+
+type ListTasksResponse = variant {
+    Ok: vec record { text; nat64; Registration };
+    Err: ListTasksError;
+};
+
 type DispenseTaskError = variant {
     NoTasksAvailable;
     Unauthorized;
@@ -166,6 +186,17 @@ type DispenseTaskResponse = variant {
     Err: DispenseTaskError;
 };
 
+type RemoveTaskError = variant {
+    NotFound;
+    Unauthorized;
+    UnexpectedError: text;
+};
+
+type RemoveTaskResponse = variant {
+    Ok;
+    Err: RemoveTaskError;
+};
+
 type ModifyAllowedPrincipalError = variant {
     Unauthorized;
     UnexpectedError: text;
@@ -207,6 +238,7 @@ service: (InitArg) -> {
     getRegistration: (Id) -> (GetRegistrationResponse) query;
     updateRegistration: (Id, UpdateType) -> (UpdateRegistrationResponse);
     removeRegistration: (Id) -> (RemoveRegistrationResponse);
+    listRegistrations: () -> (ListRegistrationsResponse) query;
 
     // Certificates
     getCertificate: (Id) -> (GetCertificateResponse) query;
@@ -218,7 +250,9 @@ service: (InitArg) -> {
     // Tasks
     queueTask: (Id, Timestamp) -> (QueueTaskResponse);
     dispenseTask: () -> (DispenseTaskResponse);
+    removeTask: (Id) -> (RemoveTaskResponse);
     peekTask: () -> (PeekTaskResponse) query;
+    listTasks: () -> (ListTasksResponse) query;
 
     // Metrics (Http Interface)
     http_request: (HttpRequest) -> (HttpResponse) query;

diff --git a/rs/boundary_node/certificate_issuance/certificate_orchestrator/src/main.rs b/rs/boundary_node/certificate_issuance/certificate_orchestrator/src/main.rs
@@ -7,10 +7,12 @@ use certificate_orchestrator_interface::{
     ExportCertificatesError, ExportCertificatesResponse, ExportPackage, GetCertificateError,
     GetCertificateResponse, GetRegistrationError, GetRegistrationResponse, HeaderField,
     HttpRequest, HttpResponse, Id, InitArg, ListAllowedPrincipalsError,
-    ListAllowedPrincipalsResponse, ModifyAllowedPrincipalError, ModifyAllowedPrincipalResponse,
+    ListAllowedPrincipalsResponse, ListRegistrationsError, ListRegistrationsResponse,
+    ListTasksError, ListTasksResponse, ModifyAllowedPrincipalError, ModifyAllowedPrincipalResponse,
     Name, PeekTaskError, PeekTaskResponse, QueueTaskError, QueueTaskResponse, Registration,
-    RemoveRegistrationError, RemoveRegistrationResponse, State, UpdateRegistrationError,
-    UpdateRegistrationResponse, UpdateType, UploadCertificateError, UploadCertificateResponse,
+    RemoveRegistrationError, RemoveRegistrationResponse, RemoveTaskError, RemoveTaskResponse,
+    State, UpdateRegistrationError, UpdateRegistrationResponse, UpdateType, UploadCertificateError,
+    UploadCertificateResponse,
 };
 use ic_cdk::{
     api::{id, time},
@@ -24,7 +26,7 @@ use ic_stable_structures::{
 };
 use priority_queue::PriorityQueue;
 use prometheus::{CounterVec, Encoder, Gauge, GaugeVec, Opts, Registry, TextEncoder};
-use work::{Peek, PeekError};
+use work::{Peek, PeekError, TaskRemover};
 
 use crate::{
     acl::{Authorize, AuthorizeError, Authorizer, WithAuthorize},
@@ -148,6 +150,13 @@ thread_local! {
         ), &["status"]).unwrap()
     });
 
+    static COUNTER_REMOVE_TASK_TOTAL: RefCell<CounterVec> = RefCell::new({
+        CounterVec::new(Opts::new(
+            format!("{SERVICE_NAME}_remove_task_total"), // name
+            "number of times remove_task was called", // help
+        ), &["status"]).unwrap()
+    });
+
     static GAUGE_REGISTRATIONS_TOTAL: RefCell<GaugeVec> = RefCell::new({
         GaugeVec::new(Opts::new(
             format!("{SERVICE_NAME}_registrations_total"), // name
@@ -214,6 +223,11 @@ thread_local! {
             r.register(c).unwrap();
         });
 
+        COUNTER_REMOVE_TASK_TOTAL.with(|c| {
+            let c = Box::new(c.borrow().to_owned());
+            r.register(c).unwrap();
+        });
+
         GAUGE_REGISTRATIONS_TOTAL.with(|g| {
             let g = Box::new(g.borrow().to_owned());
             r.register(g).unwrap();
@@ -344,6 +358,12 @@ thread_local! {
         let r = WithMetrics(r, &COUNTER_REMOVE_REGISTRATION_TOTAL);
         Box::new(r)
     });
+
+    static REGISTRATION_LISTER: RefCell<Box<dyn registration::List>> = RefCell::new({
+        let v = registration::Lister::new(&REGISTRATIONS);
+        let v = WithAuthorize(v, &ROOT_AUTHORIZER);
+        Box::new(v)
+    });
 }
 
 // Certificates
@@ -386,12 +406,25 @@ thread_local! {
         Box::new(d)
     });
 
+    static TASK_LISTER: RefCell<Box<dyn work::List>> = RefCell::new({
+        let v = work::Lister::new(&TASKS, &REGISTRATIONS);
+        let v = WithAuthorize(v, &ROOT_AUTHORIZER);
+        Box::new(v)
+    });
+
     static DISPENSER: RefCell<Box<dyn Dispense>> = RefCell::new({
         let d = Dispenser::new(&TASKS, &RETRIES);
         let d = WithAuthorize(d, &MAIN_AUTHORIZER);
         let d = WithMetrics(d, &COUNTER_DISPENSE_TASK_TOTAL);
         Box::new(d)
     });
+
+    static TASK_REMOVER: RefCell<Box<dyn work::Remove>> = RefCell::new({
+        let v = TaskRemover::new(&TASKS);
+        let v = WithAuthorize(v, &ROOT_AUTHORIZER);
+        let v = WithMetrics(v, &COUNTER_REMOVE_TASK_TOTAL);
+        Box::new(v)
+    });
 }
 
 // Expirations and retries
@@ -678,6 +711,20 @@ fn remove_registration(id: Id) -> RemoveRegistrationResponse {
     }
 }
 
+#[query(name = "listRegistrations")]
+#[candid_method(query, rename = "listRegistrations")]
+fn list_registrations() -> ListRegistrationsResponse {
+    match REGISTRATION_LISTER.with(|v| v.borrow().list()) {
+        Ok(rs) => ListRegistrationsResponse::Ok(rs),
+        Err(err) => ListRegistrationsResponse::Err(match err {
+            registration::ListError::Unauthorized => ListRegistrationsError::Unauthorized,
+            registration::ListError::UnexpectedError(err) => {
+                ListRegistrationsError::UnexpectedError(err.to_string())
+            }
+        }),
+    }
+}
+
 // Certificates
 
 #[query(name = "getCertificate")]
@@ -798,6 +845,35 @@ fn dispense_task() -> DispenseTaskResponse {
     }
 }
 
+#[update(name = "removeTask")]
+#[candid_method(update, rename = "removeTask")]
+fn remove_task(id: Id) -> RemoveTaskResponse {
+    match TASK_REMOVER.with(|v| v.borrow().remove(&id)) {
+        Ok(()) => RemoveTaskResponse::Ok,
+        Err(err) => RemoveTaskResponse::Err(match err {
+            work::RemoveError::NotFound => RemoveTaskError::NotFound,
+            work::RemoveError::Unauthorized => RemoveTaskError::Unauthorized,
+            work::RemoveError::UnexpectedError(err) => {
+                RemoveTaskError::UnexpectedError(err.to_string())
+            }
+        }),
+    }
+}
+
+#[query(name = "listTasks")]
+#[candid_method(query, rename = "listTasks")]
+fn list_tasks() -> ListTasksResponse {
+    match TASK_LISTER.with(|v| v.borrow().list()) {
+        Ok(ts) => ListTasksResponse::Ok(ts),
+        Err(err) => ListTasksResponse::Err(match err {
+            work::ListError::Unauthorized => ListTasksError::Unauthorized,
+            work::ListError::UnexpectedError(err) => {
+                ListTasksError::UnexpectedError(err.to_string())
+            }
+        }),
+    }
+}
+
 // Metrics
 
 #[query(name = "http_request")]

diff --git a/rs/boundary_node/certificate_issuance/certificate_orchestrator/src/registration.rs b/rs/boundary_node/certificate_issuance/certificate_orchestrator/src/registration.rs
@@ -380,6 +380,53 @@ impl<T: Update> Update for WithMetrics<T> {
     }
 }
 
+#[derive(Debug, thiserror::Error)]
+pub enum ListError {
+    #[error("Unauthorized")]
+    Unauthorized,
+
+    #[error(transparent)]
+    UnexpectedError(#[from] anyhow::Error),
+}
+
+pub trait List {
+    fn list(&self) -> Result<Vec<(String, Registration)>, ListError>;
+}
+
+pub struct Lister {
+    registrations: LocalRef<StableMap<StorableId, Registration>>,
+}
+
+impl Lister {
+    pub fn new(registrations: LocalRef<StableMap<StorableId, Registration>>) -> Self {
+        Self { registrations }
+    }
+}
+
+impl List for Lister {
+    fn list(&self) -> Result<Vec<(String, Registration)>, ListError> {
+        Ok(self.registrations.with(|rs| {
+            rs.borrow()
+                .iter()
+                .map(|(id, r)| (id.to_string(), r))
+                .collect()
+        }))
+    }
+}
+
+impl<T: List, A: Authorize> List for WithAuthorize<T, A> {
+    fn list(&self) -> Result<Vec<(String, Registration)>, ListError> {
+        if let Err(err) = self.1.authorize(&caller()) {
+            return Err(match err {
+                AuthorizeError::Unauthorized => ListError::Unauthorized,
+                AuthorizeError::UnexpectedError(err) => ListError::UnexpectedError(err),
+            });
+        };
+
+        self.0.list()
+    }
+}
+
 #[derive(Debug, thiserror::Error)]
 pub enum RemoveError {
     #[error("Not found")]