2 files changed, 9 insertions, 5 deletions

diff --git a/drivers/xen/balloon.c b/drivers/xen/balloon.c
index 5dfd8f8ff07f..cd8b4704ef4e 100644
--- a/drivers/xen/balloon.c
+++ b/drivers/xen/balloon.c
@@ -501,20 +501,24 @@ EXPORT_SYMBOL_GPL(balloon_set_new_target);
  * alloc_xenballooned_pages - get pages that have been ballooned out
  * @nr_pages: Number of pages to get
  * @pages: pages returned
+ * @highmem: highmem or lowmem pages
  * @return 0 on success, error otherwise
  */
-int alloc_xenballooned_pages(int nr_pages, struct page** pages)
+int alloc_xenballooned_pages(int nr_pages, struct page **pages, bool highmem)
 {
         int pgno = 0;
         struct page* page;
         mutex_lock(&balloon_mutex);
         while (pgno < nr_pages) {
-                page = balloon_retrieve(true);
-                if (page) {
+                page = balloon_retrieve(highmem);
+                if (page && PageHighMem(page) == highmem) {
                         pages[pgno++] = page;
                 } else {
                         enum bp_state st;
-                        st = decrease_reservation(nr_pages - pgno, GFP_HIGHUSER);
+                        if (page)
+                                balloon_append(page);
+                        st = decrease_reservation(nr_pages - pgno,
+                                        highmem ? GFP_HIGHUSER : GFP_USER);
                         if (st != BP_DONE)
                                 goto out_undo;
                 }
diff --git a/drivers/xen/gntdev.c b/drivers/xen/gntdev.c
index f914b26cf0c2..7b9b1d1b75a5 100644
--- a/drivers/xen/gntdev.c
+++ b/drivers/xen/gntdev.c
@@ -123,7 +123,7 @@ static struct grant_map *gntdev_alloc_map(struct gntdev_priv *priv, int count)
             NULL == add->pages)
                 goto err;
 
-        if (alloc_xenballooned_pages(count, add->pages))
+        if (alloc_xenballooned_pages(count, add->pages, false /* lowmem */))
                 goto err;
 
         for (i = 0; i < count; i++) {