tmem: Fix domain lifecycle synchronisation.

Obtaining a domain reference count is neither necessary nor
sufficient. Instead we simply check whether a domain is already dying
when it first becomes a client of tmem. If it is not then we will
correctly clean up later via tmem_destroy() called from domain_kill().

Signed-off-by: Keir Fraser <keir.fraser@citrix.com>

xen/common/tmem.c

@@ -1170,17 +1170,19 @@ static client_t *client_create(cli_id_t cli_id)
     if ( client == NULL )
     {
         printk("failed... out of memory\n");
-        return NULL;
+        goto fail;
     }
     memset(client,0,sizeof(client_t));
     if ( (client->tmh = tmh_client_init(cli_id)) == NULL )
     {
         printk("failed... can't allocate host-dependent part of client\n");
-        if ( client )
-            tmh_free_infra(client);
-        return NULL;
+        goto fail;
+    }
+    if ( !tmh_set_client_from_id(client, client->tmh, cli_id) )
+    {
+        printk("failed... can't set client\n");
+        goto fail;
     }
-    tmh_set_client_from_id(client, client->tmh, cli_id);
     client->cli_id = cli_id;
 #ifdef __i386__
     client->compress = 0;
@@ -1202,6 +1204,10 @@ static client_t *client_create(cli_id_t cli_id)
     client->succ_eph_gets = 0; client->succ_pers_gets = 0;
     printk("ok\n");
     return client;
+
+ fail:
+    tmh_free_infra(client);
+    return NULL;
 }
 
 static void client_free(client_t *client)

xen/common/tmem_xen.c

@@ -339,10 +339,10 @@ EXPORT tmh_client_t *tmh_client_init(cli_id_t cli_id)
 
 EXPORT void tmh_client_destroy(tmh_client_t *tmh)
 {
+    ASSERT(tmh->domain->is_dying);
 #ifndef __i386__
     xmem_pool_destroy(tmh->persistent_pool);
 #endif
-    put_domain(tmh->domain);
     tmh->domain = NULL;
 }
 

xen/include/xen/tmem_xen.h

@@ -302,9 +302,6 @@ typedef struct page_info pfp_t;
 extern tmh_client_t *tmh_client_init(cli_id_t);
 extern void tmh_client_destroy(tmh_client_t *);
 
-/* we don't need to take a reference to the domain here as we hold
- * one for the entire life of the client, so use rcu_lock_domain_by_id
- * variant instead of get_domain_by_id() */
 static inline struct client *tmh_client_from_cli_id(cli_id_t cli_id)
 {
     struct client *c;
@@ -333,15 +330,21 @@ static inline tmh_cli_ptr_t *tmh_get_cli_ptr_from_current(void)
     return current->domain;
 }
 
-static inline void tmh_set_client_from_id(struct client *client,
-                                          tmh_client_t *tmh, cli_id_t cli_id)
+static inline bool_t tmh_set_client_from_id(
+    struct client *client, tmh_client_t *tmh, cli_id_t cli_id)
 {
-    /* here we DO want to take/hold a reference to the domain as
-     * this routine should be called exactly once when the client is created;
-     * the matching put_domain is in tmh_client_destroy */
-    struct domain *d = get_domain_by_id(cli_id);
-    d->tmem = client;
-    tmh->domain = d;
+    struct domain *d = rcu_lock_domain_by_id(cli_id);
+    bool_t rc = 0;
+    if ( d == NULL )
+        return 0;
+    if ( !d->is_dying )
+    {
+        d->tmem = client;
+        tmh->domain = d;
+        rc = 1;
+    }
+    rcu_unlock_domain(d);
+    return rc;
 }
 
 static inline bool_t tmh_current_is_privileged(void)