diff --git a/xen/arch/x86/psr.c b/xen/arch/x86/psr.c
index a85ea99..7bc212f 100644
--- a/xen/arch/x86/psr.c
+++ b/xen/arch/x86/psr.c
@@ -125,6 +125,8 @@ struct feat_node {
     uint32_t cos_reg_val[MAX_COS_REG_CNT];
 };
 
+#define PSR_DOM_IDS_NUM ((DOMID_IDLE + 1) / sizeof(uint32_t))
+
 /*
  * PSR features are managed per socket. Below structure defines the members
  * used to manage these features.
@@ -134,9 +136,11 @@ struct feat_node {
  *             COS ID. Every entry of cos_ref corresponds to one COS ID.
  */
 struct psr_socket_info {
-    struct feat_node *features[PSR_SOCKET_MAX_FEAT];
     spinlock_t ref_lock;
+    spinlock_t dom_ids_lock;
+    struct feat_node *features[PSR_SOCKET_MAX_FEAT];
     unsigned int cos_ref[MAX_COS_REG_CNT];
+    uint32_t dom_ids[PSR_DOM_IDS_NUM];
 };
 
 struct psr_assoc {
@@ -194,26 +198,11 @@ static void free_socket_resources(unsigned int socket)
 {
     unsigned int i;
     struct psr_socket_info *info = socket_info + socket;
-    struct domain *d;
+    unsigned long flag;
 
     if ( !info )
         return;
 
-    /* Restore domain cos id to 0 when socket is offline. */
-    for_each_domain ( d )
-    {
-        unsigned int cos = d->arch.psr_cos_ids[socket];
-        if ( cos == 0 )
-            continue;
-
-        spin_lock(&info->ref_lock);
-        ASSERT(!cos || info->cos_ref[cos]);
-        info->cos_ref[cos]--;
-        spin_unlock(&info->ref_lock);
-
-        d->arch.psr_cos_ids[socket] = 0;
-    }
-
     /*
      * Free resources of features. The global feature object, e.g. feat_l3_cat,
      * may not be freed here if it is not added into array. It is simply being
@@ -221,12 +210,17 @@ static void free_socket_resources(unsigned int socket)
      */
     for ( i = 0; i < PSR_SOCKET_MAX_FEAT; i++ )
     {
-        if ( !info->features[i] )
-            continue;
-
         xfree(info->features[i]);
         info->features[i] = NULL;
     }
+
+    spin_lock(&info->ref_lock);
+    memset(info->cos_ref, 0, MAX_COS_REG_CNT * sizeof(unsigned int));
+    spin_unlock(&info->ref_lock);
+
+    spin_lock_irqsave(&info->dom_ids_lock, flag);
+    memset(info->dom_ids, 0, PSR_DOM_IDS_NUM * sizeof(uint32_t));
+    spin_unlock_irqrestore(&info->dom_ids_lock, flag);
 }
 
 static bool feat_init_done(const struct psr_socket_info *info)
@@ -682,9 +676,37 @@ void psr_ctxt_switch_to(struct domain *d)
         psr_assoc_rmid(&reg, d->arch.psr_rmid);
 
     if ( psra->cos_mask )
-        psr_assoc_cos(&reg, d->arch.psr_cos_ids ?
-                      d->arch.psr_cos_ids[cpu_to_socket(smp_processor_id())] :
-                      0, psra->cos_mask);
+    {
+        unsigned int socket = cpu_to_socket(smp_processor_id());
+        struct psr_socket_info *info = socket_info + socket;
+
+        if ( test_bit(d->domain_id, info->dom_ids) )
+            goto set_assoc;
+
+        spin_lock(&info->dom_ids_lock);
+
+        int old_bit = test_and_set_bit(d->domain_id, info->dom_ids);
+
+        /*
+         * If old_bit is 0, that means this is the first time the domain is
+         * switched to this socket or domain's COS ID has not been set since
+         * the socket is online. So, the domain's COS ID on this socket should
+         * be default value, 0. If not, that means this socket has been offline
+         * and the domain's COS ID has been set when the socket was online. So,
+         * this COS ID is invalid and we have to restore it to 0.
+         */
+        if ( d->arch.psr_cos_ids &&
+             old_bit == 0 &&
+             d->arch.psr_cos_ids[socket] != 0 )
+            d->arch.psr_cos_ids[socket] = 0;
+
+        spin_unlock(&info->dom_ids_lock);
+
+ set_assoc:
+        psr_assoc_cos(&reg,
+                      d->arch.psr_cos_ids ? d->arch.psr_cos_ids[socket] : 0,
+                      psra->cos_mask);
+    }
 
     if ( reg != psra->val )
     {
@@ -1146,40 +1168,6 @@ static int write_psr_msr(unsigned int socket, unsigned int cos,
     return 0;
 }
 
-static void restore_default_val(unsigned int socket, unsigned int cos,
-                                enum psr_feat_type feat_type)
-{
-    unsigned int i, j;
-    uint32_t default_val;
-    const struct psr_socket_info *info = get_socket_info(socket);
-
-    for ( i = 0; i < PSR_SOCKET_MAX_FEAT; i++ )
-    {
-        const struct feat_node *feat = info->features[i];
-        /*
-         * There are four judgements:
-         * 1. Input 'feat_type' is valid so we have to get feature according to
-         *    it. If current feature type (i) does not match 'feat_type', we
-         *    need skip it, so continue to check next feature.
-         * 2. Input 'feat_type' is 'PSR_SOCKET_MAX_FEAT' which means we should
-         *    handle all features in this case. So, go to next loop.
-         * 3. Do not need restore the COS value back to default if cos_num is 1,
-         *    e.g. L3 CAT. Because next value setting will overwrite it.
-         * 4. 'feat' we got is NULL, continue.
-         */
-        if ( ( feat_type != PSR_SOCKET_MAX_FEAT && feat_type != i ) ||
-             !feat || feat->props->cos_num == 1 )
-            continue;
-
-        for ( j = 0; j < feat->props->cos_num; j++ )
-        {
-            feat->props->get_val(feat, 0, feat->props->type[j], &default_val);
-            write_psr_msr(socket, cos, default_val,
-                          feat->props->type[j], i);
-        }
-    }
-}
-
 /* The whole set process is protected by domctl_lock. */
 int psr_set_val(struct domain *d, unsigned int socket,
                 uint32_t val, enum cbm_type type)
@@ -1191,6 +1179,7 @@ int psr_set_val(struct domain *d, unsigned int socket,
     struct psr_socket_info *info = get_socket_info(socket);
     unsigned int array_len;
     enum psr_feat_type feat_type;
+    unsigned long flag;
 
     if ( IS_ERR(info) )
         return PTR_ERR(info);
@@ -1286,22 +1275,6 @@ int psr_set_val(struct domain *d, unsigned int socket,
     ASSERT(!cos || ref[cos]);
     ASSERT(!old_cos || ref[old_cos]);
     ref[old_cos]--;
-
-    /*
-     * Step 6:
-     * For features,  e.g. CDP, which cos_num is more than 1, we have to
-     * restore the old_cos value back to default when ref[old_cos] is 0.
-     * Otherwise, user will see wrong values when this COS ID is reused. E.g.
-     * user wants to set DATA to 0x3ff for a new domain. He hopes to see the
-     * DATA is set to 0x3ff and CODE should be the default value, 0x7ff. But
-     * if the COS ID picked for this action is the one that has been used by
-     * other domain and the CODE has been set to 0x1ff. Then, user will see
-     * DATA: 0x3ff, CODE: 0x1ff. So, we have to restore COS values for features
-     * using multiple COSs.
-     */
-    if ( old_cos && !ref[old_cos] )
-        restore_default_val(socket, old_cos, feat_type);
-
     spin_unlock(&info->ref_lock);
 
     /*
@@ -1310,7 +1283,10 @@ int psr_set_val(struct domain *d, unsigned int socket,
      * which COS the domain is using on the socket. One domain can only use
      * one COS ID at same time on each socket.
      */
+    spin_lock_irqsave(&info->dom_ids_lock, flag);
     d->arch.psr_cos_ids[socket] = cos;
+    test_and_set_bit(d->domain_id, info->dom_ids);
+    spin_unlock_irqrestore(&info->dom_ids_lock, flag);
 
     xfree(val_array);
     return ret;
@@ -1336,6 +1312,7 @@ static void psr_free_cos(struct domain *d)
     for ( socket = 0; socket < nr_sockets; socket++ )
     {
         struct psr_socket_info *info;
+        unsigned long flag;
 
         /* cos 0 is default one which does not need be handled. */
         cos = d->arch.psr_cos_ids[socket];
@@ -1346,14 +1323,11 @@ static void psr_free_cos(struct domain *d)
         spin_lock(&info->ref_lock);
         ASSERT(!cos || info->cos_ref[cos]);
         info->cos_ref[cos]--;
-        /*
-         * The 'cos_ref[cos]' of 'd' is 0 now so we need restore corresponding
-         * COS registers to default value. Because this case happens when a
-         * domain is destroied, we need restore all features.
-         */
-        if ( !info->cos_ref[cos] )
-            restore_default_val(socket, cos, PSR_SOCKET_MAX_FEAT);
         spin_unlock(&info->ref_lock);
+
+        spin_lock_irqsave(&info->dom_ids_lock, flag);
+        clear_bit(d->domain_id, info->dom_ids);
+        spin_unlock_irqrestore(&info->dom_ids_lock, flag);
     }
 
     xfree(d->arch.psr_cos_ids);
@@ -1453,6 +1427,7 @@ static void psr_cpu_init(void)
         goto assoc_init;
 
     spin_lock_init(&info->ref_lock);
+    spin_lock_init(&info->dom_ids_lock);
 
     cpuid_count_leaf(PSR_CPUID_LEVEL_CAT, 0, &regs);
     if ( regs.b & PSR_RESOURCE_TYPE_L3 )
