From 348298b1a38f4a71d7d3d9de8e19d6324cff29bc Mon Sep 17 00:00:00 2001
From: Gleb Smirnoff <glebius@FreeBSD.org>
Date: Wed, 16 Oct 2013 05:02:01 +0000
Subject: [PATCH] For VIMAGE kernels store vnet in the struct task, and set
 vnet context during task processing.

Reported & tested by:	mm
---
 sys/kern/subr_taskqueue.c |  3 +++
 sys/sys/_task.h           |  3 +++
 sys/sys/taskqueue.h       | 13 +++++++++++++
 3 files changed, 19 insertions(+)

diff --git a/sys/kern/subr_taskqueue.c b/sys/kern/subr_taskqueue.c
index 9c7bf411aebc..b8d30a10c120 100644
--- a/sys/kern/subr_taskqueue.c
+++ b/sys/kern/subr_taskqueue.c
@@ -42,6 +42,7 @@ __FBSDID("$FreeBSD$");
 #include <sys/taskqueue.h>
 #include <sys/unistd.h>
 #include <machine/stdarg.h>
+#include <net/vnet.h>
 
 static MALLOC_DEFINE(M_TASKQUEUE, "taskqueue", "Task Queues");
 static void	*taskqueue_giant_ih;
@@ -330,7 +331,9 @@ taskqueue_run_locked(struct taskqueue *queue)
 		tb.tb_running = task;
 		TQ_UNLOCK(queue);
 
+		CURVNET_SET(task->ta_vnet);
 		task->ta_func(task->ta_context, pending);
+		CURVNET_RESTORE();
 
 		TQ_LOCK(queue);
 		tb.tb_running = NULL;
diff --git a/sys/sys/_task.h b/sys/sys/_task.h
index 11fd1bc0d093..8ef0250d8d50 100644
--- a/sys/sys/_task.h
+++ b/sys/sys/_task.h
@@ -49,6 +49,9 @@ struct task {
 	u_short	ta_priority;		/* (c) Priority */
 	task_fn_t *ta_func;		/* (c) task handler */
 	void	*ta_context;		/* (c) argument for handler */
+#ifdef VIMAGE
+	struct vnet *ta_vnet;		/* (c) vnet context */
+#endif
 };
 
 #endif /* !_SYS__TASK_H_ */
diff --git a/sys/sys/taskqueue.h b/sys/sys/taskqueue.h
index 7836262b233c..ae38e99d53f4 100644
--- a/sys/sys/taskqueue.h
+++ b/sys/sys/taskqueue.h
@@ -36,6 +36,9 @@
 #include <sys/queue.h>
 #include <sys/_task.h>
 #include <sys/_callout.h>
+#ifdef VIMAGE
+#include <net/vnet.h>
+#endif
 
 struct taskqueue;
 struct thread;
@@ -105,12 +108,22 @@ void	taskqueue_thread_enqueue(void *context);
 /*
  * Initialise a task structure.
  */
+#ifdef VIMAGE
+#define TASK_INIT(task, priority, func, context) do {	\
+	(task)->ta_pending = 0;				\
+	(task)->ta_priority = (priority);		\
+	(task)->ta_func = (func);			\
+	(task)->ta_context = (context);			\
+	(task)->ta_vnet = curvnet;			\
+} while (0)
+#else	/* !VIMAGE */
 #define TASK_INIT(task, priority, func, context) do {	\
 	(task)->ta_pending = 0;				\
 	(task)->ta_priority = (priority);		\
 	(task)->ta_func = (func);			\
 	(task)->ta_context = (context);			\
 } while (0)
+#endif	/* !VIMAGE */
 
 void _timeout_task_init(struct taskqueue *queue,
 	    struct timeout_task *timeout_task, int priority, task_fn_t func,