mutex_unlock(&priv->mutex);
 }
 
-static void ath9k_htc_sta_notify(struct ieee80211_hw *hw,
-                                struct ieee80211_vif *vif,
-                                enum sta_notify_cmd cmd,
-                                struct ieee80211_sta *sta)
+static int ath9k_htc_sta_add(struct ieee80211_hw *hw,
+                            struct ieee80211_vif *vif,
+                            struct ieee80211_sta *sta)
 {
        struct ath9k_htc_priv *priv = hw->priv;
        int ret;
 
        mutex_lock(&priv->mutex);
        ath9k_htc_ps_wakeup(priv);
+       ret = ath9k_htc_add_station(priv, vif, sta);
+       if (!ret)
+               ath9k_htc_init_rate(priv, sta);
+       ath9k_htc_ps_restore(priv);
+       mutex_unlock(&priv->mutex);
 
-       switch (cmd) {
-       case STA_NOTIFY_ADD:
-               ret = ath9k_htc_add_station(priv, vif, sta);
-               if (!ret)
-                       ath9k_htc_init_rate(priv, sta);
-               break;
-       case STA_NOTIFY_REMOVE:
-               ath9k_htc_remove_station(priv, vif, sta);
-               break;
-       default:
-               break;
-       }
+       return ret;
+}
+
+static int ath9k_htc_sta_remove(struct ieee80211_hw *hw,
+                               struct ieee80211_vif *vif,
+                               struct ieee80211_sta *sta)
+{
+       struct ath9k_htc_priv *priv = hw->priv;
+       int ret;
 
+       mutex_lock(&priv->mutex);
+       ath9k_htc_ps_wakeup(priv);
+       ret = ath9k_htc_remove_station(priv, vif, sta);
        ath9k_htc_ps_restore(priv);
        mutex_unlock(&priv->mutex);
+
+       return ret;
 }
 
 static int ath9k_htc_conf_tx(struct ieee80211_hw *hw, u16 queue,
        .remove_interface   = ath9k_htc_remove_interface,
        .config             = ath9k_htc_config,
        .configure_filter   = ath9k_htc_configure_filter,
-       .sta_notify         = ath9k_htc_sta_notify,
+       .sta_add            = ath9k_htc_sta_add,
+       .sta_remove         = ath9k_htc_sta_remove,
        .conf_tx            = ath9k_htc_conf_tx,
        .bss_info_changed   = ath9k_htc_bss_info_changed,
        .set_key            = ath9k_htc_set_key,