In some situations where xhci removal happens parallel to xhci_handshake,
we encounter a scenario where the xhci_handshake can't succeed, and it
polls until timeout.
If xhci_handshake runs until timeout it can on some platforms result in
a long wait which might lead to a watchdog timeout.
Add a helper that checks xhci status during the handshake, and exits if
set state is entered. Use this helper in places where xhci_handshake is
called unlocked and has a long timeout. For example  xhci command timeout
and xhci reset.
[commit message and code comment rewording -Mathias]
Signed-off-by: Udipto Goswami <quic_ugoswami@quicinc.com>
Signed-off-by: Mathias Nyman <mathias.nyman@linux.intel.com>
Link: https://lore.kernel.org/r/20231019102924.2797346-18-mathias.nyman@linux.intel.com
Signed-off-by: Greg Kroah-Hartman <gregkh@linuxfoundation.org>
 
         * In the future we should distinguish between -ENODEV and -ETIMEDOUT
         * and try to recover a -ETIMEDOUT with a host controller reset.
         */
-       ret = xhci_handshake(&xhci->op_regs->cmd_ring,
-                       CMD_RING_RUNNING, 0, 5 * 1000 * 1000);
+       ret = xhci_handshake_check_state(xhci, &xhci->op_regs->cmd_ring,
+                       CMD_RING_RUNNING, 0, 5 * 1000 * 1000,
+                       XHCI_STATE_REMOVING);
        if (ret < 0) {
                xhci_err(xhci, "Abort failed to stop command ring: %d\n", ret);
                xhci_halt(xhci);
 
        return ret;
 }
 
+/*
+ * xhci_handshake_check_state - same as xhci_handshake but takes an additional
+ * exit_state parameter, and bails out with an error immediately when xhc_state
+ * has exit_state flag set.
+ */
+int xhci_handshake_check_state(struct xhci_hcd *xhci, void __iomem *ptr,
+               u32 mask, u32 done, int usec, unsigned int exit_state)
+{
+       u32     result;
+       int     ret;
+
+       ret = readl_poll_timeout_atomic(ptr, result,
+                               (result & mask) == done ||
+                               result == U32_MAX ||
+                               xhci->xhc_state & exit_state,
+                               1, usec);
+
+       if (result == U32_MAX || xhci->xhc_state & exit_state)
+               return -ENODEV;
+
+       return ret;
+}
+
 /*
  * Disable interrupts and begin the xHCI halting process.
  */
        if (xhci->quirks & XHCI_INTEL_HOST)
                udelay(1000);
 
-       ret = xhci_handshake(&xhci->op_regs->command, CMD_RESET, 0, timeout_us);
+       ret = xhci_handshake_check_state(xhci, &xhci->op_regs->command,
+                               CMD_RESET, 0, timeout_us, XHCI_STATE_REMOVING);
        if (ret)
                return ret;
 
 
 /* xHCI host controller glue */
 typedef void (*xhci_get_quirks_t)(struct device *, struct xhci_hcd *);
 int xhci_handshake(void __iomem *ptr, u32 mask, u32 done, u64 timeout_us);
+int xhci_handshake_check_state(struct xhci_hcd *xhci, void __iomem *ptr,
+               u32 mask, u32 done, int usec, unsigned int exit_state);
 void xhci_quiesce(struct xhci_hcd *xhci);
 int xhci_halt(struct xhci_hcd *xhci);
 int xhci_start(struct xhci_hcd *xhci);