Index: uspace/lib/c/generic/task.c
===================================================================
--- uspace/lib/c/generic/task.c	(revision 88205440caa3e827197cd095139acca02b989d5e)
+++ uspace/lib/c/generic/task.c	(revision 32e3cdfc8986d9f6c67d7e98cd10bf6ffa0720ef)
@@ -2,4 +2,5 @@
  * Copyright (c) 2006 Jakub Jermar
  * Copyright (c) 2008 Jiri Svoboda
+ * Copyright (c) 2014 Martin Sucha
  * All rights reserved.
  *
@@ -94,4 +95,6 @@
  *
  * @param id   If not NULL, the ID of the task is stored here on success.
+ * @param wait If not NULL, setup waiting for task's return value and store
+ *             the information necessary for waiting here on success.
  * @param path Pathname of the binary to execute.
  * @param argv Command-line arguments.
@@ -100,5 +103,6 @@
  *
  */
-int task_spawnv(task_id_t *id, const char *path, const char *const args[])
+int task_spawnv(task_id_t *id, task_wait_t *wait, const char *path,
+    const char *const args[])
 {
 	/* Send default files */
@@ -125,5 +129,5 @@
 	files[3] = NULL;
 	
-	return task_spawnvf(id, path, args, files);
+	return task_spawnvf(id, wait, path, args, files);
 }
 
@@ -135,4 +139,6 @@
  *
  * @param id    If not NULL, the ID of the task is stored here on success.
+ * @param wait  If not NULL, setup waiting for task's return value and store
+ *              the information necessary for waiting here on success.
  * @param path  Pathname of the binary to execute.
  * @param argv  Command-line arguments.
@@ -142,6 +148,6 @@
  *
  */
