@@ -185,13 +185,6 @@ void umcg_clear_child(struct task_struct
 	umcg_clear_task(tsk);
 }
 
-/* Called both by normally (unregister) and abnormally exiting workers. */
-void umcg_worker_exit(void)
-{
-	umcg_unpin_pages();
-	umcg_clear_task(current);
-}
-
 /*
  * Do a state transition: @from -> @to.
  *

@@ -748,32 +741,43 @@ SYSCALL_DEFINE2(umcg_wait, u32, flags, u
  * sys_umcg_ctl: (un)register the current task as a UMCG task.
  * @flags:       ORed values from enum umcg_ctl_flag; see below;
  * @self:        a pointer to struct umcg_task that describes this
- *               task and governs the behavior of sys_umcg_wait if
- *               registering; must be NULL if unregistering.
+ *               task and governs the behavior of sys_umcg_wait.
  * @which_clock: clockid to use for timestamps and timeouts
  *
  * @flags & UMCG_CTL_REGISTER: register a UMCG task:
  *
- *         UMCG workers:
- *              - @flags & UMCG_CTL_WORKER
- *              - self->state must be UMCG_TASK_BLOCKED
- *
- *         UMCG servers:
- *              - !(@flags & UMCG_CTL_WORKER)
- *              - self->state must be UMCG_TASK_RUNNING
- *
- *         All tasks:
- *              - self->server_tid must be a valid server
- *              - self->next_tid must be zero
- *
- *         If the conditions above are met, sys_umcg_ctl() immediately returns
- *         if the registered task is a server. If the registered task is a
- *         worker it will be added to it's server's runnable_workers_ptr list
- *         and the server will be woken.
- *
- * @flags == UMCG_CTL_UNREGISTER: unregister a UMCG task. If the current task
- *           is a UMCG worker, the userspace is responsible for waking its
- *           server (before or after calling sys_umcg_ctl).
+ *	UMCG workers:
+ *	 - @flags & UMCG_CTL_WORKER
+ *	 - self->state must be UMCG_TASK_BLOCKED
+ *
+ *	UMCG servers:
+ *	 - !(@flags & UMCG_CTL_WORKER)
+ *	 - self->state must be UMCG_TASK_RUNNING
+ *
+ *	All tasks:
+ *	 - self->server_tid must be a valid server
+ *	 - self->next_tid must be zero
+ *
+ *	If the conditions above are met, sys_umcg_ctl() immediately returns
+ *	if the registered task is a server. If the registered task is a
+ *	worker it will be added to it's server's runnable_workers_ptr list
+ *	and the server will be woken.
+ *
+ * @flags & UMCG_CTL_UNREGISTER: unregister a UMCG task.
+ *
+ *	UMCG workers:
+ *	 - @flags & UMCG_CTL_WORKER
+ *
+ *	UMCG servers:
+ *	 - !(@flags & UMCG_CTL_WORKER)
+ *
+ *	All tasks:
+ *	 - self must match with UMCG_CTL_REGISTER
+ *	 - self->state must be UMCG_TASK_RUNNING
+ *	 - self->server_tid must be a valid server
+ *
+ * 	If the conditions above are met, sys_umcg_ctl() will change state to
+ * 	UMCG_TASK_NONE, and for workers, wake either next or server.
  *
  * Return:
  * 0		- success

@@ -794,16 +798,31 @@ SYSCALL_DEFINE3(umcg_ctl, u32, flags, st
 		      UMCG_CTL_WORKER))
 		return -EINVAL;
 
-	if (flags == UMCG_CTL_UNREGISTER) {
-		if (self || !current->umcg_task)
+	if (flags & UMCG_CTL_UNREGISTER) {
+		int ret;
+
+		if (!self || self != current->umcg_task)
 			return -EINVAL;
 
-		if (current->flags & PF_UMCG_WORKER) {
-			umcg_worker_exit();
-			// XXX wake server
-		} else
-			umcg_clear_task(current);
+		current->flags &= ~PF_UMCG_WORKER;
 
+		ret = umcg_pin_pages();
+		if (ret) {
+			current->flags |= PF_UMCG_WORKER;
+			return ret;
+		}
+
+		ret = umcg_update_state(current, self, UMCG_TASK_RUNNING, UMCG_TASK_NONE);
+		if (ret) {
+			current->flags |= PF_UMCG_WORKER;
+			return ret;
+		}
+
+		if (current->flags & PF_UMCG_WORKER)
+			umcg_wake(current);
+
+		umcg_unpin_pages();
+		umcg_clear_task(current);
 		return 0;
 	}