feat(PocketIC): new PocketIC operation to set certified time (#3595)

This PR introduces a new PocketIC operation to set certified time so
that query calls and read state requests are evaluated on a state with
that (certified) time. Using this new operation to set time in the
"live" mode also fixes test
[flakiness](https://github.com/dfinity/ic/actions/runs/12909846684/job/35998504336).

packages/pocket-ic/CHANGELOG.md

@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   If the status of the update call is known, but the update call was submitted by a different caller, then an error is returned.
 - The function `PocketIc::await_call_no_ticks` to await the status of an update call (submitted through an ingress message) becoming known without triggering round execution
   (round execution must be triggered separarely, e.g., on a "live" instance or by separate PocketIC library calls).
+- The function `PocketIc::set_certified_time` to set the current certified time on all subnets of the PocketIC instance.
 
 ### Changed
 - The response types `pocket_ic::WasmResult`, `pocket_ic::UserError`, and `pocket_ic::CallError` are replaced by a single reject response type `pocket_ic::RejectResponse`.

packages/pocket-ic/src/lib.rs

@@ -619,6 +619,13 @@ impl PocketIc {
         runtime.block_on(async { self.pocket_ic.set_time(time).await })
     }
 
+    /// Set the current certified time of the IC, on all subnets.
+    #[instrument(skip(self), fields(instance_id=self.pocket_ic.instance_id, time = ?time))]
+    pub fn set_certified_time(&self, time: SystemTime) {
+        let runtime = self.runtime.clone();
+        runtime.block_on(async { self.pocket_ic.set_certified_time(time).await })
+    }
+
     /// Advance the time on the IC on all subnets by some nanoseconds.
     #[instrument(skip(self), fields(instance_id=self.pocket_ic.instance_id, duration = ?duration))]
     pub fn advance_time(&self, duration: Duration) {

packages/pocket-ic/src/nonblocking.rs

@@ -302,7 +302,7 @@ impl PocketIc {
     #[instrument(skip(self), fields(instance_id=self.instance_id))]
     pub async fn auto_progress(&self) -> Url {
         let now = std::time::SystemTime::now();
-        self.set_time(now).await;
+        self.set_certified_time(now).await;
         let endpoint = "auto_progress";
         let auto_progress_config = AutoProgressConfig {
             artificial_delay_ms: None,
@@ -485,6 +485,22 @@ impl PocketIc {
         .await;
     }
 
+    /// Set the current certified time of the IC, on all subnets.
+    #[instrument(skip(self), fields(instance_id=self.instance_id, time = ?time))]
+    pub async fn set_certified_time(&self, time: SystemTime) {
+        let endpoint = "update/set_certified_time";
+        self.post::<(), _>(
+            endpoint,
+            RawTime {
+                nanos_since_epoch: time
+                    .duration_since(SystemTime::UNIX_EPOCH)
+                    .expect("Time went backwards")
+                    .as_nanos() as u64,
+            },
+        )
+        .await;
+    }
+
     /// Advance the time on the IC on all subnets by some nanoseconds.
     #[instrument(skip(self), fields(instance_id=self.instance_id, duration = ?duration))]
     pub async fn advance_time(&self, duration: Duration) {

packages/pocket-ic/tests/tests.rs

@@ -350,40 +350,6 @@ fn test_multiple_large_xnet_payloads() {
     }
 }
 
-#[test]
-fn test_get_and_set_and_advance_time() {
-    let pic = PocketIc::new();
-
-    let unix_time_secs = 1630328630;
-    let set_time = SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(unix_time_secs);
-    pic.set_time(set_time);
-    assert_eq!(pic.get_time(), set_time);
-    pic.tick();
-    assert_eq!(pic.get_time(), set_time);
-
-    pic.advance_time(std::time::Duration::from_secs(420));
-    assert_eq!(
-        pic.get_time(),
-        set_time + std::time::Duration::from_secs(420)
-    );
-    pic.tick();
-    assert_eq!(
-        pic.get_time(),
-        set_time + std::time::Duration::from_secs(420)
-    );
-}
-
-#[test]
-#[should_panic(expected = "SettingTimeIntoPast")]
-fn set_time_into_past() {
-    let pic = PocketIc::new();
-
-    let now = SystemTime::now();
-    pic.set_time(now + std::time::Duration::from_secs(1));
-
-    pic.set_time(now);
-}
-
 fn query_and_check_time(pic: &PocketIc, test_canister: Principal) {
     let current_time = pic
         .get_time()
@@ -402,25 +368,62 @@ fn query_and_check_time(pic: &PocketIc, test_canister: Principal) {
 }
 
 #[test]
-fn query_call_after_advance_time() {
+fn test_get_and_set_and_advance_time() {
     let pic = PocketIc::new();
 
     // We create a test canister.
     let canister = pic.create_canister();
     pic.add_cycles(canister, INIT_CYCLES);
     pic.install_canister(canister, test_canister_wasm(), vec![], None);
 
+    let unix_time_secs = 1650000000;
+    let time = SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(unix_time_secs);
+    pic.set_time(time);
+    // time is not certified so `query_and_check_time` would fail here
+    assert_eq!(pic.get_time(), time);
+    pic.tick();
     query_and_check_time(&pic, canister);
-
-    pic.advance_time(std::time::Duration::from_secs(420));
+    assert_eq!(pic.get_time(), time);
     pic.tick();
-
     query_and_check_time(&pic, canister);
+    assert_eq!(pic.get_time(), time + std::time::Duration::from_nanos(1));
 
-    pic.advance_time(std::time::Duration::from_secs(0));
+    let unix_time_secs = 1700000000;
+    let time = SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(unix_time_secs);
+    pic.set_certified_time(time);
+    query_and_check_time(&pic, canister);
+    assert_eq!(pic.get_time(), time);
+    pic.tick();
+    query_and_check_time(&pic, canister);
+    assert_eq!(pic.get_time(), time + std::time::Duration::from_nanos(1));
     pic.tick();
+    query_and_check_time(&pic, canister);
+    assert_eq!(pic.get_time(), time + std::time::Duration::from_nanos(2));
 
+    let time = pic.get_time();
+    pic.advance_time(std::time::Duration::from_secs(420));
+    // time is not certified so `query_and_check_time` would fail here
+    assert_eq!(pic.get_time(), time + std::time::Duration::from_secs(420));
+    pic.tick();
     query_and_check_time(&pic, canister);
+    assert_eq!(pic.get_time(), time + std::time::Duration::from_secs(420));
+    pic.tick();
+    query_and_check_time(&pic, canister);
+    assert_eq!(
+        pic.get_time(),
+        time + std::time::Duration::from_secs(420) + std::time::Duration::from_nanos(1)
+    );
+}
+
+#[test]
+#[should_panic(expected = "SettingTimeIntoPast")]
+fn set_time_into_past() {
+    let pic = PocketIc::new();
+
+    let now = SystemTime::now();
+    pic.set_time(now + std::time::Duration::from_secs(1));
+
+    pic.set_time(now);
 }
 
 #[test]

rs/pocket_ic_server/CHANGELOG.md

@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - New endpoint `/instances/<instance_id>/read/ingress_status` to fetch the status of an update call submitted through an ingress message.
   If an optional caller is provided, the status of the update call is known, but the update call was submitted by a different caller, then an error is returned.
+- New endpoint `/instances/<instance_id>/update/set_certified_time` to set the current certified time on all subnets of the PocketIC instance.
 
 ### Fixed
 - Canisters created via `provisional_create_canister_with_cycles` with the management canister ID as the effective canister ID

rs/pocket_ic_server/src/pocket_ic.rs

@@ -1166,32 +1166,55 @@ pub struct SetTime {
     pub time: Time,
 }
 
-impl Operation for SetTime {
-    fn compute(&self, pic: &mut PocketIc) -> OpOut {
-        // Time is kept in sync across subnets, so one can take any subnet.
-        let current_time: SystemTime = pic.any_subnet().time();
-        let set_time: SystemTime = self.time.into();
-        match current_time.cmp(&set_time) {
-            std::cmp::Ordering::Greater => OpOut::Error(PocketIcError::SettingTimeIntoPast((
-                systemtime_to_unix_epoch_nanos(current_time),
-                systemtime_to_unix_epoch_nanos(set_time),
-            ))),
-            std::cmp::Ordering::Equal => OpOut::NoOutput,
-            std::cmp::Ordering::Less => {
-                // Sets the time on all subnets.
-                for subnet in pic.subnets.get_all() {
+fn set_time(pic: &PocketIc, time: Time, certified: bool) -> OpOut {
+    // Time is kept in sync across subnets, so one can take any subnet.
+    let current_time: SystemTime = pic.any_subnet().time();
+    let set_time: SystemTime = time.into();
+    match current_time.cmp(&set_time) {
+        std::cmp::Ordering::Greater => OpOut::Error(PocketIcError::SettingTimeIntoPast((
+            systemtime_to_unix_epoch_nanos(current_time),
+            systemtime_to_unix_epoch_nanos(set_time),
+        ))),
+        std::cmp::Ordering::Equal => OpOut::NoOutput,
+        std::cmp::Ordering::Less => {
+            // Sets the time on all subnets.
+            for subnet in pic.subnets.get_all() {
+                if certified {
+                    subnet.state_machine.set_certified_time(set_time);
+                } else {
                     subnet.state_machine.set_time(set_time);
                 }
-                OpOut::NoOutput
             }
+            OpOut::NoOutput
         }
     }
+}
+
+impl Operation for SetTime {
+    fn compute(&self, pic: &mut PocketIc) -> OpOut {
+        set_time(pic, self.time, false)
+    }
 
     fn id(&self) -> OpId {
         OpId(format!("set_time_{}", self.time))
     }
 }
 
+#[derive(Clone, Debug)]
+pub struct SetCertifiedTime {
+    pub time: Time,
+}
+
+impl Operation for SetCertifiedTime {
+    fn compute(&self, pic: &mut PocketIc) -> OpOut {
+        set_time(pic, self.time, true)
+    }
+
+    fn id(&self) -> OpId {
+        OpId(format!("set_certified_time_{}", self.time))
+    }
+}
+
 #[derive(Copy, Clone, Debug)]
 pub struct GetTopology;
 

rs/pocket_ic_server/src/state_api/routes.rs

@@ -10,7 +10,7 @@ use crate::pocket_ic::{
     AddCycles, AwaitIngressMessage, CallRequest, CallRequestVersion, CanisterReadStateRequest,
     DashboardRequest, GetCanisterHttp, GetControllers, GetCyclesBalance, GetStableMemory,
     GetSubnet, GetTime, GetTopology, IngressMessageStatus, MockCanisterHttp, PubKey, Query,
-    QueryRequest, SetStableMemory, SetTime, StatusRequest, SubmitIngressMessage,
+    QueryRequest, SetCertifiedTime, SetStableMemory, SetTime, StatusRequest, SubmitIngressMessage,
     SubnetReadStateRequest, Tick,
 };
 use crate::{async_trait, pocket_ic::PocketIc, BlobStore, InstanceId, OpId, Operation};
@@ -99,6 +99,7 @@ where
             post(handler_await_ingress_message),
         )
         .directory_route("/set_time", post(handler_set_time))
+        .directory_route("/set_certified_time", post(handler_set_certified_time))
         .directory_route("/add_cycles", post(handler_add_cycles))
         .directory_route("/set_stable_memory", post(handler_set_stable_memory))
         .directory_route("/tick", post(handler_tick))
@@ -1035,6 +1036,20 @@ pub async fn handler_set_time(
     (code, Json(response))
 }
 
+pub async fn handler_set_certified_time(
+    State(AppState { api_state, .. }): State<AppState>,
+    Path(instance_id): Path<InstanceId>,
+    headers: HeaderMap,
+    axum::extract::Json(time): axum::extract::Json<rest::RawTime>,
+) -> (StatusCode, Json<ApiResponse<()>>) {
+    let timeout = timeout_or_default(headers);
+    let op = SetCertifiedTime {
+        time: ic_types::Time::from_nanos_since_unix_epoch(time.nanos_since_epoch),
+    };
+    let (code, response) = run_operation(api_state, instance_id, timeout, op).await;
+    (code, Json(response))
+}
+
 pub async fn handler_add_cycles(
     State(AppState { api_state, .. }): State<AppState>,
     Path(instance_id): Path<InstanceId>,

rs/state_machine_tests/src/lib.rs

@@ -2513,6 +2513,26 @@ impl StateMachine {
             .unwrap_or_else(|_| error!(self.replica_logger, "Time went backwards."));
     }
 
+    /// Certifies the specified time by modifying the time in the replicated state
+    /// and certifying that new state.
+    pub fn set_certified_time(&self, time: SystemTime) {
+        let t = time
+            .duration_since(SystemTime::UNIX_EPOCH)
+            .unwrap()
+            .as_nanos() as u64;
+        let time = Time::from_nanos_since_unix_epoch(t);
+        let (height, mut replicated_state) = self.state_manager.take_tip();
+        replicated_state.metadata.batch_time = time;
+        self.state_manager.commit_and_certify(
+            replicated_state,
+            height.increment(),
+            CertificationScope::Metadata,
+            None,
+        );
+        self.set_time(time.into());
+        *self.time_of_last_round.write().unwrap() = time;
+    }
+
     /// Returns the current state machine time.
     /// The time of a round executed by this state machine equals its current time
     /// if its current time increased since the last round.