extern gfp_t gfp_allowed_mask;
 
-extern void set_gfp_allowed_mask(gfp_t mask);
-extern gfp_t clear_gfp_allowed_mask(gfp_t mask);
+extern void pm_restrict_gfp_mask(void);
+extern void pm_restore_gfp_mask(void);
 
 #endif /* __LINUX_GFP_H */
 
 int hibernation_snapshot(int platform_mode)
 {
        int error;
-       gfp_t saved_mask;
 
        error = platform_begin(platform_mode);
        if (error)
                goto Close;
 
        suspend_console();
-       saved_mask = clear_gfp_allowed_mask(GFP_IOFS);
+       pm_restrict_gfp_mask();
        error = dpm_suspend_start(PMSG_FREEZE);
        if (error)
                goto Recover_platform;
                goto Recover_platform;
 
        error = create_image(platform_mode);
-       /* Control returns here after successful restore */
+       /*
+        * Control returns here (1) after the image has been created or the
+        * image creation has failed and (2) after a successful restore.
+        */
 
  Resume_devices:
        /* We may need to release the preallocated image pages here. */
 
        dpm_resume_end(in_suspend ?
                (error ? PMSG_RECOVER : PMSG_THAW) : PMSG_RESTORE);
-       set_gfp_allowed_mask(saved_mask);
+
+       if (error || !in_suspend)
+               pm_restore_gfp_mask();
+
        resume_console();
  Close:
        platform_end(platform_mode);
 int hibernation_restore(int platform_mode)
 {
        int error;
-       gfp_t saved_mask;
 
        pm_prepare_console();
        suspend_console();
-       saved_mask = clear_gfp_allowed_mask(GFP_IOFS);
+       pm_restrict_gfp_mask();
        error = dpm_suspend_start(PMSG_QUIESCE);
        if (!error) {
                error = resume_target_kernel(platform_mode);
                dpm_resume_end(PMSG_RECOVER);
        }
-       set_gfp_allowed_mask(saved_mask);
+       pm_restore_gfp_mask();
        resume_console();
        pm_restore_console();
        return error;
 int hibernation_platform_enter(void)
 {
        int error;
-       gfp_t saved_mask;
 
        if (!hibernation_ops)
                return -ENOSYS;
 
        entering_platform_hibernation = true;
        suspend_console();
-       saved_mask = clear_gfp_allowed_mask(GFP_IOFS);
        error = dpm_suspend_start(PMSG_HIBERNATE);
        if (error) {
                if (hibernation_ops->recover)
  Resume_devices:
        entering_platform_hibernation = false;
        dpm_resume_end(PMSG_RESTORE);
-       set_gfp_allowed_mask(saved_mask);
        resume_console();
 
  Close:
                swsusp_free();
                if (!error)
                        power_down();
+               pm_restore_gfp_mask();
        } else {
                pr_debug("PM: Image restored successfully.\n");
        }
 
 int suspend_devices_and_enter(suspend_state_t state)
 {
        int error;
-       gfp_t saved_mask;
 
        if (!suspend_ops)
                return -ENOSYS;
                        goto Close;
        }
        suspend_console();
-       saved_mask = clear_gfp_allowed_mask(GFP_IOFS);
+       pm_restrict_gfp_mask();
        suspend_test_start();
        error = dpm_suspend_start(PMSG_SUSPEND);
        if (error) {
        suspend_test_start();
        dpm_resume_end(PMSG_RESUME);
        suspend_test_finish("resume devices");
-       set_gfp_allowed_mask(saved_mask);
+       pm_restore_gfp_mask();
        resume_console();
  Close:
        if (suspend_ops->end)
 
        case SNAPSHOT_UNFREEZE:
                if (!data->frozen || data->ready)
                        break;
+               pm_restore_gfp_mask();
                thaw_processes();
                usermodehelper_enable();
                data->frozen = 0;
                        error = -EPERM;
                        break;
                }
+               pm_restore_gfp_mask();
                error = hibernation_snapshot(data->platform_support);
                if (!error)
                        error = put_user(in_suspend, (int __user *)arg);
 
  * only be modified with pm_mutex held, unless the suspend/hibernate code is
  * guaranteed not to run in parallel with that modification).
  */
-void set_gfp_allowed_mask(gfp_t mask)
+
+static gfp_t saved_gfp_mask;
+
+void pm_restore_gfp_mask(void)
 {
        WARN_ON(!mutex_is_locked(&pm_mutex));
-       gfp_allowed_mask = mask;
+       if (saved_gfp_mask) {
+               gfp_allowed_mask = saved_gfp_mask;
+               saved_gfp_mask = 0;
+       }
 }
 
-gfp_t clear_gfp_allowed_mask(gfp_t mask)
+void pm_restrict_gfp_mask(void)
 {
-       gfp_t ret = gfp_allowed_mask;
-
        WARN_ON(!mutex_is_locked(&pm_mutex));
-       gfp_allowed_mask &= ~mask;
-       return ret;
+       WARN_ON(saved_gfp_mask);
+       saved_gfp_mask = gfp_allowed_mask;
+       gfp_allowed_mask &= ~GFP_IOFS;
 }
 #endif /* CONFIG_PM_SLEEP */