[PATCH v5 10/14] gpu: nova-core: gsp: cmdq: extract await_msg_with from await_msg

[Alexandre Courbot <acourbot@nvidia.com>]

The logic of `await_msg` is also usable for other kinds of messages that
are not RPC; it will be used to wait for GMC messages as well. Extract
that logic into a new `await_msg_with` method that uses a closure to
decide whether a given message is the one we wait for, and use it in
`await_msg`.

Signed-off-by: Alexandre Courbot <acourbot@nvidia.com>
---
 drivers/gpu/nova-core/gsp/cmdq.rs | 44 ++++++++++++++++++++++++++++-----------
 1 file changed, 32 insertions(+), 12 deletions(-)

diff --git a/drivers/gpu/nova-core/gsp/cmdq.rs b/drivers/gpu/nova-core/gsp/cmdq.rs
index 58761a5f609f..7db3d28e8902 100644
--- a/drivers/gpu/nova-core/gsp/cmdq.rs
+++ b/drivers/gpu/nova-core/gsp/cmdq.rs
@@ -971,6 +971,36 @@ fn consume_element<R>(
         Ok(res)
     }
 
+    /// Consumes messages from the queue until a message that satisfies `f` is encountered, and
+    /// returns the result of `f`.
+    ///
+    /// The message is waited for up to [`Cmdq::RECEIVE_TIMEOUT`] from the call.
+    ///
+    /// `f` is expected to return `Ok` with the result to return to the caller, and `Err(ENOMSG)` to
+    /// signal that the received message was not the one expected (such messages are dropped). Any
+    /// other error returned by `f` is propagated to the caller.
+    ///
+    /// # Errors
+    ///
+    /// - `ETIMEDOUT` if no message satisfying `f` arrives before the deadline, however many other
+    ///   messages arrive while waiting.
+    /// - `EIO` if a message fails framing or checksum validation.
+    /// - Any error returned by `f` with the exception of `ENOMSG`.
+    fn await_msg_with<R>(&mut self, f: impl Fn(GspParsedMessage<'_>) -> Result<R>) -> Result<R> {
+        let deadline = Instant::<Monotonic>::now() + Cmdq::RECEIVE_TIMEOUT;
+        loop {
+            let remaining = deadline - Instant::<Monotonic>::now();
+            if remaining.is_negative() {
+                break Err(ETIMEDOUT);
+            }
+            match self.consume_element(remaining, &f)? {
+                Ok(m) => break Ok(m),
+                Err(ENOMSG) => continue,
+                Err(e) => break Err(e),
+            }
+        }
+    }
+
     /// Receives a message of type `M`, waiting up to [`Cmdq::RECEIVE_TIMEOUT`] from the call.
     ///
     /// Any other message that arrives first is logged as an event and does not extend the
@@ -1013,18 +1043,8 @@ fn await_msg<M: MessageFromGsp>(&mut self) -> Result<M>
                 Err(ENOMSG)
             }
         };
-        let deadline = Instant::<Monotonic>::now() + Cmdq::RECEIVE_TIMEOUT;
-        loop {
-            let remaining = deadline - Instant::<Monotonic>::now();
-            if remaining.is_negative() {
-                break Err(ETIMEDOUT);
-            }
-            match self.consume_element(remaining, f)? {
-                Ok(m) => break Ok(m),
-                Err(ENOMSG) => continue,
-                Err(e) => break Err(e),
-            }
-        }
+
+        self.await_msg_with(f)
     }
 
     /// Logs and consumes every message the queue holds.

-- 
2.56.0