-int task_spawnvf(task_id_t *id, const char *path, const char *const args[],
-    int *const files[])
+int task_spawnvf(task_id_t *id, task_wait_t *wait, const char *path,
+    const char *const args[], int *const files[])
 {
 	/* Connect to a program loader. */
@@ -150,4 +156,6 @@
 		return EREFUSED;
 	
+	bool wait_initialized = false;
+	
 	/* Get task ID. */
 	task_id_t task_id;
@@ -181,4 +189,12 @@
 		goto error;
 	
+	/* Setup waiting for return value if needed */
+	if (wait) {
+		rc = task_setup_wait(task_id, wait);
+		if (rc != EOK)
+			goto error;
+		wait_initialized = true;
+	}
+	
 	/* Run it. */
 	rc = loader_run(ldr);
@@ -193,4 +209,7 @@
 	
 error:
+	if (wait_initialized)
+		task_cancel_wait(wait);
+	
 	/* Error exit */
 	loader_abort(ldr);
@@ -204,4 +223,6 @@
  *
  * @param id   If not NULL, the ID of the task is stored here on success.
+ * @param wait If not NULL, setup waiting for task's return value and store
+ *             the information necessary for waiting here on success.
  * @param path Pathname of the binary to execute.
  * @param cnt  Number of arguments.
@@ -211,5 +232,6 @@
  *
  */
-int task_spawn(task_id_t *task_id, const char *path, int cnt, va_list ap)
+int task_spawn(task_id_t *task_id, task_wait_t *wait, const char *path,
+    int cnt, va_list ap)
 {
 	/* Allocate argument list. */
@@ -227,5 +249,5 @@
 	
 	/* Spawn task. */
-	int rc = task_spawnv(task_id, path, arglist);
+	int rc = task_spawnv(task_id, wait, path, arglist);
 	
 	/* Free argument list. */
@@ -240,4 +262,6 @@
  *
  * @param id   If not NULL, the ID of the task is stored here on success.
+ * @param wait If not NULL, setup waiting for task's return value and store
+ *             the information necessary for waiting here on success.
  * @param path Pathname of the binary to execute.
  * @param ...  Command-line arguments.
@@ -246,5 +270,5 @@
  *
  */
-int task_spawnl(task_id_t *task_id, const char *path, ...)
+int task_spawnl(task_id_t *task_id, task_wait_t *wait, const char *path, ...)
 {
 	/* Count the number of arguments. */
@@ -262,5 +286,5 @@
 	
 	va_start(ap, path);
-	int rc = task_spawn(task_id, path, cnt, ap);
+	int rc = task_spawn(task_id, wait, path, cnt, ap);
 	va_end(ap);
 	
@@ -268,19 +292,90 @@
 }
 
-int task_wait(task_id_t id, task_exit_t *texit, int *retval)
+/** Setup waiting for a task.
+ * 
+ * If the task finishes after this call succeeds, it is guaranteed that
+ * task_wait(wait, &texit, &retval) will return correct return value for
+ * the task.
+ *
+ * @param id   ID of the task to setup waiting for.
+ * @param wait Information necessary for the later task_wait call is stored here.
+ *
+ * @return EOK on success, else error code.
+ */
+int task_setup_wait(task_id_t id, task_wait_t *wait)
+{
+	async_exch_t *exch = async_exchange_begin(session_ns);
+	wait->aid = async_send_2(exch, NS_TASK_WAIT, LOWER32(id), UPPER32(id),
+	    &wait->result);
+	async_exchange_end(exch);
+
+	return EOK;
+}
+
+/** Cancel waiting for a task.
+ *
+ * This can be called *instead of* task_wait if the caller is not interested
+ * in waiting for the task anymore.
+ *
+ * This function cannot be called if the task_wait was already called.
+ *
+ * @param wait task_wait_t previously initialized by task_setup_wait.
+ */
+void task_cancel_wait(task_wait_t *wait) {
+	async_forget(wait->aid);
+}
+
+/** Wait for a task to finish.
+ *
+ * This function returns correct values even if the task finished in
+ * between task_setup_wait and this task_wait call.
+ *
+ * This function cannot be called more than once with the same task_wait_t
+ * (it can be reused, but must be reinitialized with task_setup_wait first)
+ *
+ * @param wait   task_wait_t previously initialized by task_setup_wait.
+ * @param texit  Store type of task exit here.
+ * @param retval Store return value of the task here.
+ * 
+ * @return EOK on success, else error code.
+ */
+int task_wait(task_wait_t *wait, task_exit_t *texit, int *retval)
 {
 	assert(texit);
 	assert(retval);
-	
-	async_exch_t *exch = async_exchange_begin(session_ns);
-	sysarg_t te, rv;
-	int rc = (int) async_req_2_2(exch, NS_TASK_WAIT, LOWER32(id),
-	    UPPER32(id), &te, &rv);
-	async_exchange_end(exch);
-	
-	*texit = te;
-	*retval = rv;
-	
-	return rc;
+
+	sysarg_t rc;
+	async_wait_for(wait->aid, &rc);
+
+	if (rc == EOK) {
+		*texit = IPC_GET_ARG1(wait->result);
+		*retval = IPC_GET_ARG2(wait->result);
+	}
+
+	return rc;
+}
+
+/** Wait for a task to finish by its id.
+ *
+ * Note that this will fail with ENOENT if the task id is not registered in ns
+ * (e.g. if the task finished). If you are spawning a task and need to wait
+ * for its completion, use wait parameter of the task_spawn* functions instead
+ * to prevent a race where the task exits before you may have a chance to wait
+ * wait for it.
+ *
+ * @param id ID of the task to wait for.
+ * @param texit  Store type of task exit here.
+ * @param retval Store return value of the task here.
+ * 
+ * @return EOK on success, else error code.
+ */
+int task_wait_task_id(task_id_t id, task_exit_t *texit, int *retval)
+{
+	task_wait_t wait;
+	int rc = task_setup_wait(id, &wait);
+	if (rc != EOK)
+		return rc;
+	
+	return task_wait(&wait, texit, retval);
 }