3 files changed, 24 insertions, 15 deletions

diff --git a/block/as-iosched.c b/block/as-iosched.c
index 8da3cf66894c..d2ee2af44b58 100644
--- a/block/as-iosched.c
+++ b/block/as-iosched.c
@@ -195,6 +195,12 @@ static void free_as_io_context(struct as_io_context *aic)
         kfree(aic);
 }
 
+static void as_trim(struct io_context *ioc)
+{
+        kfree(ioc->aic);
+        ioc->aic = NULL;
+}
+
 /* Called when the task exits */
 static void exit_as_io_context(struct as_io_context *aic)
 {
@@ -1860,6 +1866,7 @@ static struct elevator_type iosched_as = {
                 .elevator_may_queue_fn =        as_may_queue,
                 .elevator_init_fn =             as_init_queue,
                 .elevator_exit_fn =             as_exit_queue,
+                .trim =                         as_trim,
         },
 
         .elevator_ktype = &as_ktype,
diff --git a/block/cfq-iosched.c b/block/cfq-iosched.c
index 521c56d4fdbc..7102bafc98b3 100644
--- a/block/cfq-iosched.c
+++ b/block/cfq-iosched.c
@@ -1211,6 +1211,13 @@ static void cfq_free_io_context(struct cfq_io_context *cic)
         kmem_cache_free(cfq_ioc_pool, cic);
 }
 
+static void cfq_trim(struct io_context *ioc)
+{
+        ioc->set_ioprio = NULL;
+        if (ioc->cic)
+                cfq_free_io_context(ioc->cic);
+}
+
 /*
  * Called with interrupts disabled
  */
@@ -2472,6 +2479,7 @@ static struct elevator_type iosched_cfq = {
                 .elevator_may_queue_fn =        cfq_may_queue,
                 .elevator_init_fn =             cfq_init_queue,
                 .elevator_exit_fn =             cfq_exit_queue,
+                .trim =                         cfq_trim,
         },
         .elevator_ktype =       &cfq_ktype,
         .elevator_name =        "cfq",
diff --git a/block/elevator.c b/block/elevator.c
index 24b702d649a9..0232df2b16e6 100644
--- a/block/elevator.c
+++ b/block/elevator.c
@@ -675,21 +675,15 @@ void elv_unregister(struct elevator_type *e)
         /*
          * Iterate every thread in the process to remove the io contexts.
          */
-        read_lock(&tasklist_lock);
-        do_each_thread(g, p) {
-                struct io_context *ioc = p->io_context;
-                if (ioc && ioc->cic) {
-                        ioc->cic->exit(ioc->cic);
-                        ioc->cic->dtor(ioc->cic);
-                        ioc->cic = NULL;
-                }
-                if (ioc && ioc->aic) {
-                        ioc->aic->exit(ioc->aic);
-                        ioc->aic->dtor(ioc->aic);
-                        ioc->aic = NULL;
-                }
-        } while_each_thread(g, p);
-        read_unlock(&tasklist_lock);
+        if (e->ops.trim) {
+                read_lock(&tasklist_lock);
+                do_each_thread(g, p) {
+                        task_lock(p);
+                        e->ops.trim(p->io_context);
+                        task_unlock(p);
+                } while_each_thread(g, p);
+                read_unlock(&tasklist_lock);
+        }
 
         spin_lock_irq(&elv_list_lock);
         list_del_init(&e->list);