3 files changed, 30 insertions, 31 deletions

diff --git a/drivers/block/as-iosched.c b/drivers/block/as-iosched.c
index d63c92dfcf96..b19289348fb0 100644
--- a/drivers/block/as-iosched.c
+++ b/drivers/block/as-iosched.c
@@ -219,13 +219,17 @@ static struct as_io_context *alloc_as_io_context(void)
  */
 static struct io_context *as_get_io_context(void)
 {
-	struct io_context *ioc = get_io_context();
-	if (ioc && !ioc->aic)
+	struct io_context *ioc = get_io_context(GFP_ATOMIC);
+	if (ioc && !ioc->aic) {
 		ioc->aic = alloc_as_io_context();
+		if (!ioc->aic) {
+			put_io_context(ioc);
+			ioc = NULL;
+		}
+	}
 	return ioc;
 }
 
-
 /*
  * the back merge hash support functions
  */
@@ -971,32 +975,26 @@ static void as_completed_request(request_queue_t *q, struct request *rq)
 static void as_remove_queued_request(request_queue_t *q, struct request *rq)
 {
 	struct as_rq *arq = RQ_DATA(rq);
+	const int data_dir = arq->is_sync;
+	struct as_data *ad = q->elevator.elevator_data;
 
-	if (!arq)
-		BUG();
-	else {
-		const int data_dir = arq->is_sync;
-		struct as_data *ad = q->elevator.elevator_data;
-
-		WARN_ON(arq->state != AS_RQ_QUEUED);
-
-		if (arq->io_context && arq->io_context->aic) {
-			BUG_ON(!atomic_read(&arq->io_context->aic->nr_queued));
-			atomic_dec(&arq->io_context->aic->nr_queued);
-		}
-
-		/*
-		 * Update the "next_arq" cache if we are about to remove its
-		 * entry
-		 */
-		if (ad->next_arq[data_dir] == arq)
-			ad->next_arq[data_dir] = as_find_next_arq(ad, arq);
+	WARN_ON(arq->state != AS_RQ_QUEUED);
 
-		list_del_init(&arq->fifo);
-		as_remove_merge_hints(q, arq);
-		as_del_arq_rb(ad, arq);
+	if (arq->io_context && arq->io_context->aic) {
+		BUG_ON(!atomic_read(&arq->io_context->aic->nr_queued));
+		atomic_dec(&arq->io_context->aic->nr_queued);
 	}
 
+	/*
+	 * Update the "next_arq" cache if we are about to remove its
+	 * entry
+	 */
+	if (ad->next_arq[data_dir] == arq)
+		ad->next_arq[data_dir] = as_find_next_arq(ad, arq);
+
+	list_del_init(&arq->fifo);
+	as_remove_merge_hints(q, arq);
+	as_del_arq_rb(ad, arq);
 }
 
 /*
@@ -1292,7 +1290,7 @@ static void as_add_request(struct as_data *ad, struct as_rq *arq)
 
 	arq->io_context = as_get_io_context();
 
-	if (arq->io_context && arq->io_context->aic) {
+	if (arq->io_context) {
 		atomic_inc(&arq->io_context->aic->nr_queued);
 		as_update_iohist(arq->io_context->aic, arq->request);
 	}
diff --git a/drivers/block/ll_rw_blk.c b/drivers/block/ll_rw_blk.c
index 633266ee8c87..13cc6073bb47 100644
--- a/drivers/block/ll_rw_blk.c
+++ b/drivers/block/ll_rw_blk.c
@@ -1360,7 +1360,7 @@ static struct request *get_request(request_queue_t *q, int rw, int gfp_mask)
 {
 	struct request *rq = NULL;
 	struct request_list *rl = &q->rq;
-	struct io_context *ioc = get_io_context();
+	struct io_context *ioc = get_io_context(gfp_mask);
 
 	spin_lock_irq(q->queue_lock);
 	if (rl->count[rw]+1 >= q->nr_requests) {
@@ -1439,7 +1439,7 @@ static struct request *get_request_wait(request_queue_t *q, int rw)
 			struct io_context *ioc;
 
 			io_schedule();
-			ioc = get_io_context();
+			ioc = get_io_context(GFP_NOIO);
 			ioc_set_batching(ioc);
 			put_io_context(ioc);
 		}
@@ -2462,7 +2462,7 @@ void exit_io_context(void)
  * But weird things happen, so we disable local interrupts to ensure exclusive
  * access to *current.
  */
-struct io_context *get_io_context(void)
+struct io_context *get_io_context(int gfp_flags)
 {
 	struct task_struct *tsk = current;
 	unsigned long flags;
@@ -2482,8 +2482,9 @@ struct io_context *get_io_context(void)
 			tsk->io_context = ret;
 		}
 	}
+	if (ret)
+		atomic_inc(&ret->refcount);
 	local_irq_restore(flags);
-	atomic_inc(&ret->refcount);
 	return ret;
 }
 
diff --git a/include/linux/blkdev.h b/include/linux/blkdev.h
index 69178ca80d7d..2e7f92aa1dc2 100644
--- a/include/linux/blkdev.h
+++ b/include/linux/blkdev.h
@@ -70,7 +70,7 @@ struct io_context {
 
 void put_io_context(struct io_context *ioc);
 void exit_io_context(void);
-struct io_context *get_io_context(void);
+struct io_context *get_io_context(int gfp_flags);
 void copy_io_context(struct io_context **pdst, struct io_context **psrc);
 void swap_io_context(struct io_context **ioc1, struct io_context **ioc2);