#endif
 
 /* work_structs for global per-cpu drains */
+struct pcpu_drain {
+       struct zone *zone;
+       struct work_struct work;
+};
 DEFINE_MUTEX(pcpu_drain_mutex);
-DEFINE_PER_CPU(struct work_struct, pcpu_drain);
+DEFINE_PER_CPU(struct pcpu_drain, pcpu_drain);
 
 #ifdef CONFIG_GCC_PLUGIN_LATENT_ENTROPY
 volatile unsigned long latent_entropy __latent_entropy;
 
 static void drain_local_pages_wq(struct work_struct *work)
 {
+       struct pcpu_drain *drain;
+
+       drain = container_of(work, struct pcpu_drain, work);
+
        /*
         * drain_all_pages doesn't use proper cpu hotplug protection so
         * we can race with cpu offline when the WQ can move this from
         * a different one.
         */
        preempt_disable();
-       drain_local_pages(NULL);
+       drain_local_pages(drain->zone);
        preempt_enable();
 }
 
        }
 
        for_each_cpu(cpu, &cpus_with_pcps) {
-               struct work_struct *work = per_cpu_ptr(&pcpu_drain, cpu);
-               INIT_WORK(work, drain_local_pages_wq);
-               queue_work_on(cpu, mm_percpu_wq, work);
+               struct pcpu_drain *drain = per_cpu_ptr(&pcpu_drain, cpu);
+
+               drain->zone = zone;
+               INIT_WORK(&drain->work, drain_local_pages_wq);
+               queue_work_on(cpu, mm_percpu_wq, &drain->work);
        }
        for_each_cpu(cpu, &cpus_with_pcps)
-               flush_work(per_cpu_ptr(&pcpu_drain, cpu));
+               flush_work(&per_cpu_ptr(&pcpu_drain, cpu)->work);
 
        mutex_unlock(&pcpu_drain_mutex);
 }