2 files changed, 34 insertions, 36 deletions

diff --git a/src/core/timer.c b/src/core/timer.c
index f64c1294..75759cd2 100644
--- a/src/core/timer.c
+++ b/src/core/timer.c
@@ -15,16 +15,13 @@
 static void nni_timer_loop(void *);
 
 struct nni_timer {
-	// We use two mutexes.  One protects the list, and the other ensures
-	// that cancel is blocked if we are running a timeout calllback.
-	// The callback(s) are allowed to reschedule a timeout.  The list
-	// mutex is *always* acquired before the run mutex.
-	nni_mtx		t_list_mx;
-	nni_mtx		t_run_mx;
+	nni_mtx		t_mx;
 	nni_cv		t_cv;
 	nni_list	t_entries;
 	nni_thr		t_thr;
 	int		t_close;
+	int		t_waiting;
+	nni_timer_node *t_active;       // Must never ever be dereferenced!
 };
 
 typedef struct nni_timer   nni_timer;
@@ -42,9 +39,8 @@ nni_timer_sys_init(void)
 	NNI_LIST_INIT(&timer->t_entries, nni_timer_node, t_node);
 	timer->t_close = 0;
 
-	if (((rv = nni_mtx_init(&timer->t_list_mx)) != 0) ||
-	    ((rv = nni_mtx_init(&timer->t_run_mx)) != 0) ||
-	    ((rv = nni_cv_init(&timer->t_cv, &timer->t_list_mx)) != 0) ||
+	if (((rv = nni_mtx_init(&timer->t_mx)) != 0) ||
+	    ((rv = nni_cv_init(&timer->t_cv, &timer->t_mx)) != 0) ||
 	    ((rv = nni_thr_init(&timer->t_thr, nni_timer_loop, timer)) != 0)) {
 		nni_timer_sys_fini();
 		return (rv);
@@ -59,15 +55,14 @@ nni_timer_sys_fini(void)
 {
 	nni_timer *timer = &nni_global_timer;
 
-	nni_mtx_lock(&timer->t_list_mx);
+	nni_mtx_lock(&timer->t_mx);
 	timer->t_close = 1;
 	nni_cv_wake(&timer->t_cv);
-	nni_mtx_unlock(&timer->t_list_mx);
+	nni_mtx_unlock(&timer->t_mx);
 
 	nni_thr_fini(&timer->t_thr);
 	nni_cv_fini(&timer->t_cv);
-	nni_mtx_fini(&timer->t_list_mx);
-	nni_mtx_fini(&timer->t_run_mx);
+	nni_mtx_fini(&timer->t_mx);
 }
 
 
@@ -91,14 +86,15 @@ nni_timer_cancel(nni_timer_node *node)
 {
 	nni_timer *timer = &nni_global_timer;
 
-	nni_mtx_lock(&timer->t_list_mx);
-	nni_mtx_lock(&timer->t_run_mx);
-	if (node->t_sched) {
+	nni_mtx_lock(&timer->t_mx);
+	while (timer->t_active == node) {
+		timer->t_waiting = 1;
+		nni_cv_wait(&timer->t_cv);
+	}
+	if (nni_list_active(&timer->t_entries, node)) {
 		nni_list_remove(&timer->t_entries, node);
-		node->t_sched = 0;
 	}
-	nni_mtx_unlock(&timer->t_run_mx);
-	nni_mtx_unlock(&timer->t_list_mx);
+	nni_mtx_unlock(&timer->t_mx);
 }
 
 
@@ -111,7 +107,7 @@ nni_timer_schedule(nni_timer_node *node, nni_time when)
 
 	node->t_expire = when;
 
-	nni_mtx_lock(&timer->t_list_mx);
+	nni_mtx_lock(&timer->t_mx);
 
 	if (nni_list_active(&timer->t_entries, node)) {
 		nni_list_remove(&timer->t_entries, node);
@@ -127,11 +123,10 @@ nni_timer_schedule(nni_timer_node *node, nni_time when)
 	} else {
 		nni_list_append(&timer->t_entries, node);
 	}
-	node->t_sched = 1;
 	if (wake) {
 		nni_cv_wake(&timer->t_cv);
 	}
-	nni_mtx_unlock(&timer->t_list_mx);
+	nni_mtx_unlock(&timer->t_mx);
 }
 
 
@@ -144,35 +139,39 @@ nni_timer_loop(void *arg)
 	nni_timer_node *node;
 
 	for (;;) {
-		nni_mtx_lock(&timer->t_list_mx);
+		nni_mtx_lock(&timer->t_mx);
+		timer->t_active = NULL;
+		if (timer->t_waiting) {
+			timer->t_waiting = 1;
+			nni_cv_wake(&timer->t_cv);
+		}
 		if (timer->t_close) {
-			nni_mtx_unlock(&timer->t_list_mx);
+			nni_mtx_unlock(&timer->t_mx);
 			break;
 		}
 
 		now = nni_clock();
 		if ((node = nni_list_first(&timer->t_entries)) == NULL) {
 			nni_cv_wait(&timer->t_cv);
-			nni_mtx_unlock(&timer->t_list_mx);
+			nni_mtx_unlock(&timer->t_mx);
 			continue;
 		}
 		if (now < node->t_expire) {
 			// End of run, we have to wait for next.
 			nni_cv_until(&timer->t_cv, node->t_expire);
-			nni_mtx_unlock(&timer->t_list_mx);
+			nni_mtx_unlock(&timer->t_mx);
 			continue;
 		}
 
 		nni_list_remove(&timer->t_entries, node);
-		node->t_sched = 0;
-
-		// The lock ordering here is important.  We acquire the run
-		// lock before dropping the list lock.  One the run is done,
-		// we can drop the run lock too.  The reason for the second
-		// lock is so that the callback can reschedule itself.
-		nni_mtx_lock(&timer->t_run_mx);
-		nni_mtx_unlock(&timer->t_list_mx);
+
+		// Save the active node.  Note that the timer callback can
+		// free this memory or do something else with it, so it is
+		// important that we never dereference this pointer, but
+		// just compare the value of the pointer itself.
+		timer->t_active = node;
+		nni_mtx_unlock(&timer->t_mx);
+
 		node->t_cb(node->t_arg);
-		nni_mtx_unlock(&timer->t_run_mx);
 	}
 }
diff --git a/src/core/timer.h b/src/core/timer.h
index 89a5cda0..09883e0b 100644
--- a/src/core/timer.h
+++ b/src/core/timer.h
@@ -20,7 +20,6 @@ struct nni_timer_node {
 	nni_cb		t_cb;
 	void *		t_arg;
 	nni_list_node	t_node;
-	int		t_sched;
 };
 
 typedef struct nni_timer_node   nni_timer_node;