5 files changed, 48 insertions, 18 deletions

diff --git a/drivers/net/wireless/iwlwifi/mvm/mac80211.c b/drivers/net/wireless/iwlwifi/mvm/mac80211.c
index 899b56c85b5b..a5eb8c82f16a 100644
--- a/drivers/net/wireless/iwlwifi/mvm/mac80211.c
+++ b/drivers/net/wireless/iwlwifi/mvm/mac80211.c
@@ -946,7 +946,7 @@ static void iwl_mvm_mac_sta_notify(struct ieee80211_hw *hw,
 
         switch (cmd) {
         case STA_NOTIFY_SLEEP:
-                if (atomic_read(&mvmsta->pending_frames) > 0)
+                if (atomic_read(&mvm->pending_frames[mvmsta->sta_id]) > 0)
                         ieee80211_sta_block_awake(hw, sta, true);
                 /*
                  * The fw updates the STA to be asleep. Tx packets on the Tx
diff --git a/drivers/net/wireless/iwlwifi/mvm/mvm.h b/drivers/net/wireless/iwlwifi/mvm/mvm.h
index 8269bc562951..9f46b23801bc 100644
--- a/drivers/net/wireless/iwlwifi/mvm/mvm.h
+++ b/drivers/net/wireless/iwlwifi/mvm/mvm.h
@@ -292,6 +292,7 @@ struct iwl_mvm {
         struct ieee80211_sta __rcu *fw_id_to_mac_id[IWL_MVM_STATION_COUNT];
         struct work_struct sta_drained_wk;
         unsigned long sta_drained[BITS_TO_LONGS(IWL_MVM_STATION_COUNT)];
+        atomic_t pending_frames[IWL_MVM_STATION_COUNT];
 
         /* configured by mac80211 */
         u32 rts_threshold;
diff --git a/drivers/net/wireless/iwlwifi/mvm/sta.c b/drivers/net/wireless/iwlwifi/mvm/sta.c
index 0fd96e4da461..5c664ed54400 100644
--- a/drivers/net/wireless/iwlwifi/mvm/sta.c
+++ b/drivers/net/wireless/iwlwifi/mvm/sta.c
@@ -219,7 +219,7 @@ int iwl_mvm_add_sta(struct iwl_mvm *mvm,
         mvm_sta->max_agg_bufsize = LINK_QUAL_AGG_FRAME_LIMIT_DEF;
 
         /* HW restart, don't assume the memory has been zeroed */
-        atomic_set(&mvm_sta->pending_frames, 0);
+        atomic_set(&mvm->pending_frames[sta_id], 0);
         mvm_sta->tid_disable_agg = 0;
         mvm_sta->tfd_queue_msk = 0;
         for (i = 0; i < IEEE80211_NUM_ACS; i++)
@@ -407,14 +407,21 @@ int iwl_mvm_rm_sta(struct iwl_mvm *mvm,
         }
 
         /*
+         * Make sure that the tx response code sees the station as -EBUSY and
+         * calls the drain worker.
+         */
+        spin_lock_bh(&mvm_sta->lock);
+        /*
          * There are frames pending on the AC queues for this station.
          * We need to wait until all the frames are drained...
          */
-        if (atomic_read(&mvm_sta->pending_frames)) {
-                ret = iwl_mvm_drain_sta(mvm, mvm_sta, true);
+        if (atomic_read(&mvm->pending_frames[mvm_sta->sta_id])) {
                 rcu_assign_pointer(mvm->fw_id_to_mac_id[mvm_sta->sta_id],
                                    ERR_PTR(-EBUSY));
+                spin_unlock_bh(&mvm_sta->lock);
+                ret = iwl_mvm_drain_sta(mvm, mvm_sta, true);
         } else {
+                spin_unlock_bh(&mvm_sta->lock);
                 ret = iwl_mvm_rm_sta_common(mvm, mvm_sta->sta_id);
                 rcu_assign_pointer(mvm->fw_id_to_mac_id[mvm_sta->sta_id], NULL);
         }
diff --git a/drivers/net/wireless/iwlwifi/mvm/sta.h b/drivers/net/wireless/iwlwifi/mvm/sta.h
index 12abd2d71835..a4ddce77aaae 100644
--- a/drivers/net/wireless/iwlwifi/mvm/sta.h
+++ b/drivers/net/wireless/iwlwifi/mvm/sta.h
@@ -274,7 +274,6 @@ struct iwl_mvm_tid_data {
  * @bt_reduced_txpower: is reduced tx power enabled for this station
  * @lock: lock to protect the whole struct. Since %tid_data is access from Tx
  * and from Tx response flow, it needs a spinlock.
- * @pending_frames: number of frames for this STA on the shared Tx queues.
  * @tid_data: per tid data. Look at %iwl_mvm_tid_data.
  *
  * When mac80211 creates a station it reserves some space (hw->sta_data_size)
@@ -290,7 +289,6 @@ struct iwl_mvm_sta {
         u8 max_agg_bufsize;
         bool bt_reduced_txpower;
         spinlock_t lock;
-        atomic_t pending_frames;
         struct iwl_mvm_tid_data tid_data[IWL_MAX_TID_COUNT];
         struct iwl_lq_sta lq_sta;
         struct ieee80211_vif *vif;
diff --git a/drivers/net/wireless/iwlwifi/mvm/tx.c b/drivers/net/wireless/iwlwifi/mvm/tx.c
index 479074303bd7..f212f16502ff 100644
--- a/drivers/net/wireless/iwlwifi/mvm/tx.c
+++ b/drivers/net/wireless/iwlwifi/mvm/tx.c
@@ -416,9 +416,8 @@ int iwl_mvm_tx_skb(struct iwl_mvm *mvm, struct sk_buff *skb,
 
         spin_unlock(&mvmsta->lock);
 
-        if (mvmsta->vif->type == NL80211_IFTYPE_AP &&
-            txq_id < IWL_MVM_FIRST_AGG_QUEUE)
-                atomic_inc(&mvmsta->pending_frames);
+        if (txq_id < IWL_MVM_FIRST_AGG_QUEUE)
+                atomic_inc(&mvm->pending_frames[mvmsta->sta_id]);
 
         return 0;
 
@@ -680,16 +679,41 @@ static void iwl_mvm_rx_tx_cmd_single(struct iwl_mvm *mvm,
         /*
          * If the txq is not an AMPDU queue, there is no chance we freed
          * several skbs. Check that out...
-         * If there are no pending frames for this STA, notify mac80211 that
-         * this station can go to sleep in its STA table.
          */
-        if (txq_id < IWL_MVM_FIRST_AGG_QUEUE && mvmsta &&
-            !WARN_ON(skb_freed > 1) &&
-            mvmsta->vif->type == NL80211_IFTYPE_AP &&
-            atomic_sub_and_test(skb_freed, &mvmsta->pending_frames)) {
-                ieee80211_sta_block_awake(mvm->hw, sta, false);
-                set_bit(sta_id, mvm->sta_drained);
-                schedule_work(&mvm->sta_drained_wk);
+        if (txq_id < IWL_MVM_FIRST_AGG_QUEUE && !WARN_ON(skb_freed > 1) &&
+            atomic_sub_and_test(skb_freed, &mvm->pending_frames[sta_id])) {
+                if (mvmsta) {
+                        /*
+                         * If there are no pending frames for this STA, notify
+                         * mac80211 that this station can go to sleep in its
+                         * STA table.
+                         */
+                        if (mvmsta->vif->type == NL80211_IFTYPE_AP)
+                                ieee80211_sta_block_awake(mvm->hw, sta, false);
+                        /*
+                         * We might very well have taken mvmsta pointer while
+                         * the station was being removed. The remove flow might
+                         * have seen a pending_frame (because we didn't take
+                         * the lock) even if now the queues are drained. So make
+                         * really sure now that this the station is not being
+                         * removed. If it is, run the drain worker to remove it.
+                         */
+                        spin_lock_bh(&mvmsta->lock);
+                        sta = rcu_dereference(mvm->fw_id_to_mac_id[sta_id]);
+                        if (IS_ERR_OR_NULL(sta)) {
+                                /*
+                                 * Station disappeared in the meantime:
+                                 * so we are draining.
+                                 */
+                                set_bit(sta_id, mvm->sta_drained);
+                                schedule_work(&mvm->sta_drained_wk);
+                        }
+                        spin_unlock_bh(&mvmsta->lock);
+                } else if (!mvmsta) {
+                        /* Tx response without STA, so we are draining */
+                        set_bit(sta_id, mvm->sta_drained);
+                        schedule_work(&mvm->sta_drained_wk);
+                }
         }
 
         rcu_read_unlock();