diff --git a/drivers/net/wireless/intel/iwlwifi/mvm/mac80211.c b/drivers/net/wireless/intel/iwlwifi/mvm/mac80211.c
index 09afb6d557146e3fad709a3d4130d5da7f4ded93..f428f1cb00346ea36904f92d6264547e740464be 100644
--- a/drivers/net/wireless/intel/iwlwifi/mvm/mac80211.c
+++ b/drivers/net/wireless/intel/iwlwifi/mvm/mac80211.c
@@ -967,6 +967,7 @@ static void iwl_mvm_restart_cleanup(struct iwl_mvm *mvm)
 	mvm->calibrating = false;
 
 	/* just in case one was running */
+	iwl_mvm_cleanup_roc_te(mvm);
 	ieee80211_remain_on_channel_expired(mvm->hw);
 
 	/*
diff --git a/drivers/net/wireless/intel/iwlwifi/mvm/time-event.c b/drivers/net/wireless/intel/iwlwifi/mvm/time-event.c
index 87a04c32cb926194d2a7c4ddc0d607766b86f217..924dd6a416266b52dbeb85cf0eccb7aa9e014960 100644
--- a/drivers/net/wireless/intel/iwlwifi/mvm/time-event.c
+++ b/drivers/net/wireless/intel/iwlwifi/mvm/time-event.c
@@ -792,11 +792,9 @@ int iwl_mvm_start_p2p_roc(struct iwl_mvm *mvm, struct ieee80211_vif *vif,
 	return iwl_mvm_time_event_send_add(mvm, vif, te_data, &time_cmd);
 }
 
-void iwl_mvm_stop_roc(struct iwl_mvm *mvm)
+static struct iwl_mvm_time_event_data *iwl_mvm_get_roc_te(struct iwl_mvm *mvm)
 {
-	struct iwl_mvm_vif *mvmvif = NULL;
 	struct iwl_mvm_time_event_data *te_data;
-	bool is_p2p = false;
 
 	lockdep_assert_held(&mvm->mutex);
 
@@ -810,11 +808,8 @@ void iwl_mvm_stop_roc(struct iwl_mvm *mvm)
 	 * request
 	 */
 	list_for_each_entry(te_data, &mvm->time_event_list, list) {
-		if (te_data->vif->type == NL80211_IFTYPE_P2P_DEVICE) {
-			mvmvif = iwl_mvm_vif_from_mac80211(te_data->vif);
-			is_p2p = true;
-			goto remove_te;
-		}
+		if (te_data->vif->type == NL80211_IFTYPE_P2P_DEVICE)
+			goto out;
 	}
 
 	/* There can only be at most one AUX ROC time event, we just use the
@@ -823,18 +818,35 @@ void iwl_mvm_stop_roc(struct iwl_mvm *mvm)
 	te_data = list_first_entry_or_null(&mvm->aux_roc_te_list,
 					   struct iwl_mvm_time_event_data,
 					   list);
+out:
+	spin_unlock_bh(&mvm->time_event_lock);
+	return te_data;
+}
+
+void iwl_mvm_cleanup_roc_te(struct iwl_mvm *mvm)
+{
+	struct iwl_mvm_time_event_data *te_data;
+	u32 uid;
+
+	te_data = iwl_mvm_get_roc_te(mvm);
 	if (te_data)
-		mvmvif = iwl_mvm_vif_from_mac80211(te_data->vif);
+		__iwl_mvm_remove_time_event(mvm, te_data, &uid);
+}
 
-remove_te:
-	spin_unlock_bh(&mvm->time_event_lock);
+void iwl_mvm_stop_roc(struct iwl_mvm *mvm)
+{
+	struct iwl_mvm_vif *mvmvif;
+	struct iwl_mvm_time_event_data *te_data;
 
-	if (!mvmvif) {
+	te_data = iwl_mvm_get_roc_te(mvm);
+	if (!te_data) {
 		IWL_WARN(mvm, "No remain on channel event\n");
 		return;
 	}
 
-	if (is_p2p)
+	mvmvif = iwl_mvm_vif_from_mac80211(te_data->vif);
+
+	if (te_data->vif->type == NL80211_IFTYPE_P2P_DEVICE)
 		iwl_mvm_remove_time_event(mvm, mvmvif, te_data);
 	else
 		iwl_mvm_remove_aux_roc_te(mvm, mvmvif, te_data);
diff --git a/drivers/net/wireless/intel/iwlwifi/mvm/time-event.h b/drivers/net/wireless/intel/iwlwifi/mvm/time-event.h
index 61d7cd791b6e2aa5a4ce4ac64ec7caff943b326b..99d9a35ad5b1b30cc3a7e09e3d6dd1d826756f44 100644
--- a/drivers/net/wireless/intel/iwlwifi/mvm/time-event.h
+++ b/drivers/net/wireless/intel/iwlwifi/mvm/time-event.h
@@ -215,6 +215,7 @@ void iwl_mvm_remove_time_event(struct iwl_mvm *mvm,
 void iwl_mvm_te_clear_data(struct iwl_mvm *mvm,
 			   struct iwl_mvm_time_event_data *te_data);
 
+void iwl_mvm_cleanup_roc_te(struct iwl_mvm *mvm);
 void iwl_mvm_roc_done_wk(struct work_struct *wk);
 
 /**