uwb_core: add timeout for waiting notification

If the UciManager doesn't send session status notification when
initializing session, UwbSession's initialize() method will be stuck
forever. This CL handles this kind of situation by adding timeout.

Bug: 228136039
Test: cargo test
Test: cargo clippy
Change-Id: Ic5841254b5a4f2f5026550f37923317c943ca3e6

3 files changed, 45 insertions, 3 deletions

diff --git a/src/rust/uwb_core/src/session/error.rs b/src/rust/uwb_core/src/session/error.rs
index eacd03a..39ffe37 100644
--- a/src/rust/uwb_core/src/session/error.rs
+++ b/src/rust/uwb_core/src/session/error.rs
@@ -30,6 +30,8 @@ pub enum Error {
     InvalidArguments,
     #[error("Wrong SessionState: {0}")]
     WrongState(SessionState),
+    #[error("Notification is not received in timeout")]
+    Timeout,
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
diff --git a/src/rust/uwb_core/src/session/session_manager.rs b/src/rust/uwb_core/src/session/session_manager.rs
index 60ab908..b29bd33 100644
--- a/src/rust/uwb_core/src/session/session_manager.rs
+++ b/src/rust/uwb_core/src/session/session_manager.rs
@@ -289,4 +289,29 @@ mod tests {
 
         assert!(mock_uci_manager.wait_expected_calls_done().await);
     }
+
+    #[tokio::test]
+    async fn test_init_session_timeout() {
+        let session_id = 0x123;
+        let session_type = SessionType::FiraRangingSession;
+        let params = generate_params();
+
+        let session_id_clone = session_id;
+        let session_type_clone = session_type;
+        let (mut session_manager, mut mock_uci_manager) =
+            setup_session_manager(move |uci_manager| {
+                uci_manager.expect_session_init(
+                    session_id_clone,
+                    session_type_clone,
+                    vec![], // Not sending SessionStatus notification.
+                    Ok(()),
+                );
+            })
+            .await;
+
+        let result = session_manager.init_session(session_id, session_type, params).await;
+        assert_eq!(result, Err(Error::Timeout));
+
+        assert!(mock_uci_manager.wait_expected_calls_done().await);
+    }
 }
diff --git a/src/rust/uwb_core/src/session/uwb_session.rs b/src/rust/uwb_core/src/session/uwb_session.rs
index 93f894b..37f3824 100644
--- a/src/rust/uwb_core/src/session/uwb_session.rs
+++ b/src/rust/uwb_core/src/session/uwb_session.rs
@@ -12,8 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::time::Duration;
+
 use log::{debug, error, warn};
 use tokio::sync::{mpsc, oneshot, watch};
+use tokio::time::timeout;
 
 use crate::session::error::{Error, Result};
 use crate::session::params::AppConfigParams;
@@ -147,10 +150,22 @@ impl<T: UciManager> UwbSessionActor<T> {
     }
 
     async fn wait_state(&mut self, expected_state: SessionState) -> Result<()> {
-        if self.state_receiver.changed().await.is_err() {
-            debug!("UwbSession is about to drop.");
-            return Err(Error::TokioFailure);
+        const WAIT_STATE_TIMEOUT_MS: u64 = 1000;
+        match timeout(Duration::from_millis(WAIT_STATE_TIMEOUT_MS), self.state_receiver.changed())
+            .await
+        {
+            Ok(result) => {
+                if result.is_err() {
+                    debug!("UwbSession is about to drop.");
+                    return Err(Error::TokioFailure);
+                }
+            }
+            Err(_) => {
+                error!("Timeout waiting for the session status notification");
+                return Err(Error::Timeout);
+            }
         }
+
         let state = *self.state_receiver.borrow();
         if state != expected_state {
             error!(