feat: [MR-552] Implement callback expiration

rs/protobuf/def/state/queues/v1/queues.proto

@@ -180,6 +180,7 @@ message CanisterQueues {
     uint64 id = 1;
     uint64 callback_id = 2;
   }
+  repeated CallbackReference expired_callbacks = 13;
   repeated CallbackReference shed_responses = 12;
 
   enum NextInputQueue {

rs/protobuf/src/gen/state/state.queues.v1.rs

@@ -216,6 +216,8 @@ pub struct CanisterQueues {
     pub canister_queues: ::prost::alloc::vec::Vec<canister_queues::CanisterQueuePair>,
     #[prost(message, optional, tag = "10")]
     pub pool: ::core::option::Option<MessagePool>,
+    #[prost(message, repeated, tag = "13")]
+    pub expired_callbacks: ::prost::alloc::vec::Vec<canister_queues::CallbackReference>,
     #[prost(message, repeated, tag = "12")]
     pub shed_responses: ::prost::alloc::vec::Vec<canister_queues::CallbackReference>,
     #[prost(enumeration = "canister_queues::NextInputQueue", tag = "6")]

rs/protobuf/src/gen/types/state.queues.v1.rs

@@ -216,6 +216,8 @@ pub struct CanisterQueues {
     pub canister_queues: ::prost::alloc::vec::Vec<canister_queues::CanisterQueuePair>,
     #[prost(message, optional, tag = "10")]
     pub pool: ::core::option::Option<MessagePool>,
+    #[prost(message, repeated, tag = "13")]
+    pub expired_callbacks: ::prost::alloc::vec::Vec<canister_queues::CallbackReference>,
     #[prost(message, repeated, tag = "12")]
     pub shed_responses: ::prost::alloc::vec::Vec<canister_queues::CallbackReference>,
     #[prost(enumeration = "canister_queues::NextInputQueue", tag = "6")]

rs/replicated_state/src/canister_state/queues/message_pool.rs

@@ -117,6 +117,18 @@ impl Id {
             Class::BestEffort
         }
     }
+
+    /// Tests whether this `Id` represents an inbound best-effort response.
+    fn is_inbound_best_effort_response(&self) -> bool {
+        self.0 & (Context::BIT | Class::BIT | Kind::BIT)
+            == (Context::Inbound as u64 | Class::BestEffort as u64 | Kind::Response as u64)
+    }
+
+    /// Tests whether this `Id` represents an outbound guaranteed-response request.
+    fn is_outbound_guaranteed_request(&self) -> bool {
+        self.0 & (Context::BIT | Class::BIT | Kind::BIT)
+            == (Context::Outbound as u64 | Class::GuaranteedResponse as u64 | Kind::Request as u64)
+    }
 }
 
 /// A typed reference -- inbound (`CanisterInput`) or outbound
@@ -142,13 +154,19 @@ impl<T> Reference<T> {
         Id::from(self).kind()
     }
 
-    pub(super) fn context(&self) -> Context {
+    fn context(&self) -> Context {
         Id::from(self).context()
     }
 
-    pub(super) fn class(&self) -> Class {
+    #[allow(dead_code)]
+    fn class(&self) -> Class {
         Id::from(self).class()
     }
+
+    /// Tests whether this is a reference to an inbound best-effort response.
+    pub(super) fn is_inbound_best_effort_response(&self) -> bool {
+        Id::from(self).is_inbound_best_effort_response()
+    }
 }
 
 impl<T> Clone for Reference<T> {
@@ -324,10 +342,7 @@ impl TryFrom<pb_queues::canister_queues::CallbackReference> for CallbackReferenc
     type Error = ProxyDecodeError;
     fn try_from(item: pb_queues::canister_queues::CallbackReference) -> Result<Self, Self::Error> {
         let reference = Reference(item.id, PhantomData);
-        if reference.context() == Context::Inbound
-            && reference.class() == Class::BestEffort
-            && reference.kind() == Kind::Response
-        {
+        if reference.is_inbound_best_effort_response() {
             Ok(CallbackReference(reference, item.callback_id.into()))
         } else {
             Err(ProxyDecodeError::Other(
@@ -411,6 +426,14 @@ impl MessagePool {
         self.insert_impl(msg, actual_deadline, Context::Inbound)
     }
 
+    /// Reserves an `InboundReference` for a timeout reject response for a
+    /// best-effort callback.
+    ///
+    /// This is equivalent to inserting and then immediately removing the response.
+    pub(super) fn make_inbound_timeout_response_reference(&mut self) -> InboundReference {
+        self.next_reference(Class::BestEffort, Kind::Response)
+    }
+
     /// Inserts an outbound request (one that is to be enqueued in an output queue)
     /// into the pool. Returns the reference assigned to the request.
     ///
@@ -640,7 +663,9 @@ impl MessagePool {
             .into_iter()
             .map(|(_, id)| {
                 let msg = self.take_impl(id).unwrap();
-                self.outbound_guaranteed_request_deadlines.remove(&id);
+                if id.is_outbound_guaranteed_request() {
+                    self.outbound_guaranteed_request_deadlines.remove(&id);
+                }
                 self.remove_from_size_queue(id, &msg);
                 (id.into(), msg)
             })
@@ -736,10 +761,7 @@ impl MessagePool {
         // guaranteed response requests (and nothing else).
         let mut expected_outbound_guaranteed_request_ids = BTreeSet::new();
         self.messages.keys().for_each(|id| {
-            if id.context() == Context::Outbound
-                && id.class() == Class::GuaranteedResponse
-                && id.kind() == Kind::Request
-            {
+            if id.is_outbound_guaranteed_request() {
                 expected_outbound_guaranteed_request_ids.insert(id);
             }
         });

rs/replicated_state/src/canister_state/queues/message_pool/tests.rs

@@ -135,6 +135,33 @@ fn test_get() {
 
 #[test]
 fn test_take() {
+    fn test_take_impl<T>(
+        request_id: Reference<T>,
+        response_id: Reference<T>,
+        request: Request,
+        response: Response,
+        pool: &mut MessagePool,
+    ) {
+        let request: RequestOrResponse = request.into();
+        let response: RequestOrResponse = response.into();
+
+        // Ensure that the messages are now in the pool.
+        assert_eq!(Some(&request), pool.get(request_id));
+        assert_eq!(Some(&response), pool.get(response_id));
+
+        // Actually take the messages.
+        assert_eq!(Some(request), pool.take(request_id));
+        assert_eq!(Some(response), pool.take(response_id));
+
+        // Messages are gone.
+        assert_eq!(None, pool.get(request_id));
+        assert_eq!(None, pool.get(response_id));
+
+        // And cannot be taken out again.
+        assert_eq!(None, pool.take(request_id));
+        assert_eq!(None, pool.take(response_id));
+    }
+
     let mut pool = MessagePool::default();
 
     for deadline in [NO_DEADLINE, time(13)] {
@@ -155,33 +182,6 @@ fn test_take() {
                     test_take_impl(request_id, response_id, request, response, &mut pool);
                 }
             }
-
-            fn test_take_impl<T>(
-                request_id: Reference<T>,
-                response_id: Reference<T>,
-                request: Request,
-                response: Response,
-                pool: &mut MessagePool,
-            ) {
-                let request: RequestOrResponse = request.into();
-                let response: RequestOrResponse = response.into();
-
-                // Ensure that the messages are now in the pool.
-                assert_eq!(Some(&request), pool.get(request_id));
-                assert_eq!(Some(&response), pool.get(response_id));
-
-                // Actually take the messages.
-                assert_eq!(Some(request), pool.take(request_id));
-                assert_eq!(Some(response), pool.take(response_id));
-
-                // Messages are gone.
-                assert_eq!(None, pool.get(request_id));
-                assert_eq!(None, pool.get(response_id));
-
-                // And cannot be taken out again.
-                assert_eq!(None, pool.take(request_id));
-                assert_eq!(None, pool.take(response_id));
-            }
         }
     }
 
@@ -562,13 +562,68 @@ fn test_id_from_reference_roundtrip() {
             assert_eq!(reference.0, id.0);
             assert_eq!(id, reference.into());
             assert_eq!(SomeReference::Inbound(reference), SomeReference::from(id));
+            assert_eq!(
+                (id.context(), id.class(), id.kind()),
+                (reference.context(), reference.class(), reference.kind())
+            );
 
             // Outbound.
             let reference = OutboundReference::new(class, kind, 13);
             let id = Id::from(reference);
             assert_eq!(reference.0, id.0);
             assert_eq!(id, reference.into());
             assert_eq!(SomeReference::Outbound(reference), SomeReference::from(id));
+            assert_eq!(
+                (id.context(), id.class(), id.kind()),
+                (reference.context(), reference.class(), reference.kind())
+            );
+        }
+    }
+}
+
+#[test]
+fn test_is_inbound_best_effort_response() {
+    use Class::*;
+    use Kind::*;
+
+    for kind in [Request, Response] {
+        for class in [GuaranteedResponse, BestEffort] {
+            let reference = InboundReference::new(class, kind, 13);
+            let id = Id::from(reference);
+            assert_eq!(
+                class == BestEffort && kind == Response,
+                id.is_inbound_best_effort_response()
+            );
+            assert_eq!(
+                class == BestEffort && kind == Response,
+                reference.is_inbound_best_effort_response()
+            );
+
+            let reference = OutboundReference::new(class, kind, 13);
+            let id = Id::from(reference);
+            assert!(!id.is_inbound_best_effort_response());
+            assert!(!reference.is_inbound_best_effort_response());
+        }
+    }
+}
+
+#[test]
+fn test_is_outbound_guaranteed_request() {
+    use Class::*;
+    use Kind::*;
+
+    for kind in [Request, Response] {
+        for class in [GuaranteedResponse, BestEffort] {
+            let reference = InboundReference::new(class, kind, 13);
+            let id = Id::from(reference);
+            assert!(!id.is_outbound_guaranteed_request());
+
+            let reference = OutboundReference::new(class, kind, 13);
+            let id = Id::from(reference);
+            assert_eq!(
+                class == GuaranteedResponse && kind == Request,
+                id.is_outbound_guaranteed_request()
+            );
         }
     }
 }