#define __MLX5_EN_XSK_RX_H__
 
 #include "en.h"
+#include <net/xdp_sock.h>
 
 /* RX data path */
 
                                              struct mlx5e_wqe_frag_info *wi,
                                              u32 cqe_bcnt);
 
+static inline bool mlx5e_xsk_update_rx_wakeup(struct mlx5e_rq *rq, bool alloc_err)
+{
+       if (!xsk_umem_uses_need_wakeup(rq->umem))
+               return alloc_err;
+
+       if (unlikely(alloc_err))
+               xsk_set_rx_need_wakeup(rq->umem);
+       else
+               xsk_clear_rx_need_wakeup(rq->umem);
+
+       return false;
+}
+
 #endif /* __MLX5_EN_XSK_RX_H__ */
 
 #define __MLX5_EN_XSK_TX_H__
 
 #include "en.h"
+#include <net/xdp_sock.h>
 
 /* TX data path */
 
 
 bool mlx5e_xsk_tx(struct mlx5e_xdpsq *sq, unsigned int budget);
 
+static inline void mlx5e_xsk_update_tx_wakeup(struct mlx5e_xdpsq *sq)
+{
+       if (!xsk_umem_uses_need_wakeup(sq->umem))
+               return;
+
+       if (sq->pc != sq->cc)
+               xsk_clear_tx_need_wakeup(sq->umem);
+       else
+               xsk_set_tx_need_wakeup(sq->umem);
+}
+
 #endif /* __MLX5_EN_XSK_TX_H__ */
 
        rq->mpwqe.umr_in_progress += rq->mpwqe.umr_last_bulk;
        rq->mpwqe.actual_wq_head   = head;
 
-       /* If XSK Fill Ring doesn't have enough frames, busy poll by
-        * rescheduling the NAPI poll.
+       /* If XSK Fill Ring doesn't have enough frames, report the error, so
+        * that one of the actions can be performed:
+        * 1. If need_wakeup is used, signal that the application has to kick
+        * the driver when it refills the Fill Ring.
+        * 2. Otherwise, busy poll by rescheduling the NAPI poll.
         */
        if (unlikely(alloc_err == -ENOMEM && rq->umem))
                return true;
 
 #include <linux/irq.h>
 #include "en.h"
 #include "en/xdp.h"
+#include "en/xsk/rx.h"
 #include "en/xsk/tx.h"
 
 static inline bool mlx5e_channel_no_affinity_change(struct mlx5e_channel *c)
 
 static bool mlx5e_napi_xsk_post(struct mlx5e_xdpsq *xsksq, struct mlx5e_rq *xskrq)
 {
-       bool busy_xsk = false;
-
+       bool busy_xsk = false, xsk_rx_alloc_err;
+
+       /* Handle the race between the application querying need_wakeup and the
+        * driver setting it:
+        * 1. Update need_wakeup both before and after the TX. If it goes to
+        * "yes", it can only happen with the first update.
+        * 2. If the application queried need_wakeup before we set it, the
+        * packets will be transmitted anyway, even w/o a wakeup.
+        * 3. Give a chance to clear need_wakeup after new packets were queued
+        * for TX.
+        */
+       mlx5e_xsk_update_tx_wakeup(xsksq);
        busy_xsk |= mlx5e_xsk_tx(xsksq, MLX5E_TX_XSK_POLL_BUDGET);
-       busy_xsk |= xskrq->post_wqes(xskrq);
+       mlx5e_xsk_update_tx_wakeup(xsksq);
+
+       xsk_rx_alloc_err = xskrq->post_wqes(xskrq);
+       busy_xsk |= mlx5e_xsk_update_rx_wakeup(xskrq, xsk_rx_alloc_err);
 
        return busy_xsk;
 }