Snap for 9861064 from d02c83686e9ea3d55e704ea705d917dd302de0ab to mainline-uwb-releaseaml_uwb_331820070android13-mainline-uwb-release

Change-Id: I7011f49bc0702bee36b3ed3fcb88fa28b8ec0ccb

8 files changed, 292 insertions, 94 deletions

diff --git a/src/rust/uci_hal_android/uci_hal_android.rs b/src/rust/uci_hal_android/uci_hal_android.rs
index 1b32344..aac73b0 100644
--- a/src/rust/uci_hal_android/uci_hal_android.rs
+++ b/src/rust/uci_hal_android/uci_hal_android.rs
@@ -29,7 +29,7 @@ use android_hardware_uwb::binder::{
 };
 use async_trait::async_trait;
 use binder_tokio::{Tokio, TokioRuntime};
-use log::error;
+use log::{debug, error};
 use tokio::runtime::Handle;
 use tokio::sync::{mpsc, Mutex};
 use uwb_core::error::{Error as UwbCoreError, Result as UwbCoreResult};
@@ -71,6 +71,7 @@ fn send_device_state_error_notification(
 struct RawUciCallback {
     uci_sender: mpsc::UnboundedSender<UciHalPacket>,
     hal_open_result_sender: mpsc::Sender<Result<()>>,
+    hal_core_init_result_sender: mpsc::Sender<Result<()>>,
     hal_close_result_sender: mpsc::Sender<Result<()>>,
 }
 
@@ -78,9 +79,15 @@ impl RawUciCallback {
     pub fn new(
         uci_sender: mpsc::UnboundedSender<UciHalPacket>,
         hal_open_result_sender: mpsc::Sender<Result<()>>,
+        hal_core_init_result_sender: mpsc::Sender<Result<()>>,
         hal_close_result_sender: mpsc::Sender<Result<()>>,
     ) -> Self {
-        Self { uci_sender, hal_open_result_sender, hal_close_result_sender }
+        Self {
+            uci_sender,
+            hal_open_result_sender,
+            hal_core_init_result_sender,
+            hal_close_result_sender,
+        }
     }
 }
 
@@ -89,6 +96,7 @@ impl Interface for RawUciCallback {}
 #[async_trait]
 impl IUwbClientCallbackAsyncServer for RawUciCallback {
     async fn onHalEvent(&self, event: UwbEvent, _event_status: UwbStatus) -> BinderResult<()> {
+        debug!("onHalEvent: {:?}, {:?}", event, _event_status);
         match event {
             // UwbEvent::ERROR is processed differently by UciHalAndroid depending on its state.
             //
@@ -109,10 +117,16 @@ impl IUwbClientCallbackAsyncServer for RawUciCallback {
                 send_device_state_error_notification(&self.uci_sender)
                     .map_err(|e| BinderStatus::from(Error::from(e)))
             }
-            UwbEvent::POST_INIT_CPLT => self.hal_open_result_sender.try_send(Ok(())).map_err(|e| {
-                error!("Failed sending POST_INIT_CPLT: {:?}", e);
+            UwbEvent::OPEN_CPLT => self.hal_open_result_sender.try_send(Ok(())).map_err(|e| {
+                error!("Failed sending OPEN_CPLT: {:?}", e);
                 BinderStatus::new_exception(ExceptionCode::TRANSACTION_FAILED, None)
             }),
+            UwbEvent::POST_INIT_CPLT => {
+                self.hal_core_init_result_sender.try_send(Ok(())).map_err(|e| {
+                    error!("Failed sending POST_INIT_CPLT: {:?}", e);
+                    BinderStatus::new_exception(ExceptionCode::TRANSACTION_FAILED, None)
+                })
+            }
             UwbEvent::CLOSE_CPLT => self.hal_close_result_sender.try_send(Ok(())).map_err(|e| {
                 error!("Failed sending CLOSE_CPLT: {:?}", e);
                 BinderStatus::new_exception(ExceptionCode::TRANSACTION_FAILED, None)
@@ -201,20 +215,30 @@ impl UciHal for UciHalAndroid {
 
         // Connect callback to packet_sender.
         let (hal_open_result_sender, mut hal_open_result_receiver) = mpsc::channel::<Result<()>>(1);
+        let (hal_core_init_result_sender, mut hal_core_init_result_receiver) =
+            mpsc::channel::<Result<()>>(1);
         let (hal_close_result_sender, hal_close_result_receiver) = mpsc::channel::<Result<()>>(1);
         let m_cback = BnUwbClientCallback::new_async_binder(
             RawUciCallback::new(
                 packet_sender.clone(),
                 hal_open_result_sender,
+                hal_core_init_result_sender,
                 hal_close_result_sender,
             ),
             TokioRuntime(Handle::current()),
             BinderFeatures::default(),
         );
+        // Open and wait for OPEN_CPLT.
+        debug!("Trigger IUwbChip.open()");
         i_uwb_chip.open(&m_cback).await.map_err(|e| UwbCoreError::from(Error::from(e)))?;
+        if (hal_open_result_receiver.recv().await).is_none() {
+            error!("OPEN_CPLT event is not received");
+            return Err(UwbCoreError::Unknown);
+        }
         // Initialize core and wait for POST_INIT_CPLT.
+        debug!("Trigger IUwbChip.coreInit()");
         i_uwb_chip.coreInit().await.map_err(|e| UwbCoreError::from(Error::from(e)))?;
-        match hal_open_result_receiver.recv().await {
+        match hal_core_init_result_receiver.recv().await {
             Some(Ok(())) => {
                 // Workaround while http://b/243140882 is not fixed:
                 // Send DEVICE_STATE_READY notification as chip is not sending this notification.
diff --git a/src/rust/uwb_core/src/params/uci_packets.rs b/src/rust/uwb_core/src/params/uci_packets.rs
index e66856a..17f7d0c 100644
--- a/src/rust/uwb_core/src/params/uci_packets.rs
+++ b/src/rust/uwb_core/src/params/uci_packets.rs
@@ -30,7 +30,7 @@ pub use uwb_uci_packets::{
     ShortAddressOwrAoaRangingMeasurement, ShortAddressTwoWayRangingMeasurement, StatusCode,
     UpdateMulticastListAction,
 };
-pub(crate) use uwb_uci_packets::{UciControlPacket, UciDataPacket, UciDataPacketHal, UciDataSnd};
+pub(crate) use uwb_uci_packets::{UciControlPacket, UciDataPacket, UciDataPacketHal};
 
 use crate::error::Error;
 
diff --git a/src/rust/uwb_core/src/uci/command.rs b/src/rust/uwb_core/src/uci/command.rs
index eb3b3df..210447e 100644
--- a/src/rust/uwb_core/src/uci/command.rs
+++ b/src/rust/uwb_core/src/uci/command.rs
@@ -68,6 +68,9 @@ pub enum UciCommand {
         session_id: u32,
         ranging_round_indexes: Vec<u8>,
     },
+    SessionQueryMaxDataSize {
+        session_id: SessionId,
+    },
     SessionStart {
         session_id: SessionId,
     },
@@ -173,6 +176,9 @@ impl TryFrom<UciCommand> for uwb_uci_packets::UciControlPacket {
             UciCommand::SessionGetRangingCount { session_id } => {
                 uwb_uci_packets::SessionGetRangingCountCmdBuilder { session_id }.build().into()
             }
+            UciCommand::SessionQueryMaxDataSize { session_id } => {
+                uwb_uci_packets::SessionQueryMaxDataSizeCmdBuilder { session_id }.build().into()
+            }
         };
         Ok(packet)
     }
diff --git a/src/rust/uwb_core/src/uci/mock_uci_manager.rs b/src/rust/uwb_core/src/uci/mock_uci_manager.rs
index a7aedb0..4b620a3 100644
--- a/src/rust/uwb_core/src/uci/mock_uci_manager.rs
+++ b/src/rust/uwb_core/src/uci/mock_uci_manager.rs
@@ -279,6 +279,20 @@ impl MockUciManager {
         );
     }
 
+    /// Prepare Mock to expect for session_query_max_data_size.
+    ///
+    /// MockUciManager expects call, returns out as response.
+    pub fn expect_session_query_max_data_size(
+        &mut self,
+        expected_session_id: SessionId,
+        out: Result<u16>,
+    ) {
+        self.expected_calls
+            .lock()
+            .unwrap()
+            .push_back(ExpectedCall::SessionQueryMaxDataSize { expected_session_id, out });
+    }
+
     /// Prepare Mock to expect range_start.
     ///
     /// MockUciManager expects call with parameters, returns out as response, followed by notfs
@@ -738,6 +752,23 @@ impl UciManager for MockUciManager {
         }
     }
 
+     async fn session_query_max_data_size(&self, session_id: SessionId) -> Result<u16> {
+        let mut expected_calls = self.expected_calls.lock().unwrap();
+        match expected_calls.pop_front() {
+            Some(ExpectedCall::SessionQueryMaxDataSize {expected_session_id, out})
+                if expected_session_id == session_id =>
+            {
+                self.expect_call_consumed.notify_one();
+                out
+            }
+            Some(call) => {
+                expected_calls.push_front(call);
+                Err(Error::MockUndefined)
+            }
+            None => Err(Error::MockUndefined),
+        }
+    }
+
     async fn range_start(&self, session_id: SessionId) -> Result<()> {
         let mut expected_calls = self.expected_calls.lock().unwrap();
         match expected_calls.pop_front() {
@@ -958,6 +989,10 @@ enum ExpectedCall {
         expected_ranging_round_indexes: Vec<u8>,
         out: Result<SessionUpdateActiveRoundsDtTagResponse>,
     },
+    SessionQueryMaxDataSize {
+        expected_session_id: SessionId,
+        out: Result<u16>,
+    },
     RangeStart {
         expected_session_id: SessionId,
         notfs: Vec<UciNotification>,
diff --git a/src/rust/uwb_core/src/uci/response.rs b/src/rust/uwb_core/src/uci/response.rs
index 6c1afab..9cb19bf 100644
--- a/src/rust/uwb_core/src/uci/response.rs
+++ b/src/rust/uwb_core/src/uci/response.rs
@@ -43,6 +43,7 @@ pub(super) enum UciResponse {
     SessionGetState(Result<SessionState>),
     SessionUpdateControllerMulticastList(Result<()>),
     SessionUpdateActiveRoundsDtTag(Result<SessionUpdateActiveRoundsDtTagResponse>),
+    SessionQueryMaxDataSize(Result<u16>),
     SessionStart(Result<()>),
     SessionStop(Result<()>),
     SessionGetRangingCount(Result<usize>),
@@ -77,6 +78,8 @@ impl UciResponse {
 
             Self::CoreSetConfig(resp) => Self::matches_status_retry(&resp.status),
             Self::SessionSetAppConfig(resp) => Self::matches_status_retry(&resp.status),
+
+            Self::SessionQueryMaxDataSize(result) => Self::matches_result_retry(result),
         }
     }
 
@@ -191,6 +194,9 @@ impl TryFrom<uwb_uci_packets::SessionConfigResponse> for UciResponse {
                     }),
                 ))
             }
+            SessionConfigResponseChild::SessionQueryMaxDataSizeRsp(evt) =>
+                 Ok(UciResponse::SessionQueryMaxDataSize(Ok(evt.get_max_data_size()),
+                )),
             _ => Err(Error::Unknown),
         }
     }
diff --git a/src/rust/uwb_core/src/uci/uci_manager.rs b/src/rust/uwb_core/src/uci/uci_manager.rs
index c158d20..e8fe393 100644
--- a/src/rust/uwb_core/src/uci/uci_manager.rs
+++ b/src/rust/uwb_core/src/uci/uci_manager.rs
@@ -25,11 +25,10 @@ use crate::uci::command::UciCommand;
 use crate::error::{Error, Result};
 use crate::params::uci_packets::{
     AppConfigTlv, AppConfigTlvType, CapTlv, Controlees, CoreSetConfigResponse, CountryCode,
-    CreditAvailability, DataTransferNtfStatusCode, DeviceConfigId, DeviceConfigTlv, DeviceState,
-    FiraComponent, GetDeviceInfoResponse, GroupId, MessageType, PowerStats, RawUciMessage,
-    ResetConfig, SessionId, SessionState, SessionType, SessionUpdateActiveRoundsDtTagResponse,
-    SetAppConfigResponse, UciControlPacket, UciDataPacket, UciDataPacketHal, UciDataSnd,
-    UpdateMulticastListAction,
+    CreditAvailability, DeviceConfigId, DeviceConfigTlv, DeviceState, FiraComponent,
+    GetDeviceInfoResponse, GroupId, MessageType, PowerStats, RawUciMessage, ResetConfig, SessionId,
+    SessionState, SessionType, SessionUpdateActiveRoundsDtTagResponse, SetAppConfigResponse,
+    UciControlPacket, UciDataPacket, UciDataPacketHal, UpdateMulticastListAction,
 };
 use crate::params::utils::bytes_to_u64;
 use crate::uci::message::UciMessage;
@@ -116,6 +115,8 @@ pub trait UciManager: 'static + Send + Sync + Clone {
         ranging_round_indexes: Vec<u8>,
     ) -> Result<SessionUpdateActiveRoundsDtTagResponse>;
 
+    async fn session_query_max_data_size(&self, session_id: SessionId) -> Result<u16>;
+
     async fn range_start(&self, session_id: SessionId) -> Result<()>;
     async fn range_stop(&self, session_id: SessionId) -> Result<()>;
     async fn range_get_ranging_count(&self, session_id: SessionId) -> Result<usize>;
@@ -149,8 +150,17 @@ pub trait UciManager: 'static + Send + Sync + Clone {
 #[derive(Clone)]
 pub struct UciManagerImpl {
     cmd_sender: mpsc::UnboundedSender<(UciManagerCmd, oneshot::Sender<Result<UciResponse>>)>,
-    data_packet_sender:
-        mpsc::UnboundedSender<(UciDataSnd, oneshot::Sender<DataTransferNtfStatusCode>)>,
+    // Store a sender for sending a oneshot channel sender "data_transfer_status_ntf_sender". This
+    // is done once for each DataSnd packet, before sending it's fragments. After all the DataSnd
+    // packet fragments are sent, UciManagerImpl will wait to receive a UCI DataTransferStatusNtf
+    // packet, on the corresponding oneshot channel receiver "data_transfer_status_ntf_receiver".
+    data_transfer_status_oneshot_sender:
+        mpsc::UnboundedSender<oneshot::Sender<SessionNotification>>,
+    // Store a sender for sending a tuple of DataSnd packet fragment (which should be written to
+    // the HAL), a oneshot channel sender to send back the UCI DataCreditNtf for that DataSnd
+    // packet fragment.
+    data_packet_fragment_sender:
+        mpsc::UnboundedSender<(UciDataPacketHal, oneshot::Sender<SessionNotification>)>,
 }
 
 impl UciManagerImpl {
@@ -161,12 +171,21 @@ impl UciManagerImpl {
         logger_mode: UciLoggerMode,
     ) -> Self {
         let (cmd_sender, cmd_receiver) = mpsc::unbounded_channel();
-        let (data_packet_sender, data_packet_receiver) = mpsc::unbounded_channel();
-        let mut actor =
-            UciManagerActor::new(hal, logger, logger_mode, cmd_receiver, data_packet_receiver);
+        let (data_transfer_status_oneshot_sender, data_transfer_status_oneshot_receiver) =
+            mpsc::unbounded_channel();
+        let (data_packet_fragment_sender, data_packet_fragment_receiver) =
+            mpsc::unbounded_channel();
+        let mut actor = UciManagerActor::new(
+            hal,
+            logger,
+            logger_mode,
+            cmd_receiver,
+            data_transfer_status_oneshot_receiver,
+            data_packet_fragment_receiver,
+        );
         tokio::spawn(async move { actor.run().await });
 
-        Self { cmd_sender, data_packet_sender }
+        Self { cmd_sender, data_transfer_status_oneshot_sender, data_packet_fragment_sender }
     }
 
     // Send the |cmd| to the UciManagerActor.
@@ -392,6 +411,15 @@ impl UciManager for UciManagerImpl {
         }
     }
 
+    async fn session_query_max_data_size(&self, session_id: SessionId) -> Result<u16> {
+        let cmd = UciCommand::SessionQueryMaxDataSize { session_id };
+        match self.send_cmd(UciManagerCmd::SendUciCommand { cmd }).await {
+            Ok(UciResponse::SessionQueryMaxDataSize(resp)) => resp,
+            Ok(_) => Err(Error::Unknown),
+            Err(e) => Err(e),
+        }
+    }
+
     async fn range_start(&self, session_id: SessionId) -> Result<()> {
         let cmd = UciCommand::SessionStart { session_id };
         match self.send_cmd(UciManagerCmd::SendUciCommand { cmd }).await {
@@ -461,6 +489,10 @@ impl UciManager for UciManagerImpl {
         uci_sequence_number: u8,
         data: Vec<u8>,
     ) -> Result<()> {
+        debug!(
+            "send_data_packet(): will Tx a data packet, session_id {}, sequence_number {}",
+            session_id, uci_sequence_number
+        );
         let dest_mac_address =
             bytes_to_u64(dest_mac_address_bytes).ok_or(Error::BadParameters).unwrap();
         let data_packet = uwb_uci_packets::UciDataSndBuilder {
@@ -471,26 +503,83 @@ impl UciManager for UciManagerImpl {
             data,
         }
         .build();
+
+        // We don't expect any other Data Packet Tx to be in progress when we come here.
+        //
+        // Create a oneshot channel to receive the DataTransferStatusNtf packet, and send the
+        // "sender" end of the channel to the UciManagerActor.
         let (data_transfer_status_ntf_sender, data_transfer_status_ntf_receiver) =
             oneshot::channel();
-        match self.data_packet_sender.send((data_packet, data_transfer_status_ntf_sender)) {
-            Ok(()) => {
-                // Wait to receive a DATA_TRANSFER_STATUS_NTF from UWBS. This indicates that the
-                // last fragment of the data packet has been received by the UWBS.
-                //
-                // TODO(b/261886903): Convert the DataTransferNtfStatusCode to Result.
-                data_transfer_status_ntf_receiver.await.unwrap();
-                Ok(())
-            }
+        match self.data_transfer_status_oneshot_sender.send(data_transfer_status_ntf_sender) {
+            Ok(()) => {}
             Err(e) => {
                 error!(
-                    "Failed to send data packet - got error{} for uci_sequence_number: {},\
-                       session_id: {}",
-                    e, uci_sequence_number, session_id
+                    "Error {} in sending a data_transfer_status_oneshot_sender for UCI Data packet\
+                     session_id {}, sequence_number {}",
+                    e, session_id, uci_sequence_number
                 );
-                Err(Error::PacketTxError)
+                return Err(Error::PacketTxError);
+            }
+        }
+
+        // We expect data Credit should be available when we start here, for all UWB Sessions as:
+        // - it's available by default for a UWB Session when it becomes active, and,
+        // - Data packet send completed for earlier packets only after the host received both
+        //   DATA_TRANSFER_STATUS and DATA_CREDIT notifications for them. The latter would have
+        //   indicated credit availability for the UWB session.
+        //
+        // TODO(b/261886903): Use a Map<SessionId, CreditAvailability> to explicitly confirm
+        // credit availability here (before sending any data packet fragment). The map should also
+        // be updated (in handle_notification()), when UWBS unilaterally sends a DATA_CREDIT_NTF.
+        let data_packet_session_id = data_packet.get_session_id();
+        let fragmented_packets: Vec<UciDataPacketHal> = data_packet.into();
+        for packet in fragmented_packets.into_iter() {
+            let (data_credit_ntf_sender, data_credit_ntf_receiver) = oneshot::channel();
+
+            match self.data_packet_fragment_sender.send((packet, data_credit_ntf_sender)) {
+                Ok(()) => {
+                    // Wait to receive a DATA_CREDIT_NTF from UWBS. This indicates that the
+                    // data packet fragment was received and Host can send the next fragment.
+                    let result = data_credit_ntf_receiver.await;
+                    if result.is_err() {
+                        error!("DataCreditNtf oneshot sender is dropped.");
+                        return Err(Error::PacketTxError);
+                    }
+
+                    if let SessionNotification::DataCredit { session_id, credit_availability } =
+                        result.unwrap()
+                    {
+                        if session_id != data_packet_session_id {
+                            error!(
+                                "Received Data Credit NTF for sessionID {}, was expected for \
+                                   sessionID {}",
+                                session_id, data_packet_session_id
+                            );
+                            return Err(Error::PacketTxError);
+                        }
+                        if credit_availability != CreditAvailability::CreditAvailable {
+                            error!("Received Data Credit NTF with no availability.");
+                            return Err(Error::PacketTxError);
+                        }
+                    }
+                }
+                Err(e) => {
+                    error!(
+                        "Failed to send data packet fragment - got error{} for \
+                        uci_sequence_number: {}, session_id: {}",
+                        e, uci_sequence_number, session_id
+                    );
+                    return Err(Error::PacketTxError);
+                }
             }
         }
+
+        let result = data_transfer_status_ntf_receiver.await;
+        if result.is_err() {
+            error!("DataTransferStatusNtf oneshot sender is dropped.");
+            return Err(Error::PacketTxError);
+        }
+        Ok(())
     }
 }
 
@@ -513,10 +602,17 @@ struct UciManagerActor<T: UciHal, U: UciLogger> {
     // Receive the commands and the corresponding response senders from UciManager.
     cmd_receiver: mpsc::UnboundedReceiver<(UciManagerCmd, oneshot::Sender<Result<UciResponse>>)>,
 
-    // Receive Data packets (to be sent to UWBS) and the corresponding status sender from
-    // UciManager.
-    data_packet_receiver:
-        mpsc::UnboundedReceiver<(UciDataSnd, oneshot::Sender<DataTransferNtfStatusCode>)>,
+    // Receive the oneshot channel sender, for sending back the expected DataTransferStatusNtf
+    // packet, after sending all the Data packet fragments is completed. This receiver is used
+    // to read the "data_transfer_status_ntf_sender" from UciManager, which is required once
+    // for every Tx data packet.
+    data_transfer_status_oneshot_receiver:
+        mpsc::UnboundedReceiver<oneshot::Sender<SessionNotification>>,
+
+    // Receive Data packet fragment(s) (to be sent to UWBS) and the corresponding DataCreditNtf
+    // sender, from the UciManager.
+    data_packet_fragment_receiver:
+        mpsc::UnboundedReceiver<(UciDataPacketHal, oneshot::Sender<SessionNotification>)>,
 
     // Set to true when |hal| is opened successfully.
     is_hal_opened: bool,
@@ -530,9 +626,10 @@ struct UciManagerActor<T: UciHal, U: UciLogger> {
     // notification.
     open_hal_result_sender: Option<oneshot::Sender<Result<UciResponse>>>,
 
-    // Send out the notifications received from UWBS, in response to sending a Data packet to it.
+    // Store the senders for the Data packet send UWBS notifications. These send back the received
+    // UWBS DATA_CREDIT_NTF and DATA_TRANSFER_STATUS_NTF packets back to the UciManager.
     data_credit_ntf_sender: Option<oneshot::Sender<SessionNotification>>,
-    data_transfer_status_ntf_sender: Option<oneshot::Sender<DataTransferNtfStatusCode>>,
+    data_transfer_status_ntf_sender: Option<oneshot::Sender<SessionNotification>>,
 
     // The timeout of waiting for the notification of device ready notification.
     wait_device_status_timeout: PinSleep,
@@ -564,16 +661,20 @@ impl<T: UciHal, U: UciLogger> UciManagerActor<T, U> {
             UciManagerCmd,
             oneshot::Sender<Result<UciResponse>>,
         )>,
-        data_packet_receiver: mpsc::UnboundedReceiver<(
-            UciDataSnd,
-            oneshot::Sender<DataTransferNtfStatusCode>,
+        data_transfer_status_oneshot_receiver: mpsc::UnboundedReceiver<
+            oneshot::Sender<SessionNotification>,
+        >,
+        data_packet_fragment_receiver: mpsc::UnboundedReceiver<(
+            UciDataPacketHal,
+            oneshot::Sender<SessionNotification>,
         )>,
     ) -> Self {
         Self {
             hal: TimeoutUciHal::new(hal),
             logger: UciLoggerWrapper::new(logger, logger_mode),
             cmd_receiver,
-            data_packet_receiver,
+            data_transfer_status_oneshot_receiver,
+            data_packet_fragment_receiver,
             is_hal_opened: false,
             packet_receiver: mpsc::unbounded_channel().1,
             defrager: Default::default(),
@@ -608,17 +709,50 @@ impl<T: UciHal, U: UciLogger> UciManagerActor<T, U> {
                     }
                 }
 
-                // Handle a data packet - this is to be sent from the Host to UWBS.
-                data = self.data_packet_receiver.recv(),
+                // Listen for and store a oneshot channel sender, for sending back a
+                // DataTransferStatusNtf to the UciManager. This is required once for each Data
+                // Tx packet (and is the first step in being ready to send it).
+                //
+                // The additional check confirms there is currently no stored oneshot channel
+                // sender, which also acts as a flow control mechanism (between UciManager and
+                // UciManagerActor).
+                channel_sender = self.data_transfer_status_oneshot_receiver.recv(),
                     if !self.is_waiting_data_packet_send_status() => {
+                    match channel_sender {
+                        Some(data_transfer_status_ntf_sender) => {
+                            self.data_transfer_status_ntf_sender = Some(data_transfer_status_ntf_sender);
+                        }
+                        None => {
+                            error!("Unexpected error as no data transfer status sender received\
+                                   from UciManager");
+                        }
+                    }
+                }
+
+                // Handle a data packet - this is to be sent from the Host to UWBS. Listen for a
+                // data packet only after receiving the oneshot sender on which the
+                // DataTransferStatusNtf can be sent back.
+                data = self.data_packet_fragment_receiver.recv(),
+                    if self.is_waiting_data_packet_send_status() => {
                     match data {
-                        Some((data_packet_send, data_transfer_status_ntf_sender)) => {
-                            self.data_transfer_status_ntf_sender =
-                                Some(data_transfer_status_ntf_sender);
-                            self.handle_data_packet_send(data_packet_send).await;
+                        Some((data_packet_fragment, data_credit_ntf_sender)) => {
+                            self.data_credit_ntf_sender = Some(data_credit_ntf_sender);
+                            let result = self.hal.send_packet(data_packet_fragment.to_vec()).await;
+                            if result.is_err() {
+                                error!("Result {:?} in sending data packet to HAL", result);
+                                if self.data_credit_ntf_sender.is_some() {
+                                    drop(self.data_credit_ntf_sender.take());
+                                }
+                                if self.data_transfer_status_ntf_sender.is_some() {
+                                    drop(self.data_transfer_status_ntf_sender.take());
+                                }
+                            }
                         },
                         None => {
-                            debug!("Unexpected error as no data packet to send from UciManager");
+                            error!("Unexpected error as no data packet to send from UciManager");
+                            if self.data_transfer_status_ntf_sender.is_some() {
+                                drop(self.data_transfer_status_ntf_sender.take());
+                            }
                         }
                     }
                 }
@@ -747,49 +881,6 @@ impl<T: UciHal, U: UciLogger> UciManagerActor<T, U> {
         }
     }
 
-    async fn handle_data_packet_send(&mut self, data_packet: UciDataSnd) {
-        // We expect data Credit should be available when we start here, for all UWB Sessions as:
-        // - it's available by default for a UWB Session when it becomes active, and,
-        // - Data packet send completed for earlier packets only after the host received both
-        //   DATA_TRANSFER_STATUS and DATA_CREDIT notifications for them. The latter would have
-        //   indicated credit availability for the UWB session.
-        //
-        // TODO(b/261886903): Use a Map<SessionId, CreditAvailability> to explicitly confirm
-        // credit availability here (before sending any data packet fragment). The map should also
-        // be updated (in handle_notification()), when UWBS unilaterally sends a DATA_CREDIT_NTF.
-        let data_packet_session_id = data_packet.get_session_id();
-        let fragmented_packets: Vec<UciDataPacketHal> = data_packet.into();
-        for packet in fragmented_packets.into_iter() {
-            let (data_credit_ntf_sender, data_credit_ntf_receiver) = oneshot::channel();
-            self.data_credit_ntf_sender = Some(data_credit_ntf_sender);
-            let result = self.hal.send_packet(packet.to_vec()).await;
-            if result.is_err() {
-                error!("Error in sending data packet to HAL.");
-                return;
-            }
-
-            let result = data_credit_ntf_receiver.await;
-            if result.is_err() {
-                error!("oneshot sender is dropped.");
-                return;
-            }
-
-            if let SessionNotification::DataCredit { session_id, credit_availability } =
-                result.unwrap()
-            {
-                // TODO(b/261886903): More relevant error codes.
-                if session_id != data_packet_session_id {
-                    error!("Received Data Credit NTF for different sessionID.");
-                    return;
-                }
-                if credit_availability != CreditAvailability::CreditAvailable {
-                    error!("Received Data Credit NTF with no availability.");
-                    return;
-                }
-            }
-        }
-    }
-
     async fn retry_command(&mut self) {
         if let Some(mut retryer) = self.retryer.take() {
             if !retryer.could_retry() {
@@ -917,24 +1008,37 @@ impl<T: UciHal, U: UciLogger> UciManagerActor<T, U> {
                         warn!("notify_session_initialized() failed: {:?}", e);
                     }
                 }
-                if let SessionNotification::DataCredit { session_id: _, credit_availability: _ } =
+                if let SessionNotification::DataCredit { session_id, credit_availability: _ } =
                     session_notf
                 {
                     if let Some(data_credit_ntf_sender) = self.data_credit_ntf_sender.take() {
                         let _ = data_credit_ntf_sender.send(session_notf);
+                    } else {
+                        debug!(
+                            "No data_credit_ntf_sender present, when received DataCreditNtf for\
+                               session_id {}",
+                            session_id
+                        );
                     }
                     return; // We consume these here and don't need to send to upper layer.
                 }
                 if let SessionNotification::DataTransferStatus {
-                    session_id: _,
-                    uci_sequence_number: _,
+                    session_id,
+                    uci_sequence_number,
                     status,
                 } = session_notf
                 {
                     if let Some(data_transfer_status_ntf_sender) =
                         self.data_transfer_status_ntf_sender.take()
                     {
-                        let _ = data_transfer_status_ntf_sender.send(status);
+                        let _ = data_transfer_status_ntf_sender.send(session_notf);
+                    } else {
+                        debug!(
+                            "No data_transfer_status_ntf_sender present, when received a \
+                               DataTransferStatusNtf with session_id {}, UCI sequence number {},\
+                                status {:?}",
+                            session_id, uci_sequence_number, status
+                        );
                     }
                     return; // We consume these here and don't need to send to upper layer.
                 }
diff --git a/src/rust/uwb_core/src/uci/uci_manager_sync.rs b/src/rust/uwb_core/src/uci/uci_manager_sync.rs
index dd41d2c..971687e 100644
--- a/src/rust/uwb_core/src/uci/uci_manager_sync.rs
+++ b/src/rust/uwb_core/src/uci/uci_manager_sync.rs
@@ -306,6 +306,11 @@ impl<U: UciManager> UciManagerSync<U> {
         )
     }
 
+    /// Send UCI command for getting max data size for session.
+    pub fn session_query_max_data_size(&self, session_id: SessionId) -> Result<u16> {
+        self.runtime_handle.block_on(self.uci_manager.session_query_max_data_size(session_id))
+    }
+
     /// Send UCI command for starting ranging of the session.
     pub fn range_start(&self, session_id: SessionId) -> Result<()> {
         self.runtime_handle.block_on(self.uci_manager.range_start(session_id))
diff --git a/src/rust/uwb_uci_packets/uci_packets.pdl b/src/rust/uwb_uci_packets/uci_packets.pdl
index dcf7daa..af41c5d 100644
--- a/src/rust/uwb_uci_packets/uci_packets.pdl
+++ b/src/rust/uwb_uci_packets/uci_packets.pdl
@@ -253,6 +253,7 @@ enum AppConfigTlvType : 8 {
     CCC_UWB_CONFIG_ID = 0xA4,
     CCC_PULSESHAPE_COMBO = 0xA5,
     CCC_URSK_TTL = 0xA6,
+    CCC_LAST_INDEX_USED  = 0xA8,
 
     // Interleaving ratio if AOA_RESULT_REQ is set to 0xF0.
     NB_OF_RANGE_MEASUREMENTS = 0xE3,
@@ -928,6 +929,23 @@ test DataTransferStatusNtf {
     "\x62\x05\x00\x06\x00\x00\x00\x00\x00\x00\x01\x01\x00",
 }
 
+packet SessionQueryMaxDataSizeCmd : SessionConfigCommand (opcode = 0xB) { //QUERY_MAX_DATA_SIZE
+    session_id: 32,
+}
+
+test SessionQueryMaxDataSizeCmd {
+ "\x21\x0B\x00\x04\x00\x00\x00\x00",
+}
+
+packet SessionQueryMaxDataSizeRsp : SessionConfigResponse (opcode = 0xB) { //QUER_MAX_DATA_SIZE
+    session_id: 32,
+    max_data_size: 16,
+}
+
+test SessionQueryMaxDataSizeRsp {
+  "\x41\x0B\x00\x06\x00\x00\x00\x00\x0E7\0x07",
+}
+
 packet SessionStartCmd : SessionControlCommand (opcode = 0x0) { //RANGE_START
     session_id: 32,
 }