qcacmn: MLO Peer changes

This change implements timer based msgq processing to defer
the context for MLO peer. It also fixes some of the issues in
MLO peer path, It adds utils APIs for MLO params

Change-Id: I94384c91adfb8785d833bff1ba5541dc98cfc383
CRs-Fixed: 3059837

umac/cmn_services/obj_mgr/src/wlan_objmgr_global_obj.c

@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2016-2021 The Linux Foundation. All rights reserved.
+ * Copyright (c) 2021 Qualcomm Innovation Center, Inc. All rights reserved.
  *
  * Permission to use, copy, modify, and/or distribute this software for
  * any purpose with or without fee is hereby granted, provided that the
@@ -871,6 +872,8 @@ struct mlo_mgr_context *wlan_objmgr_get_mlo_ctx(void)
 	return g_umac_glb_obj->mlo_ctx;
 }
 
+qdf_export_symbol(wlan_objmgr_get_mlo_ctx);
+
 void wlan_objmgr_set_mlo_ctx(struct mlo_mgr_context *ctx)
 {
 	g_umac_glb_obj->mlo_ctx = ctx;

umac/mlme/include/wlan_vdev_mlme.h

@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2018-2021 The Linux Foundation. All rights reserved.
+ * Copyright (c) 2021 Qualcomm Innovation Center, Inc. All rights reserved.
  *
  * Permission to use, copy, modify, and/or distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -353,11 +354,13 @@ struct vdev_mlme_mgmt_generic {
 /*
  * struct wlan_vdev_aid_mgr – AID manager
  * @aid_bitmap: AID bitmap array
+ * @start_aid: start of AID index
  * @max_aid: Max allowed AID
  * @ref_cnt:  to share AID across VDEVs for MBSSID
  */
 struct wlan_vdev_aid_mgr {
 	qdf_bitmap(aid_bitmap, WLAN_UMAC_MAX_AID);
+	uint16_t start_aid;
 	uint16_t max_aid;
 	qdf_atomic_t ref_cnt;
 };

umac/mlo_mgr/inc/wlan_mlo_mgr_ap.h

@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2021, The Linux Foundation. All rights reserved.
+ * Copyright (c) 2021 Qualcomm Innovation Center, Inc. All rights reserved.
  *
  * Permission to use, copy, modify, and/or distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -23,6 +24,8 @@
 #include <wlan_mlo_mgr_cmn.h>
 #include <wlan_mlo_mgr_public_structs.h>
 
+#define WLAN_RESV_AID_BITS 0xc000
+#define WLAN_AID(b)    ((b) & ~0xc000)
 /**
  * mlo_ap_vdev_attach() - update vdev obj and vdev count to
  *                         wlan_mlo_dev_context
@@ -82,6 +85,30 @@ void mlo_ap_vdev_detach(struct wlan_objmgr_vdev *vdev);
  * Return: None
  */
 void mlo_ap_link_down_cmpl_notify(struct wlan_objmgr_vdev *vdev);
+
+/**
+ * wlan_vdev_mlme_aid_mgr_max_aid_set() - set VDEV Max AID
+ * @vdev: vdev pointer
+ * @max_aid: max AID
+ *
+ * This function sets max AID for the VDEV
+ *
+ * Return: void
+ */
+void wlan_vdev_mlme_aid_mgr_max_aid_set(struct wlan_objmgr_vdev *vdev,
+					uint16_t max_aid);
+
+/**
+ * wlan_vdev_mlme_set_start_aid() - set VDEV start AID
+ * @vdev: vdev pointer
+ * @start_aid: start AID
+ *
+ * This function sets start AID for the VDEV
+ *
+ * Return: void
+ */
+QDF_STATUS wlan_vdev_mlme_set_start_aid(struct wlan_objmgr_vdev *vdev,
+					uint16_t start_aid);
 /**
  * wlan_vdev_aid_mgr_init() - VDEV AID mgr init
  * @max_aid: max AID
@@ -191,6 +218,18 @@ QDF_STATUS mlo_free_aid(struct wlan_objmgr_vdev *vdev, uint16_t assoc_id);
  */
 uint16_t mlme_get_aid(struct wlan_objmgr_vdev *vdev);
 
+/**
+ * mlme_is_aid_set() - Check whether the AID is already allocated
+ * @vdev: VDEV
+ * @assoc_id: Assoc ID
+ *
+ * This function checks whether the AID is already allocated
+ *
+ * Return: 1 for AID is already allocated
+ *         0 for AID is available
+ */
+int mlme_is_aid_set(struct wlan_objmgr_vdev *vdev, uint16_t assoc_id);
+
 /**
  * wlan_mlo_peer_free_aid() - Free assoc id
  * @ml_aid_mgr: MLO AID mgr
@@ -243,6 +282,42 @@ QDF_STATUS mlo_peer_free_aid(struct wlan_mlo_dev_context *ml_dev,
  */
 void mlme_free_aid(struct wlan_objmgr_vdev *vdev, uint16_t assoc_id);
 
+/**
+ * mlo_set_aid() - public API to reserve AID
+ * @vdev: VDEV object
+ * @assoc_id: Assoc id to be reserved
+ *
+ * This function reserves AID of MLO VDEV
+ *
+ * Return: SUCCESS, if it is reserved
+ *         FAILURE, if it is already allocated
+ */
+QDF_STATUS mlo_set_aid(struct wlan_objmgr_vdev *vdev,
+		       uint16_t assoc_id);
+
+/**
+ * mlme_set_aid() - public API to reserve AID
+ * @vdev: VDEV object
+ * @assoc_id: Assoc id to be reserved
+ *
+ * This function reserves AID of VDEV
+ *
+ * Return: SUCCESS, if it is reserved
+ *         FAILURE, if it is already allocated
+ */
+QDF_STATUS mlme_set_aid(struct wlan_objmgr_vdev *vdev,
+			uint16_t assoc_id);
+
+/**
+ * wlan_mlme_get_aid_count() - public API to get AID count
+ * @vdev: VDEV object
+ *
+ * This function counts number AIDs allocated for the VDEV
+ *
+ * Return: aid count value
+ */
+uint16_t wlan_mlme_get_aid_count(struct wlan_objmgr_vdev *vdev);
+
 /**
  * mlo_ap_ml_peerid_alloc() - public API to allocate MLO peer id
  *

umac/mlo_mgr/inc/wlan_mlo_mgr_cmn.h

@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2021, The Linux Foundation. All rights reserved.
+ * Copyright (c) 2021 Qualcomm Innovation Center, Inc. All rights reserved.
  *
  * Permission to use, copy, modify, and/or distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -218,6 +219,14 @@ void mlo_mlme_peer_assoc_resp(struct wlan_objmgr_peer *peer);
 qdf_nbuf_t mlo_mlme_get_link_assoc_req(struct wlan_objmgr_peer *peer,
 				       uint8_t link_ix);
 
+/**
+ * mlo_mlme_peer_deauth() - Initiate deauth on link peer
+ * @peer: Object manager peer
+ *
+ * Return: void
+ */
+void mlo_mlme_peer_deauth(struct wlan_objmgr_peer *peer);
+
 /**
  * mlo_get_link_vdev_ix() - Get index of link VDEV in MLD
  * @ml_dev: ML device context

umac/mlo_mgr/inc/wlan_mlo_mgr_msgq.h

@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2021, The Linux Foundation. All rights reserved.
+ * Copyright (c) 2021 Qualcomm Innovation Center, Inc. All rights reserved.
  *
  * Permission to use, copy, modify, and/or distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -34,6 +35,7 @@ struct ctxt_switch_mgr {
 	qdf_list_t msgq_list;
 	qdf_spinlock_t ctxt_lock;
 	bool timer_started;
+	bool allow_msg;
 	uint16_t max_messages_procd;
 };
 
@@ -43,12 +45,14 @@ struct ctxt_switch_mgr {
  * @MLO_PEER_ASSOC:    Partner peer ASSOC
  * @MLO_PEER_ASSOC_FAIL:  Partner peer ASSOC failure
  * @MLO_PEER_DISCONNECT:  Partner peer Disconnect
+ * @MLO_PEER_DEAUTH:  Initiate Deauth for ML connection
  */
 enum mlo_msg_type {
 	MLO_PEER_CREATE,
 	MLO_PEER_ASSOC,
 	MLO_PEER_ASSOC_FAIL,
 	MLO_PEER_DISCONNECT,
+	MLO_PEER_DEAUTH,
 };
 
 /*
@@ -89,18 +93,28 @@ struct peer_discon_notify_s {
 	struct wlan_objmgr_peer *peer;
 };
 
+/*
+ * struct peer_deauth_notify_s - MLO partner peer deauth notification
+ * @peer: Link peer on which Peer deauth to be sent
+ */
+struct peer_deauth_notify_s {
+	struct wlan_objmgr_peer *peer;
+};
+
 /*
  * union msg_payload - MLO message payload
  * @peer_create: peer create notification structure
  * @peer_assoc: peer assoc notification structure
  * @peer_assoc_fail: peer assoc fail notification structure
  * @peer_disconn: peer disconnect notification structure
+ * @peer_deauth: peer deauth notification structure
  */
 union msg_payload {
 	struct peer_create_notif_s peer_create;
 	struct peer_assoc_notify_s peer_assoc;
 	struct peer_assoc_fail_notify_s peer_assoc_fail;
 	struct peer_discon_notify_s peer_disconn;
+	struct peer_deauth_notify_s peer_deauth;
 };
 
 #define MLO_MAX_MSGQ_SIZE 256
@@ -131,4 +145,22 @@ struct mlo_ctxt_switch_msg_s {
 QDF_STATUS mlo_msgq_post(enum mlo_msg_type type,
 			 struct wlan_mlo_dev_context *ml_dev,
 			 void *payload);
+
+/**
+ * mlo_msgq_init() - Init MLO message queue
+ *
+ * This function initializes MLO msg queue module
+ *
+ * Return: void
+ */
+void mlo_msgq_init(void);
+
+/**
+ * mlo_msgq_free() - Free MLO message queue
+ *
+ * This function frees MLO msg queue module
+ *
+ * Return: void
+ */
+void mlo_msgq_free(void);
 #endif

umac/mlo_mgr/inc/wlan_mlo_mgr_peer.h

@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2021, The Linux Foundation. All rights reserved.
+ * Copyright (c) 2021 Qualcomm Innovation Center, Inc. All rights reserved.
  *
  * Permission to use, copy, modify, and/or distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -144,6 +145,31 @@ QDF_STATUS wlan_mlo_peer_is_assoc_done(struct wlan_mlo_peer_context *ml_peer);
 struct wlan_objmgr_peer *wlan_mlo_peer_get_assoc_peer(
 					struct wlan_mlo_peer_context *ml_peer);
 
+/**
+ * mlo_peer_is_assoc_peer() - check whether the peer is assoc peer
+ * @ml_peer: MLO peer
+ * @peer: Link peer
+ *
+ * This function checks whether the peer is assoc peer of MLO peer,
+ * This API doesn't have lock protection, caller needs to take the lock
+ *
+ * Return: true, if it is assoc peer
+ */
+bool mlo_peer_is_assoc_peer(struct wlan_mlo_peer_context *ml_peer,
+			    struct wlan_objmgr_peer *peer);
+
+/**
+ * wlan_mlo_peer_is_assoc_peer() - check whether the peer is assoc peer
+ * @ml_peer: MLO peer
+ * @peer: Link peer
+ *
+ * This function checks whether the peer is assoc peer of MLO peer
+ *
+ * Return: true, if it is assoc peer
+ */
+bool wlan_mlo_peer_is_assoc_peer(struct wlan_mlo_peer_context *ml_peer,
+				 struct wlan_objmgr_peer *peer);
+
 /**
  * wlan_mlo_partner_peer_assoc_post() - Notify partner peer assoc
  * @peer: Link peer
@@ -154,6 +180,16 @@ struct wlan_objmgr_peer *wlan_mlo_peer_get_assoc_peer(
  */
 void wlan_mlo_partner_peer_assoc_post(struct wlan_objmgr_peer *assoc_peer);
 
+/**
+ * wlan_mlo_peer_deauth_init() - Initiate Deauth of MLO peer
+ * @ml_peer: MLO peer
+ *
+ * This function initiates deauth on MLO peer and its links peers
+ *
+ * Return: void
+ */
+void wlan_mlo_peer_deauth_init(struct wlan_mlo_peer_context *ml_peer);
+
 /**
  * wlan_mlo_partner_peer_create_failed_notify() - Notify peer creation fail
  * @ml_peer: MLO peer
@@ -195,7 +231,7 @@ QDF_STATUS wlan_mlo_peer_create(struct wlan_objmgr_vdev *vdev,
 				uint16_t aid);
 
 /**
- * mlo_peer_free() - Free MLO peer
+ * mlo_peer_cleanup() - Free MLO peer
  * @ml_peer: MLO peer
  *
  * This function frees MLO peer and resets MLO peer associations
@@ -204,7 +240,7 @@ QDF_STATUS wlan_mlo_peer_create(struct wlan_objmgr_vdev *vdev,
  *
  * Return: void
  */
-void mlo_peer_free(struct wlan_mlo_peer_context *ml_peer);
+void mlo_peer_cleanup(struct wlan_mlo_peer_context *ml_peer);
 
 /**
  * wlan_mlo_peer_get_ref() - Get ref of MLO peer
@@ -231,20 +267,22 @@ static inline void wlan_mlo_peer_release_ref(
 					struct wlan_mlo_peer_context *ml_peer)
 {
 	if (qdf_atomic_dec_and_test(&ml_peer->ref_cnt))
-		mlo_peer_free(ml_peer);
+		mlo_peer_cleanup(ml_peer);
 }
 
 /**
  * wlan_mlo_link_peer_attach() - MLO link peer attach
  * @ml_peer: MLO peer
  * @peer: Link peer
+ * @frm_buf: Assoc resp buffer of non-assoc link
  *
  * This function attaches link peer to MLO peer
  *
  * Return: SUCCESS, if peer is successfully attached to MLO peer
  */
 QDF_STATUS wlan_mlo_link_peer_attach(struct wlan_mlo_peer_context *ml_peer,
-				     struct wlan_objmgr_peer *peer);
+				     struct wlan_objmgr_peer *peer,
+				     qdf_nbuf_t frm_buf);
 
 /**
  * wlan_mlo_link_peer_delete() - MLO link peer delete
@@ -268,6 +306,42 @@ qdf_nbuf_t mlo_peer_get_link_peer_assoc_req_buf(
 			struct wlan_mlo_peer_context *ml_peer,
 			uint8_t link_ix);
 
+/**
+ * mlo_peer_get_link_peer_assoc_resp_buf() - get MLO link peer assoc resp buf
+ * @ml_peer: MLO peer
+ * @link_ix: Link index of the link peer
+ *
+ * This function retrieves stored assoc resp buffer of link peer
+ *
+ * Return: resp_buf, if link_peer is available
+ *         NULL, if link_peer is not present
+ */
+qdf_nbuf_t mlo_peer_get_link_peer_assoc_resp_buf(
+		struct wlan_mlo_peer_context *ml_peer,
+		uint8_t link_ix);
+
+/**
+ * wlan_mlo_peer_free_all_link_assoc_resp_buf() - Free all assoc resp buffers
+ * @peer: Link peer
+ *
+ * This function frees all assoc resp link buffers
+ *
+ * Return: void
+ */
+void wlan_mlo_peer_free_all_link_assoc_resp_buf(struct wlan_objmgr_peer *peer);
+
+/**
+ * wlan_mlo_peer_get_links_info() - get MLO peer partner links info
+ * @peer: Link peer
+ * @ml_links: structure to be filled with partner link info
+ *
+ * This function retrieves partner link info of link peer
+ *
+ * Return: void
+ */
+void wlan_mlo_peer_get_links_info(struct wlan_objmgr_peer *peer,
+				  struct mlo_tgt_partner_info *ml_links);
+
 /**
  ** APIs to operations on ML peer object
  */

umac/mlo_mgr/inc/wlan_mlo_mgr_public_structs.h

@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2021, The Linux Foundation. All rights reserved.
+ * Copyright (c) 2021 Qualcomm Innovation Center, Inc. All rights reserved.
  *
  * Permission to use, copy, modify, and/or distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -102,11 +103,13 @@ struct mlo_mgr_context {
 /*
  * struct wlan_ml_vdev_aid_mgr – ML AID manager
  * @aid_bitmap: AID bitmap array
+ * @start_aid: start of AID index
  * @max_aid: Max allowed AID
  * @aid_mgr[]:  Array of link vdev aid mgr
  */
 struct wlan_ml_vdev_aid_mgr {
 	qdf_bitmap(aid_bitmap, WLAN_UMAC_MAX_AID);
+	uint16_t start_aid;
 	uint16_t max_aid;
 	struct wlan_vdev_aid_mgr *aid_mgr[WLAN_UMAC_MLO_MAX_VDEVS];
 };
@@ -207,6 +210,7 @@ struct wlan_mlo_dev_context {
  * @link_ix: Link index
  * @is_primary: sets true if the peer is primary UMAC’s peer
  * @hw_link_id: HW Link id of peer
+ * @assoc_rsp_buf: Assoc resp buffer
  */
 struct wlan_mlo_link_peer_entry {
 	struct wlan_objmgr_peer *link_peer;
@@ -214,6 +218,7 @@ struct wlan_mlo_link_peer_entry {
 	uint8_t link_ix;
 	bool is_primary;
 	uint8_t hw_link_id;
+	qdf_nbuf_t assoc_rsp_buf;
 };
 
 /*
@@ -243,6 +248,7 @@ enum mlo_peer_state {
  * @ref_cnt: Reference counter to avoid use after free
  * @ml_dev: MLO dev context
  * @mlpeer_state: MLO peer state
+ * @avg_link_rssi: avg RSSI of ML peer
  */
 struct wlan_mlo_peer_context {
 	qdf_list_node_t peer_node;
@@ -262,6 +268,7 @@ struct wlan_mlo_peer_context {
 	qdf_atomic_t ref_cnt;
 	struct wlan_mlo_dev_context *ml_dev;
 	enum mlo_peer_state mlpeer_state;
+	int8_t avg_link_rssi;
 };
 
 /*
@@ -286,6 +293,26 @@ struct mlo_partner_info {
 	struct mlo_link_info partner_link_info[WLAN_UMAC_MLO_MAX_VDEVS];
 };
 
+/*
+ * struct mlo_tgt_link_info – ML target link info
+ * @vdev_id: link peer vdev id
+ * @hw_mld_link_id: HW link id
+ */
+struct mlo_tgt_link_info {
+	uint8_t vdev_id;
+	uint8_t hw_mld_link_id;
+};
+
+/*
+ * struct mlo_tgt_partner_info – mlo target partner link info
+ * @num_partner_links: no. of partner links
+ * @link_info: per partner link info
+ */
+struct mlo_tgt_partner_info {
+	uint8_t num_partner_links;
+	struct mlo_tgt_link_info link_info[WLAN_UMAC_MLO_MAX_VDEVS];
+};
+
 /*
  * struct mlo_mlme_ext_ops - MLME callback functions
  * @mlo_mlme_ext_validate_conn_req: Callback to validate connect request
@@ -296,6 +323,7 @@ struct mlo_partner_info {
  * @mlo_mlme_ext_peer_delete: Callback to initiate link peer delete
  * @mlo_mlme_ext_assoc_resp: Callback to initiate assoc resp
  * @mlo_mlme_get_link_assoc_req: Calback to get link assoc req buffer
+ * @mlo_mlme_ext_deauth: Callback to initiate deauth
  */
 struct mlo_mlme_ext_ops {
 	QDF_STATUS (*mlo_mlme_ext_validate_conn_req)(
@@ -312,5 +340,6 @@ struct mlo_mlme_ext_ops {
 	void (*mlo_mlme_ext_assoc_resp)(struct wlan_objmgr_peer *peer);
 	qdf_nbuf_t (*mlo_mlme_get_link_assoc_req)(struct wlan_objmgr_peer *peer,
 						  uint8_t link_ix);
+	void (*mlo_mlme_ext_deauth)(struct wlan_objmgr_peer *peer);
 };
 #endif

umac/mlo_mgr/src/wlan_mlo_mgr_aid.c

@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2021, The Linux Foundation. All rights reserved.
+ * Copyright (c) 2021 Qualcomm Innovation Center, Inc. All rights reserved.
  *
  * Permission to use, copy, modify, and/or distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -44,6 +45,7 @@ static uint16_t wlan_mlo_peer_alloc_aid(
 	uint16_t i, j;
 	struct wlan_vdev_aid_mgr *vdev_aid_mgr;
 	uint16_t first_aid = 0;
+	uint16_t start_aid;
 	struct mlo_mgr_context *mlo_mgr_ctx = wlan_objmgr_get_mlo_ctx();
 
 	if (!mlo_mgr_ctx)
@@ -54,7 +56,8 @@ static uint16_t wlan_mlo_peer_alloc_aid(
 	/* TODO check locking strategy */
 	ml_aid_lock_acquire(mlo_mgr_ctx);
 
-	for (i = 0; i < ml_aid_mgr->max_aid; i++) {
+	start_aid = ml_aid_mgr->start_aid;
+	for (i = start_aid; i < ml_aid_mgr->max_aid; i++) {
 		if (qdf_test_bit(i, ml_aid_mgr->aid_bitmap))
 			continue;
 
@@ -77,7 +80,7 @@ static uint16_t wlan_mlo_peer_alloc_aid(
 			}
 		} else {
 			vdev_aid_mgr = ml_aid_mgr->aid_mgr[link_ix];
-			if (vdev_aid_mgr)
+			if (!vdev_aid_mgr)
 				break;
 
 			if (qdf_test_bit(i, vdev_aid_mgr->aid_bitmap))
@@ -89,7 +92,9 @@ static uint16_t wlan_mlo_peer_alloc_aid(
 			for (j = 0; j < WLAN_UMAC_MLO_MAX_VDEVS; j++) {
 				if (j == link_ix)
 					continue;
-
+				/* Check whether this bit used by other VDEV
+				 * Non-MLO peers
+				 */
 				vdev_aid_mgr = ml_aid_mgr->aid_mgr[j];
 				if (vdev_aid_mgr &&
 				    qdf_test_bit(i, vdev_aid_mgr->aid_bitmap)) {
@@ -97,6 +102,9 @@ static uint16_t wlan_mlo_peer_alloc_aid(
 					break;
 				}
 			}
+			/* Assoc ID is used by other link, return this aid
+			 * to caller
+			 */
 			if (assoc_id == i + 1) {
 				vdev_aid_mgr = ml_aid_mgr->aid_mgr[link_ix];
 				qdf_set_bit(i, vdev_aid_mgr->aid_bitmap);
@@ -108,10 +116,13 @@ static uint16_t wlan_mlo_peer_alloc_aid(
 
 	if ((!is_mlo_peer) && first_aid) {
 		vdev_aid_mgr = ml_aid_mgr->aid_mgr[link_ix];
-		qdf_set_bit(first_aid, vdev_aid_mgr->aid_bitmap);
+		qdf_set_bit(first_aid - 1, vdev_aid_mgr->aid_bitmap);
 		assoc_id = first_aid;
 	}
 
+	if ((assoc_id == (uint16_t)-1) && (i == ml_aid_mgr->max_aid))
+		mlo_err("MLO aid allocation failed (reached max)");
+
 	ml_aid_lock_release(mlo_mgr_ctx);
 
 	return assoc_id;
@@ -123,6 +134,7 @@ static uint16_t wlan_mlme_peer_alloc_aid(
 {
 	uint16_t assoc_id = (uint16_t)-1;
 	uint16_t i;
+	uint16_t start_aid;
 	struct mlo_mgr_context *mlo_mgr_ctx = wlan_objmgr_get_mlo_ctx();
 
 	if (!mlo_mgr_ctx)
@@ -131,20 +143,108 @@ static uint16_t wlan_mlme_peer_alloc_aid(
 	if (!no_lock)
 		ml_aid_lock_acquire(mlo_mgr_ctx);
 
-	for (i = 0; i < vdev_aid_mgr->max_aid; i++) {
+	start_aid = vdev_aid_mgr->start_aid;
+	for (i = start_aid; i < vdev_aid_mgr->max_aid; i++) {
 		if (qdf_test_bit(i, vdev_aid_mgr->aid_bitmap))
 			continue;
 
 		assoc_id = i + 1;
 		qdf_set_bit(i, vdev_aid_mgr->aid_bitmap);
+		break;
 	}
 
 	if (!no_lock)
 		ml_aid_lock_release(mlo_mgr_ctx);
 
+	if (i == vdev_aid_mgr->max_aid)
+		return (uint16_t)-1;
+
 	return assoc_id;
 }
 
+static QDF_STATUS wlan_mlo_peer_set_aid(
+		struct wlan_ml_vdev_aid_mgr *ml_aid_mgr,
+		bool is_mlo_peer,
+		uint8_t link_ix,
+		uint16_t assoc_id)
+{
+	uint16_t j;
+	struct wlan_vdev_aid_mgr *vdev_aid_mgr;
+	struct mlo_mgr_context *mlo_mgr_ctx = wlan_objmgr_get_mlo_ctx();
+
+	if (!mlo_mgr_ctx)
+		return QDF_STATUS_E_FAILURE;
+
+	if (!is_mlo_peer && link_ix == 0xff)
+		return QDF_STATUS_E_FAILURE;
+	/* TODO check locking strategy */
+	ml_aid_lock_acquire(mlo_mgr_ctx);
+
+	if (qdf_test_bit(WLAN_AID(assoc_id) - 1,  ml_aid_mgr->aid_bitmap)) {
+		ml_aid_lock_release(mlo_mgr_ctx);
+		return QDF_STATUS_E_FAILURE;
+	}
+
+	if (is_mlo_peer) {
+		for (j = 0; j < WLAN_UMAC_MLO_MAX_VDEVS; j++) {
+			vdev_aid_mgr = ml_aid_mgr->aid_mgr[j];
+			if (vdev_aid_mgr &&
+			    qdf_test_bit(WLAN_AID(assoc_id) - 1,
+					 vdev_aid_mgr->aid_bitmap)) {
+				ml_aid_lock_release(mlo_mgr_ctx);
+				return QDF_STATUS_E_FAILURE;
+			}
+			/* AID is free */
+			if (j == WLAN_UMAC_MLO_MAX_VDEVS - 1)
+				mlo_peer_set_aid_bit(ml_aid_mgr,
+						     WLAN_AID(assoc_id) - 1);
+		}
+		qdf_set_bit(WLAN_AID(assoc_id) - 1, ml_aid_mgr->aid_bitmap);
+	} else {
+		vdev_aid_mgr = ml_aid_mgr->aid_mgr[link_ix];
+		if (!vdev_aid_mgr) {
+			ml_aid_lock_release(mlo_mgr_ctx);
+			return QDF_STATUS_E_FAILURE;
+		}
+
+		if (qdf_test_bit(WLAN_AID(assoc_id) - 1,
+				 vdev_aid_mgr->aid_bitmap)) {
+			ml_aid_lock_release(mlo_mgr_ctx);
+			return QDF_STATUS_E_FAILURE;
+		}
+
+		qdf_set_bit(WLAN_AID(assoc_id) - 1, vdev_aid_mgr->aid_bitmap);
+	}
+
+	ml_aid_lock_release(mlo_mgr_ctx);
+
+	return QDF_STATUS_SUCCESS;
+}
+
+static QDF_STATUS wlan_mlme_peer_set_aid(
+		struct wlan_vdev_aid_mgr *vdev_aid_mgr,
+		bool no_lock, uint16_t assoc_id)
+{
+	struct mlo_mgr_context *mlo_mgr_ctx = wlan_objmgr_get_mlo_ctx();
+	QDF_STATUS status = QDF_STATUS_E_FAILURE;
+
+	if (!mlo_mgr_ctx)
+		return status;
+
+	if (!no_lock)
+		ml_aid_lock_acquire(mlo_mgr_ctx);
+
+	if (!qdf_test_bit(WLAN_AID(assoc_id) - 1, vdev_aid_mgr->aid_bitmap)) {
+		qdf_set_bit(WLAN_AID(assoc_id) - 1, vdev_aid_mgr->aid_bitmap);
+		status = QDF_STATUS_SUCCESS;
+	}
+
+	if (!no_lock)
+		ml_aid_lock_release(mlo_mgr_ctx);
+
+	return status;
+}
+
 QDF_STATUS wlan_mlo_peer_free_aid(
 		struct wlan_ml_vdev_aid_mgr *ml_aid_mgr,
 		uint8_t link_ix,
@@ -160,7 +260,7 @@ QDF_STATUS wlan_mlo_peer_free_aid(
 
 	/* TODO check locking strategy */
 	ml_aid_lock_acquire(mlo_mgr_ctx);
-	assoc_id_ix = assoc_id - 1;
+	assoc_id_ix = WLAN_AID(assoc_id) - 1;
 	if (qdf_test_bit(assoc_id_ix, ml_aid_mgr->aid_bitmap)) {
 		qdf_clear_bit(assoc_id_ix, ml_aid_mgr->aid_bitmap);
 		for (j = 0; j < WLAN_UMAC_MLO_MAX_VDEVS; j++) {
@@ -183,6 +283,26 @@ QDF_STATUS wlan_mlo_peer_free_aid(
 	return QDF_STATUS_SUCCESS;
 }
 
+static int wlan_mlme_peer_aid_is_set(struct wlan_vdev_aid_mgr *vdev_aid_mgr,
+				     bool no_lock, uint16_t assoc_id)
+{
+	struct mlo_mgr_context *mlo_mgr_ctx = wlan_objmgr_get_mlo_ctx();
+	int isset = 0;
+
+	if (!mlo_mgr_ctx)
+		return isset;
+
+	if (!no_lock)
+		ml_aid_lock_acquire(mlo_mgr_ctx);
+
+	isset = qdf_test_bit(WLAN_AID(assoc_id) - 1, vdev_aid_mgr->aid_bitmap);
+
+	if (!no_lock)
+		ml_aid_lock_release(mlo_mgr_ctx);
+
+	return isset;
+}
+
 void wlan_mlme_peer_free_aid(
 		struct wlan_vdev_aid_mgr *vdev_aid_mgr,
 		bool no_lock, uint16_t assoc_id)
@@ -195,12 +315,30 @@ void wlan_mlme_peer_free_aid(
 	if (!no_lock)
 		ml_aid_lock_acquire(mlo_mgr_ctx);
 
-	qdf_clear_bit(assoc_id - 1, vdev_aid_mgr->aid_bitmap);
+	qdf_clear_bit(WLAN_AID(assoc_id) - 1, vdev_aid_mgr->aid_bitmap);
 
 	if (!no_lock)
 		ml_aid_lock_release(mlo_mgr_ctx);
 }
 
+uint16_t wlan_mlme_get_aid_count(struct wlan_objmgr_vdev *vdev)
+{
+	uint16_t i;
+	uint16_t aid_count = 0;
+	struct wlan_vdev_aid_mgr *vdev_aid_mgr;
+
+	vdev_aid_mgr = wlan_vdev_mlme_get_aid_mgr(vdev);
+	if (!vdev_aid_mgr)
+		return (uint16_t)-1;
+
+	for (i = 0; i < vdev_aid_mgr->max_aid; i++) {
+		if (qdf_test_bit(i, vdev_aid_mgr->aid_bitmap))
+			aid_count++;
+	}
+
+	return aid_count;
+}
+
 QDF_STATUS mlo_peer_allocate_aid(
 		struct wlan_mlo_dev_context *ml_dev,
 		struct wlan_mlo_peer_context *ml_peer)
@@ -268,17 +406,71 @@ QDF_STATUS mlo_free_aid(struct wlan_objmgr_vdev *vdev, uint16_t assoc_id)
 {
 	struct wlan_mlo_dev_context *ml_dev;
 	struct wlan_ml_vdev_aid_mgr *ml_aid_mgr;
+	struct wlan_mlo_ap *ap_ctx;
 
 	ml_dev = vdev->mlo_dev_ctx;
 
 	if (!ml_dev)
 		return QDF_STATUS_E_INVAL;
 
-	ml_aid_mgr = ml_dev->ap_ctx->ml_aid_mgr;
+	ap_ctx = ml_dev->ap_ctx;
+	if (!ap_ctx)
+		return QDF_STATUS_E_INVAL;
+
+	ml_aid_mgr = ap_ctx->ml_aid_mgr;
 	if (!ml_aid_mgr)
 		return QDF_STATUS_E_INVAL;
 
-	return wlan_mlo_peer_free_aid(ml_aid_mgr, true, assoc_id);
+	return wlan_mlo_peer_free_aid(ml_aid_mgr, 0xff, assoc_id);
+}
+
+static QDF_STATUS mlo_peer_set_aid(struct wlan_mlo_dev_context *ml_dev,
+				   uint16_t assoc_id)
+{
+	struct wlan_ml_vdev_aid_mgr *ml_aid_mgr;
+	QDF_STATUS status;
+	struct wlan_mlo_ap *ap_ctx;
+
+	ap_ctx = ml_dev->ap_ctx;
+	if (!ap_ctx)
+		return QDF_STATUS_E_FAILURE;
+
+	ml_aid_mgr = ap_ctx->ml_aid_mgr;
+	if (!ml_aid_mgr)
+		return QDF_STATUS_E_FAILURE;
+
+	status = wlan_mlo_peer_set_aid(ml_aid_mgr, true, 0xff, assoc_id);
+
+	return status;
+}
+
+QDF_STATUS mlo_set_aid(struct wlan_objmgr_vdev *vdev, uint16_t assoc_id)
+{
+	struct wlan_mlo_dev_context *ml_dev;
+
+	ml_dev = vdev->mlo_dev_ctx;
+
+	if (!ml_dev)
+		return QDF_STATUS_E_FAILURE;
+
+	return mlo_peer_set_aid(ml_dev, assoc_id);
+}
+
+int mlme_is_aid_set(struct wlan_objmgr_vdev *vdev, uint16_t assoc_id)
+{
+	struct wlan_vdev_aid_mgr *vdev_aid_mgr;
+	bool no_lock = true;
+
+	vdev_aid_mgr = wlan_vdev_mlme_get_aid_mgr(vdev);
+	if (vdev_aid_mgr) {
+		if (qdf_atomic_read(&vdev_aid_mgr->ref_cnt) > 1)
+			no_lock = false;
+
+		return wlan_mlme_peer_aid_is_set(vdev_aid_mgr, no_lock,
+						 assoc_id);
+	}
+
+	return 0;
 }
 
 uint16_t mlme_get_aid(struct wlan_objmgr_vdev *vdev)
@@ -313,6 +505,39 @@ uint16_t mlme_get_aid(struct wlan_objmgr_vdev *vdev)
 	return assoc_id;
 }
 
+QDF_STATUS mlme_set_aid(struct wlan_objmgr_vdev *vdev,
+			uint16_t assoc_id)
+{
+	struct wlan_mlo_dev_context *ml_dev;
+	struct wlan_ml_vdev_aid_mgr *ml_aid_mgr;
+	struct wlan_vdev_aid_mgr *vdev_aid_mgr;
+	bool no_lock = true;
+	uint8_t link_id;
+
+	ml_dev = vdev->mlo_dev_ctx;
+
+	if (!ml_dev) {
+		vdev_aid_mgr = wlan_vdev_mlme_get_aid_mgr(vdev);
+		if (vdev_aid_mgr) {
+			if (qdf_atomic_read(&vdev_aid_mgr->ref_cnt) > 1)
+				no_lock = false;
+
+			return wlan_mlme_peer_set_aid(vdev_aid_mgr, no_lock,
+						      assoc_id);
+		} else {
+			return QDF_STATUS_E_FAILURE;
+		}
+	}
+
+	ml_aid_mgr = ml_dev->ap_ctx->ml_aid_mgr;
+	if (!ml_aid_mgr)
+		return QDF_STATUS_E_FAILURE;
+
+	link_id = mlo_get_link_vdev_ix(ml_dev, vdev);
+
+	return wlan_mlo_peer_set_aid(ml_aid_mgr, false, link_id, assoc_id);
+}
+
 void mlme_free_aid(struct wlan_objmgr_vdev *vdev, uint16_t assoc_id)
 {
 	struct wlan_mlo_dev_context *ml_dev;
@@ -344,6 +569,54 @@ void mlme_free_aid(struct wlan_objmgr_vdev *vdev, uint16_t assoc_id)
 	wlan_mlo_peer_free_aid(ml_aid_mgr, link_id, assoc_id);
 }
 
+void wlan_vdev_mlme_aid_mgr_max_aid_set(struct wlan_objmgr_vdev *vdev,
+					uint16_t max_aid)
+{
+	struct wlan_vdev_aid_mgr *aid_mgr;
+
+	aid_mgr = wlan_vdev_mlme_get_aid_mgr(vdev);
+	if (!aid_mgr)
+		return;
+
+	aid_mgr->max_aid = max_aid;
+}
+
+QDF_STATUS wlan_vdev_mlme_set_start_aid(struct wlan_objmgr_vdev *vdev,
+					uint16_t start_aid)
+{
+	struct wlan_vdev_aid_mgr *vdev_aid_mgr;
+	struct wlan_ml_vdev_aid_mgr *ml_aid_mgr;
+	struct wlan_mlo_dev_context *ml_dev;
+	uint16_t j, max_aid_start = 0;
+
+	vdev_aid_mgr = wlan_vdev_mlme_get_aid_mgr(vdev);
+	if (!vdev_aid_mgr)
+		return QDF_STATUS_E_FAILURE;
+
+	vdev_aid_mgr->start_aid = start_aid;
+
+	ml_dev = vdev->mlo_dev_ctx;
+	if (ml_dev) {
+		ml_aid_mgr = ml_dev->ap_ctx->ml_aid_mgr;
+		if (!ml_aid_mgr)
+			return QDF_STATUS_E_FAILURE;
+
+		/* Derive higher start_aid */
+		for (j = 0; j < WLAN_UMAC_MLO_MAX_VDEVS; j++) {
+			vdev_aid_mgr = ml_aid_mgr->aid_mgr[j];
+			if (!vdev_aid_mgr)
+				continue;
+
+			if (max_aid_start < vdev_aid_mgr->start_aid)
+				max_aid_start = vdev_aid_mgr->start_aid;
+		}
+
+		ml_aid_mgr->start_aid = max_aid_start;
+	}
+
+	return QDF_STATUS_SUCCESS;
+}
+
 struct wlan_vdev_aid_mgr *wlan_vdev_aid_mgr_init(uint16_t max_aid)
 {
 	struct wlan_vdev_aid_mgr *aid_mgr;
@@ -352,6 +625,7 @@ struct wlan_vdev_aid_mgr *wlan_vdev_aid_mgr_init(uint16_t max_aid)
 	if (!aid_mgr)
 		return NULL;
 
+	aid_mgr->start_aid = 0;
 	aid_mgr->max_aid = max_aid;
 	qdf_atomic_init(&aid_mgr->ref_cnt);
 	/* Take reference before returning */
@@ -482,6 +756,7 @@ QDF_STATUS wlan_mlo_vdev_aid_mgr_init(struct wlan_mlo_dev_context *ml_dev)
 		return QDF_STATUS_E_NOMEM;
 	}
 
+	ml_aidmgr->start_aid = 0;
 	ml_aidmgr->max_aid = max_aid;
 	ml_dev->ap_ctx->ml_aid_mgr = ml_aidmgr;
 

umac/mlo_mgr/src/wlan_mlo_mgr_cmn.c

@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2021, The Linux Foundation. All rights reserved.
+ * Copyright (c) 2021 Qualcomm Innovation Center, Inc. All rights reserved.
  *
  * Permission to use, copy, modify, and/or distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -207,6 +208,17 @@ qdf_nbuf_t mlo_mlme_get_link_assoc_req(struct wlan_objmgr_peer *peer,
 	return mlo_ctx->mlme_ops->mlo_mlme_get_link_assoc_req(peer, link_ix);
 }
 
+void mlo_mlme_peer_deauth(struct wlan_objmgr_peer *peer)
+{
+	struct mlo_mgr_context *mlo_ctx = wlan_objmgr_get_mlo_ctx();
+
+	if (!mlo_ctx || !mlo_ctx->mlme_ops ||
+	    !mlo_ctx->mlme_ops->mlo_mlme_ext_deauth)
+		return;
+
+	mlo_ctx->mlme_ops->mlo_mlme_ext_deauth(peer);
+}
+
 uint8_t mlo_get_link_vdev_ix(struct wlan_mlo_dev_context *ml_dev,
 			     struct wlan_objmgr_vdev *vdev)
 {

umac/mlo_mgr/src/wlan_mlo_mgr_main.c

@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2021, The Linux Foundation. All rights reserved.
+ * Copyright (c) 2021 Qualcomm Innovation Center, Inc. All rights reserved.
  *
  * Permission to use, copy, modify, and/or distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -25,6 +26,7 @@
 #include <wlan_mlo_mgr_ap.h>
 #include <wlan_mlo_mgr_peer.h>
 #include <wlan_cm_public_struct.h>
+#include "wlan_mlo_mgr_msgq.h"
 
 static void mlo_global_ctx_deinit(void)
 {
@@ -36,6 +38,7 @@ static void mlo_global_ctx_deinit(void)
 	if (qdf_list_empty(&mlo_mgr_ctx->ml_dev_list))
 		mlo_err("ML dev list is not empty");
 
+	mlo_msgq_free();
 	ml_peerid_lock_destroy(mlo_mgr_ctx);
 	ml_link_lock_destroy(mlo_mgr_ctx);
 	ml_aid_lock_destroy(mlo_mgr_ctx);
@@ -68,6 +71,8 @@ static void mlo_global_ctx_init(void)
 	ml_peerid_lock_create(mlo_mgr_ctx);
 	ml_link_lock_create(mlo_mgr_ctx);
 	ml_aid_lock_create(mlo_mgr_ctx);
+	mlo_mgr_ctx->mlo_is_force_primary_umac = 0;
+	mlo_msgq_init();
 }
 
 QDF_STATUS wlan_mlo_mgr_init(void)

umac/mlo_mgr/src/wlan_mlo_mgr_msgq.c

@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2021, The Linux Foundation. All rights reserved.
+ * Copyright (c) 2021 Qualcomm Innovation Center, Inc. All rights reserved.
  *
  * Permission to use, copy, modify, and/or distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -29,6 +30,7 @@ QDF_STATUS mlo_msgq_post(enum mlo_msg_type type,
 	struct peer_assoc_notify_s *peer_assoc;
 	struct peer_assoc_fail_notify_s *peer_assoc_fail;
 	struct peer_discon_notify_s *peer_disconn;
+	struct peer_deauth_notify_s *peer_deauth;
 
 	switch (type) {
 	case MLO_PEER_CREATE:
@@ -64,13 +66,67 @@ QDF_STATUS mlo_msgq_post(enum mlo_msg_type type,
 					     WLAN_MLO_MGR_ID);
 		break;
 
+	case MLO_PEER_DEAUTH:
+		peer_deauth = (struct peer_deauth_notify_s *)payload;
+		mlo_mlme_peer_deauth(peer_deauth->peer);
+		wlan_objmgr_peer_release_ref(peer_deauth->peer,
+					     WLAN_MLO_MGR_ID);
+		break;
+
 	default:
 		break;
 	}
 
 	return QDF_STATUS_SUCCESS;
 }
+
+void mlo_msgq_init(void)
+{
+}
+
+void mlo_msgq_free(void)
+{
+}
 #else
+static void mlo_msgq_timer_start(void)
+{
+	struct ctxt_switch_mgr *msgq_ctx;
+	struct mlo_mgr_context *mlo_ctx = wlan_objmgr_get_mlo_ctx();
+	bool start_timer = true;
+
+	if (!mlo_ctx)
+		return;
+
+	msgq_ctx = mlo_ctx->msgq_ctx;
+
+	qdf_spin_lock_bh(&msgq_ctx->ctxt_lock);
+	if (!msgq_ctx->timer_started)
+		msgq_ctx->timer_started = true;
+	else
+		start_timer = false;
+	qdf_spin_unlock_bh(&msgq_ctx->ctxt_lock);
+
+	if (start_timer)
+		qdf_timer_start(&msgq_ctx->ctxt_mgr_timer, 0);
+}
+
+static void mlo_msgq_timer_stop(void)
+{
+	struct ctxt_switch_mgr *msgq_ctx;
+	struct mlo_mgr_context *mlo_ctx = wlan_objmgr_get_mlo_ctx();
+
+	if (!mlo_ctx)
+		return;
+
+	msgq_ctx = mlo_ctx->msgq_ctx;
+
+	qdf_timer_stop(&msgq_ctx->ctxt_mgr_timer);
+
+	qdf_spin_lock_bh(&msgq_ctx->ctxt_lock);
+	msgq_ctx->timer_started = false;
+	qdf_spin_unlock_bh(&msgq_ctx->ctxt_lock);
+}
+
 QDF_STATUS mlo_msgq_post(enum mlo_msg_type type,
 			 struct wlan_mlo_dev_context *ml_dev,
 			 void *payload)
@@ -80,6 +136,16 @@ QDF_STATUS mlo_msgq_post(enum mlo_msg_type type,
 	struct peer_assoc_notify_s *peer_assoc, *peer_assoc_l;
 	struct peer_assoc_fail_notify_s *peer_assoc_fail, *peer_assoc_fail_l;
 	struct peer_discon_notify_s *peer_disconn, *peer_disconn_l;
+	struct ctxt_switch_mgr *msgq_ctx;
+	struct mlo_mgr_context *mlo_ctx = wlan_objmgr_get_mlo_ctx();
+
+	if (!mlo_ctx)
+		return QDF_STATUS_E_FAILURE;
+
+	msgq_ctx = mlo_ctx->msgq_ctx;
+
+	if (!msgq_ctx->allow_msg)
+		return QDF_STATUS_E_FAILURE;
 
 	msg = qdf_mem_malloc(sizeof(*msg));
 	if (!msg)
@@ -116,15 +182,25 @@ QDF_STATUS mlo_msgq_post(enum mlo_msg_type type,
 		peer_disconn->peer = peer_disconn_l->peer;
 		break;
 
+	case MLO_PEER_DEAUTH:
+		peer_deauth = &msg->m.peer_deauth;
+		peer_deauth_l = (struct peer_deauth_notify_s *)payload;
+		peer_deauth->peer = peer_deauth_l->peer;
+		break;
+
 	default:
 		break;
 	}
-	/* TODO queue message buffer to qdf_list */
+
+	qdf_spin_lock_bh(&msgq_ctx->ctxt_lock);
+	qdf_list_insert_back(&msgq_ctx->msgq_list, &msg->node);
+	qdf_spin_unlock_bh(&msgq_ctx->ctxt_lock);
+	mlo_msgq_timer_start();
 
 	return QDF_STATUS_SUCCESS;
 }
 
-void mlo_msgq_msg_process_hdlr(struct mlo_ctxt_switch_msg_s *msg)
+static void mlo_msgq_msg_process_hdlr(struct mlo_ctxt_switch_msg_s *msg)
 {
 	enum mlo_msg_type type;
 	struct peer_create_notif_s *peer_create;
@@ -166,13 +242,20 @@ void mlo_msgq_msg_process_hdlr(struct mlo_ctxt_switch_msg_s *msg)
 					     WLAN_MLO_MGR_ID);
 		break;
 
+	case MLO_PEER_DEAUTH:
+		peer_deauth = &msg->m.peer_deauth;
+		mlo_mlme_peer_deauth(peer_deauth->peer);
+		wlan_objmgr_peer_release_ref(peer_deauth->peer,
+					     WLAN_MLO_MGR_ID);
+		break;
+
 	default:
 		break;
 	}
 	qdf_mem_free(msg);
 }
 
-void mlo_msgq_msg_flush_hdlr(struct mlo_ctxt_switch_msg_s *msg)
+static void mlo_msgq_msg_flush_hdlr(struct mlo_ctxt_switch_msg_s *msg)
 {
 	enum mlo_msg_type type;
 	struct peer_create_notif_s *peer_create;
@@ -208,9 +291,128 @@ void mlo_msgq_msg_flush_hdlr(struct mlo_ctxt_switch_msg_s *msg)
 					     WLAN_MLO_MGR_ID);
 		break;
 
+	case MLO_PEER_DEAUTH:
+		peer_deauth = &msg->m.peer_deauth;
+		wlan_objmgr_peer_release_ref(peer_deauth->peer,
+					     WLAN_MLO_MGR_ID);
+		break;
+
 	default:
 		break;
 	}
 	qdf_mem_free(msg);
 }
+
+static void mlo_msgq_msg_flush(void)
+{
+	struct ctxt_switch_mgr *msgq_ctx;
+	struct mlo_mgr_context *mlo_ctx = wlan_objmgr_get_mlo_ctx();
+	qdf_list_node_t *msgbuf_node = NULL;
+	struct mlo_ctxt_switch_msg_s *msg;
+	QDF_STATUS status;
+
+	if (!mlo_ctx)
+		return;
+
+	msgq_ctx = mlo_ctx->msgq_ctx;
+	do {
+		msg = NULL;
+		qdf_spin_lock_bh(&msgq_ctx->ctxt_lock);
+		status = qdf_list_peek_front(&msgq_ctx->msgq_list,
+					     &msgbuf_node);
+		if (status != QDF_STATUS_E_EMPTY) {
+			qdf_list_remove_node(&msgq_ctx->msgq_list,
+					     msgbuf_node);
+			msg = qdf_container_of(msgbuf_node,
+					       struct mlo_ctxt_switch_msg_s,
+					       node);
+		}
+		qdf_spin_unlock_bh(&msgq_ctx->ctxt_lock);
+
+		if (!msg)
+			break;
+
+		mlo_msgq_msg_flush_hdlr(msg);
+
+	} while (true);
+}
+
+static void mlo_msgq_msg_handler(void *arg)
+{
+	struct ctxt_switch_mgr *msgq_ctx;
+	struct mlo_mgr_context *mlo_ctx = wlan_objmgr_get_mlo_ctx();
+	qdf_list_node_t *msgbuf_node = NULL;
+	struct mlo_ctxt_switch_msg_s *msg;
+	QDF_STATUS status;
+
+	if (!mlo_ctx)
+		return;
+
+	msgq_ctx = mlo_ctx->msgq_ctx;
+	do {
+		msg = NULL;
+		qdf_spin_lock_bh(&msgq_ctx->ctxt_lock);
+		status = qdf_list_peek_front(&msgq_ctx->msgq_list,
+					     &msgbuf_node);
+		if (status != QDF_STATUS_E_EMPTY) {
+			qdf_list_remove_node(&msgq_ctx->msgq_list,
+					     &msgbuf_node);
+			msg = qdf_container_of(msgbuf_node,
+					       struct mlo_ctxt_switch_msg_s,
+					       node);
+		} else {
+			msgq_ctx->timer_started = false;
+		}
+		qdf_spin_unlock_bh(&msgq_ctx->ctxt_lock);
+
+		if (!msg)
+			break;
+
+		mlo_msgq_msg_process_hdlr(msg);
+
+	} while (true);
+}
+
+void mlo_msgq_init(void)
+{
+	struct ctxt_switch_mgr *msgq_ctx;
+	struct mlo_mgr_context *mlo_ctx = wlan_objmgr_get_mlo_ctx();
+
+	msgq_ctx = qdf_mem_malloc(sizeof(*msgq_ctx));
+	if (!msgq_ctx) {
+		mlo_err(" message queue context allocation failed");
+		return;
+	}
+
+	qdf_spinlock_create(&msgq_ctx->ctxt_lock);
+	/* Initialize timer with timeout handler */
+	qdf_timer_init(NULL, &msgq_ctx->ctxt_mgr_timer,
+		       mlo_msgq_msg_handler,
+		       NULL, QDF_TIMER_TYPE_WAKE_APPS);
+
+	msgq_ctx->timer_started = false;
+	msgq_ctx->allow_msg = true;
+	qdf_list_create(&msgq_ctx->msgq_list, MLO_MAX_MSGQ_SIZE);
+
+	mlo_ctx->msgq_ctx = msgq_ctx;
+}
+
+void mlo_msgq_free(void)
+{
+	struct ctxt_switch_mgr *msgq_ctx;
+	struct mlo_mgr_context *mlo_ctx = wlan_objmgr_get_mlo_ctx();
+
+	if (!mlo_ctx)
+		return;
+
+	msgq_ctx = mlo_ctx->msgq_ctx;
+
+	msgq_ctx->timer_started = false;
+	msgq_ctx->allow_msg = false;
+	mlo_msgq_msg_flush();
+	qdf_list_destroy(&msgq_ctx->msgq_list);
+	qdf_timer_free(&msgq_ctx->ctxt_mgr_timer);
+	qdf_spinlock_destroy(&msgq_ctx->ctxt_lock);
+	qdf_mem_free(msgq_ctx);
+}
 #endif

umac/mlo_mgr/src/wlan_mlo_mgr_peer.c

@@ -16,13 +16,14 @@
  */
 
 #include "wlan_mlo_mgr_main.h"
+#include "qdf_module.h"
 #include "qdf_types.h"
 #include "wlan_cmn.h"
 #include "wlan_mlo_mgr_msgq.h"
 #include "wlan_objmgr_peer_obj.h"
 #include "wlan_mlo_mgr_peer.h"
 #include "wlan_mlo_mgr_ap.h"
-#include "qdf_module.h"
+#include "wlan_crypto_global_api.h"
 
 static void mlo_partner_peer_create_post(struct wlan_mlo_dev_context *ml_dev,
 					 struct wlan_objmgr_vdev *vdev_link,
@@ -117,6 +118,18 @@ static void mlo_link_peer_disconnect_notify(struct wlan_mlo_dev_context *ml_dev,
 	}
 }
 
+static void mlo_link_peer_deauth_init(struct wlan_mlo_dev_context *ml_dev,
+				      struct wlan_objmgr_peer *peer)
+{
+	struct peer_deauth_notify_s peer_deauth;
+	QDF_STATUS status;
+
+	peer_deauth.peer = peer;
+	status = mlo_msgq_post(MLO_PEER_DEAUTH, ml_dev, &peer_deauth);
+	if (status != QDF_STATUS_SUCCESS)
+		wlan_objmgr_peer_release_ref(peer, WLAN_MLO_MGR_ID);
+}
+
 QDF_STATUS
 wlan_mlo_peer_is_disconnect_progress(struct wlan_mlo_peer_context *ml_peer)
 {
@@ -168,6 +181,40 @@ struct wlan_objmgr_peer *wlan_mlo_peer_get_assoc_peer(
 	return assoc_peer;
 }
 
+bool mlo_peer_is_assoc_peer(struct wlan_mlo_peer_context *ml_peer,
+			    struct wlan_objmgr_peer *peer)
+{
+	struct wlan_mlo_link_peer_entry *peer_entry;
+	bool is_assoc_peer = false;
+
+	if (!ml_peer || !peer)
+		return is_assoc_peer;
+
+	peer_entry = &ml_peer->peer_list[0];
+
+	if (peer_entry->link_peer != peer)
+		is_assoc_peer = true;
+
+	return is_assoc_peer;
+}
+
+bool wlan_mlo_peer_is_assoc_peer(struct wlan_mlo_peer_context *ml_peer,
+				 struct wlan_objmgr_peer *peer)
+{
+	bool is_assoc_peer = false;
+
+	if (!ml_peer || !peer)
+		return is_assoc_peer;
+
+	mlo_peer_lock_acquire(ml_peer);
+
+	is_assoc_peer = mlo_peer_is_assoc_peer(ml_peer, peer);
+
+	mlo_peer_lock_release(ml_peer);
+
+	return is_assoc_peer;
+}
+
 void wlan_mlo_partner_peer_assoc_post(struct wlan_objmgr_peer *assoc_peer)
 {
 	struct wlan_mlo_dev_context *ml_dev;
@@ -217,6 +264,59 @@ void wlan_mlo_partner_peer_assoc_post(struct wlan_objmgr_peer *assoc_peer)
 	}
 }
 
+void
+wlan_mlo_peer_deauth_init(struct wlan_mlo_peer_context *ml_peer)
+{
+	struct wlan_mlo_dev_context *ml_dev;
+	struct wlan_objmgr_peer *link_peer;
+	struct wlan_objmgr_peer *link_peers[MAX_MLO_LINK_PEERS];
+	struct wlan_mlo_link_peer_entry *peer_entry;
+	uint16_t i;
+
+	mlo_peer_lock_acquire(ml_peer);
+
+	if (ml_peer->mlpeer_state == ML_PEER_DISCONN_INITIATED) {
+		mlo_peer_lock_release(ml_peer);
+		return;
+	}
+
+	ml_peer->mlpeer_state = ML_PEER_DISCONN_INITIATED;
+	ml_dev = ml_peer->ml_dev;
+
+	for (i = 0; i < MAX_MLO_LINK_PEERS; i++) {
+		link_peers[i] = NULL;
+		peer_entry = &ml_peer->peer_list[i];
+		if (!peer_entry->link_peer)
+			continue;
+
+		link_peer = peer_entry->link_peer;
+		if (wlan_objmgr_peer_try_get_ref(link_peer, WLAN_MLO_MGR_ID) !=
+						QDF_STATUS_SUCCESS)
+			continue;
+
+		link_peers[i] = link_peer;
+	}
+	mlo_peer_lock_release(ml_peer);
+
+	for (i = 0; i < MAX_MLO_LINK_PEERS; i++) {
+		if (!link_peers[i])
+			continue;
+
+		/* Prepare and queue message */
+		if (i == 0) {
+			/* Skip Deauth if PMF is enabled for the station */
+			if (wlan_crypto_is_pmf_enabled(
+					wlan_peer_get_vdev(link_peers[i]),
+					link_peers[i]))
+				break;
+
+			mlo_link_peer_deauth_init(ml_dev, link_peers[i]);
+		} else {
+			mlo_link_peer_disconnect_notify(ml_dev, link_peers[i]);
+		}
+	}
+}
+
 void
 wlan_mlo_partner_peer_create_failed_notify(
 				struct wlan_mlo_peer_context *ml_peer)
@@ -274,6 +374,9 @@ void wlan_mlo_partner_peer_disconnect_notify(struct wlan_objmgr_peer *src_peer)
 	uint16_t i;
 
 	ml_peer = src_peer->mlo_peer_ctx;
+	if (!ml_peer)
+		return;
+
 	mlo_peer_lock_acquire(ml_peer);
 
 	if (ml_peer->mlpeer_state == ML_PEER_DISCONN_INITIATED) {
@@ -331,7 +434,7 @@ static void mlo_reset_link_peer(
 	mlo_peer_lock_release(ml_peer);
 }
 
-void mlo_peer_free(struct wlan_mlo_peer_context *ml_peer)
+static void mlo_peer_free(struct wlan_mlo_peer_context *ml_peer)
 {
 	struct wlan_mlo_dev_context *ml_dev;
 
@@ -345,18 +448,41 @@ void mlo_peer_free(struct wlan_mlo_peer_context *ml_peer)
 	mlo_ap_ml_peerid_free(ml_peer->mlo_peer_id);
 	mlo_peer_free_aid(ml_dev, ml_peer);
 	mlo_peer_free_primary_umac(ml_dev, ml_peer);
-	mlo_dev_mlpeer_detach(ml_dev, ml_peer);
 	qdf_mem_free(ml_peer);
 }
 
+void mlo_peer_cleanup(struct wlan_mlo_peer_context *ml_peer)
+{
+	struct wlan_mlo_dev_context *ml_dev;
+
+	ml_dev = ml_peer->ml_dev;
+	if (!ml_dev) {
+		mlo_err("ML DEV is NULL");
+		return;
+	}
+
+	mlo_dev_mlpeer_detach(ml_dev, ml_peer);
+	mlo_peer_free(ml_peer);
+}
+
 static QDF_STATUS mlo_peer_attach_link_peer(
 		struct wlan_mlo_peer_context *ml_peer,
-		struct wlan_objmgr_peer *link_peer)
+		struct wlan_objmgr_peer *link_peer,
+		qdf_nbuf_t frm_buf)
 {
 	struct wlan_mlo_link_peer_entry *peer_entry;
 	QDF_STATUS status = QDF_STATUS_E_RESOURCES;
+	struct wlan_objmgr_pdev *pdev;
+	struct wlan_objmgr_vdev *vdev;
 	uint16_t i;
 
+	if (!link_peer)
+		return QDF_STATUS_E_FAILURE;
+
+	vdev = wlan_peer_get_vdev(link_peer);
+	if (!vdev)
+		return QDF_STATUS_E_FAILURE;
+
 	mlo_peer_lock_acquire(ml_peer);
 
 	if (ml_peer->mlpeer_state != ML_PEER_CREATED) {
@@ -377,10 +503,14 @@ static QDF_STATUS mlo_peer_attach_link_peer(
 		qdf_copy_macaddr(&peer_entry->link_addr,
 				 (struct qdf_mac_addr *)&link_peer->macaddr[0]);
 
-		peer_entry->link_ix = i + 1;
-		peer_entry->hw_link_id = 1;
-		/*wlan_peer_get_hw_link_id(link_peer)TODO*/
+		peer_entry->link_ix = wlan_vdev_get_link_id(vdev);
+		pdev = wlan_vdev_get_pdev(wlan_peer_get_vdev(link_peer));
+		peer_entry->hw_link_id = wlan_mlo_get_pdev_hw_link_id(pdev);
 		mlo_peer_assign_primary_umac(ml_peer, peer_entry);
+		if (frm_buf)
+			peer_entry->assoc_rsp_buf = frm_buf;
+		else
+			peer_entry->assoc_rsp_buf = NULL;
 
 		status = QDF_STATUS_SUCCESS;
 		break;
@@ -393,6 +523,67 @@ static QDF_STATUS mlo_peer_attach_link_peer(
 	return status;
 }
 
+qdf_nbuf_t mlo_peer_get_link_peer_assoc_resp_buf(
+		struct wlan_mlo_peer_context *ml_peer,
+		uint8_t link_ix)
+{
+	struct wlan_mlo_link_peer_entry *peer_entry;
+	qdf_nbuf_t frm_buf = NULL;
+	uint8_t i;
+
+	if (!ml_peer)
+		return NULL;
+
+	if (link_ix > MAX_MLO_LINK_PEERS)
+		return NULL;
+
+	mlo_peer_lock_acquire(ml_peer);
+	if ((ml_peer->mlpeer_state != ML_PEER_CREATED) &&
+	    (ml_peer->mlpeer_state != ML_PEER_ASSOC_DONE)) {
+		mlo_peer_lock_release(ml_peer);
+		return NULL;
+	}
+
+	for (i = 0; i < MAX_MLO_LINK_PEERS; i++) {
+		peer_entry = &ml_peer->peer_list[i];
+
+		if (!peer_entry->link_peer)
+			continue;
+
+		if (peer_entry->link_ix == link_ix) {
+			frm_buf = qdf_nbuf_clone(peer_entry->assoc_rsp_buf);
+			break;
+		}
+	}
+	mlo_peer_lock_release(ml_peer);
+
+	return frm_buf;
+}
+
+void wlan_mlo_peer_free_all_link_assoc_resp_buf(
+			struct wlan_objmgr_peer *link_peer)
+{
+	struct wlan_mlo_link_peer_entry *peer_entry;
+	struct wlan_mlo_peer_context *ml_peer;
+	uint8_t i;
+
+	ml_peer = link_peer->mlo_peer_ctx;
+	if (!ml_peer)
+		return;
+
+	mlo_peer_lock_acquire(ml_peer);
+
+	for (i = 0; i < MAX_MLO_LINK_PEERS; i++) {
+		peer_entry = &ml_peer->peer_list[i];
+
+		if (peer_entry->assoc_rsp_buf) {
+			qdf_nbuf_free(peer_entry->assoc_rsp_buf);
+			peer_entry->assoc_rsp_buf = NULL;
+		}
+	}
+	mlo_peer_lock_release(ml_peer);
+}
+
 static QDF_STATUS mlo_peer_detach_link_peer(
 		struct wlan_mlo_peer_context *ml_peer,
 		struct wlan_objmgr_peer *link_peer)
@@ -416,6 +607,11 @@ static QDF_STATUS mlo_peer_detach_link_peer(
 		if (peer_entry->link_peer != link_peer)
 			continue;
 
+		if (peer_entry->assoc_rsp_buf) {
+			qdf_nbuf_free(peer_entry->assoc_rsp_buf);
+			peer_entry->assoc_rsp_buf = NULL;
+		}
+
 		wlan_objmgr_peer_release_ref(link_peer, WLAN_MLO_MGR_ID);
 		peer_entry->link_peer = NULL;
 		ml_peer->link_peer_cnt--;
@@ -556,16 +752,22 @@ QDF_STATUS wlan_mlo_peer_create(struct wlan_objmgr_vdev *vdev,
 		qdf_copy_macaddr((struct qdf_mac_addr *)&ml_peer->peer_mld_addr,
 				 (struct qdf_mac_addr *)&link_peer->mldaddr[0]);
 		/* Allocate AID */
-		if (aid == (uint16_t)-1)
-			mlo_peer_allocate_aid(ml_dev, ml_peer);
-		else
+		if (aid == (uint16_t)-1) {
+			status = mlo_peer_allocate_aid(ml_dev, ml_peer);
+			if (status != QDF_STATUS_SUCCESS) {
+				mlo_peer_free(ml_peer);
+				mlo_dev_release_link_vdevs(link_vdevs);
+				return status;
+			}
+		} else {
 			ml_peer->assoc_id = aid;
+		}
 	}
 
 	/* Populate Link peer pointer, peer MAC address,
 	 * MLD address. HW link ID, update ref count
 	 */
-	mlo_peer_attach_link_peer(ml_peer, link_peer);
+	mlo_peer_attach_link_peer(ml_peer, link_peer, NULL);
 
 	/* Allocate Primary UMAC */
 	mlo_peer_allocate_primary_umac(ml_dev, ml_peer, link_vdevs);
@@ -606,7 +808,8 @@ QDF_STATUS wlan_mlo_peer_create(struct wlan_objmgr_vdev *vdev,
 }
 
 QDF_STATUS wlan_mlo_link_peer_attach(struct wlan_mlo_peer_context *ml_peer,
-				     struct wlan_objmgr_peer *peer)
+				     struct wlan_objmgr_peer *peer,
+				     qdf_nbuf_t frm_buf)
 {
 	QDF_STATUS status;
 	struct wlan_objmgr_peer *assoc_peer;
@@ -614,7 +817,7 @@ QDF_STATUS wlan_mlo_link_peer_attach(struct wlan_mlo_peer_context *ml_peer,
 	/* Populate Link peer pointer, peer MAC address,
 	 * MLD address. HW link ID, update ref count
 	 */
-	status = mlo_peer_attach_link_peer(ml_peer, peer);
+	status = mlo_peer_attach_link_peer(ml_peer, peer, frm_buf);
 	if (QDF_IS_STATUS_ERROR(status))
 		return status;
 
@@ -663,3 +866,53 @@ qdf_nbuf_t mlo_peer_get_link_peer_assoc_req_buf(
 
 	return assocbuf;
 }
+
+void wlan_mlo_peer_get_links_info(struct wlan_objmgr_peer *peer,
+				  struct mlo_tgt_partner_info *ml_links)
+{
+	struct wlan_mlo_peer_context *ml_peer;
+	struct wlan_mlo_link_peer_entry *peer_entry;
+	struct wlan_objmgr_peer *link_peer;
+	struct wlan_objmgr_vdev *link_vdev;
+	uint8_t i, ix;
+
+	ml_peer = peer->mlo_peer_ctx;
+	ml_links->num_partner_links = 0;
+
+	if (!ml_peer)
+		return;
+
+	mlo_peer_lock_acquire(ml_peer);
+
+	if ((ml_peer->mlpeer_state != ML_PEER_CREATED) &&
+	    (ml_peer->mlpeer_state != ML_PEER_ASSOC_DONE)) {
+		mlo_peer_lock_release(ml_peer);
+		return;
+	}
+
+	for (i = 0; i < MAX_MLO_LINK_PEERS; i++) {
+		peer_entry = &ml_peer->peer_list[i];
+		link_peer = peer_entry->link_peer;
+
+		if (!link_peer)
+			continue;
+
+		if (link_peer == peer)
+			continue;
+
+		link_vdev = wlan_peer_get_vdev(link_peer);
+		if (!link_vdev)
+			continue;
+
+		if (ml_links->num_partner_links >= WLAN_UMAC_MLO_MAX_VDEVS)
+			break;
+
+		ix = ml_links->num_partner_links;
+		ml_links->link_info[ix].vdev_id = wlan_vdev_get_id(link_vdev);
+		ml_links->link_info[ix].hw_mld_link_id = peer_entry->hw_link_id;
+		ml_links->num_partner_links++;
+	}
+	mlo_peer_lock_release(ml_peer);
+}
+
+qdf_export_symbol(wlan_mlo_peer_get_links_info);

umac/mlo_mgr/src/wlan_mlo_mgr_peer_list.c

@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2021, The Linux Foundation. All rights reserved.
+ * Copyright (c) 2021 Qualcomm Innovation Center, Inc. All rights reserved.
  *
  * Permission to use, copy, modify, and/or distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -155,10 +156,8 @@ wlan_find_mlpeer_link_mac_addr(struct wlan_mlo_dev_context *ml_dev,
 	uint8_t i;
 
 	ml_peer = (struct wlan_mlo_peer_context *)iter_ml_peer;
-	for (i = 0; i < MAX_MLO_PEER; i++) {
+	for (i = 0; i < MAX_MLO_LINK_PEERS; i++) {
 		link_peer = &ml_peer->peer_list[i];
-		if (!link_peer)
-			continue;
 
 		if (qdf_is_macaddr_equal(&link_mac_arg->mac_addr,
 					 &link_peer->link_addr)) {

umac/mlo_mgr/src/wlan_mlo_mgr_primary_umac.c

@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2021, The Linux Foundation. All rights reserved.
+ * Copyright (c) 2021 Qualcomm Innovation Center, Inc. All rights reserved.
  *
  * Permission to use, copy, modify, and/or distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -19,16 +20,41 @@
 #include "wlan_cmn.h"
 #include "wlan_mlo_mgr_peer.h"
 #include <wlan_mlo_mgr_ap.h>
+#include <wlan_utility.h>
+#include <wlan_reg_services_api.h>
+
+static void
+mld_get_best_primary_umac_w_rssi(struct wlan_mlo_peer_context *ml_peer,
+				 struct wlan_objmgr_vdev *link_vdevs[])
+{
+	struct wlan_objmgr_peer *assoc_peer;
+
+	assoc_peer = wlan_mlo_peer_get_assoc_peer(ml_peer);
+	ml_peer->primary_umac_psoc_id = wlan_peer_get_psoc_id(assoc_peer);
+}
 
 void mlo_peer_assign_primary_umac(
 		struct wlan_mlo_peer_context *ml_peer,
 		struct wlan_mlo_link_peer_entry *peer_entry)
 {
-	if (wlan_peer_get_psoc_id(peer_entry->link_peer) ==
-			ml_peer->primary_umac_psoc_id)
-		peer_entry->is_primary = true;
-	else
-		peer_entry->is_primary = false;
+	/* If MLD is within single SOC, then assoc link becomes
+	 * primary umac
+	 */
+	if (ml_peer->primary_umac_psoc_id == ML_PRIMARY_UMAC_ID_INVAL) {
+		if (mlo_peer_is_assoc_peer(ml_peer, peer_entry->link_peer)) {
+			peer_entry->is_primary = true;
+			ml_peer->primary_umac_psoc_id =
+				wlan_peer_get_psoc_id(peer_entry->link_peer);
+		} else {
+			peer_entry->is_primary = false;
+		}
+	} else {
+		if (wlan_peer_get_psoc_id(peer_entry->link_peer) ==
+				ml_peer->primary_umac_psoc_id)
+			peer_entry->is_primary = true;
+		else
+			peer_entry->is_primary = false;
+	}
 }
 
 QDF_STATUS mlo_peer_allocate_primary_umac(
@@ -36,20 +62,60 @@ QDF_STATUS mlo_peer_allocate_primary_umac(
 		struct wlan_mlo_peer_context *ml_peer,
 		struct wlan_objmgr_vdev *link_vdevs[])
 {
-	struct wlan_objmgr_vdev *vdev;
-	uint16_t link_load[WLAN_UMAC_MLO_MAX_VDEVS];
-	uint8_t i;
-
-	for (i = 0; i < WLAN_UMAC_MLO_MAX_VDEVS; i++) {
-		vdev = ml_dev->wlan_vdev_list[i];
-		if (!vdev) {
-			link_load[i] = 0;
-			continue;
+	struct wlan_mlo_link_peer_entry *peer_entry;
+	struct wlan_objmgr_peer *assoc_peer = NULL;
+	int32_t rssi;
+	struct mlo_mgr_context *mlo_ctx = wlan_objmgr_get_mlo_ctx();
+	uint8_t first_link_id = 0;
+	bool primary_umac_set = false;
+	uint8_t i, psoc_id;
+
+	peer_entry = &ml_peer->peer_list[0];
+	assoc_peer = peer_entry->link_peer;
+	if (!assoc_peer)
+		return QDF_STATUS_E_FAILURE;
+
+	/* For Station mode, assign assoc peer as primary umac */
+	if (wlan_peer_get_peer_type(assoc_peer) == WLAN_PEER_AP) {
+		mlo_peer_assign_primary_umac(ml_peer, peer_entry);
+		return QDF_STATUS_SUCCESS;
+	}
+
+	/* If MLD is single chip MLO then assoc link becomes primary UMAC */
+	/*
+	 * if (ml_dev->single_chip_mlo) {
+	 *	mlo_peer_assign_primary_umac(ml_peer, peer_entry);
+	 *	return QDF_STATUS_SUCCESS;
+	 * }
+	 */
+
+	if (mlo_ctx->mlo_is_force_primary_umac) {
+		for (i = 0; i < WLAN_UMAC_MLO_MAX_VDEVS; i++) {
+			if (!link_vdevs[i])
+				continue;
+
+			psoc_id = wlan_vdev_get_psoc_id(link_vdevs[i]);
+			if (!first_link_id)
+				first_link_id = psoc_id;
+
+			if (psoc_id == mlo_ctx->mlo_forced_primary_umac_id) {
+				ml_peer->primary_umac_psoc_id = psoc_id;
+				primary_umac_set = true;
+				break;
+			}
 		}
 
-		link_load[i] = 1/* TODO mlo_get_umac_load(vdev)*/;
+		if (!primary_umac_set)
+			ml_peer->primary_umac_psoc_id = first_link_id;
+
+		return QDF_STATUS_SUCCESS;
 	}
-	ml_peer->primary_umac_psoc_id = 0;
+
+	rssi = wlan_peer_get_rssi(assoc_peer);
+
+	mld_get_best_primary_umac_w_rssi(ml_peer, link_vdevs);
+
+	mlo_peer_assign_primary_umac(ml_peer, peer_entry);
 
 	return QDF_STATUS_SUCCESS;
 }