1 files changed, 55 insertions, 8 deletions

diff --git a/src/fiber.c b/src/fiber.c
index b678842..3f8bb15 100644
--- a/src/fiber.c
+++ b/src/fiber.c
@@ -16,15 +16,20 @@
 #include <libtf/fiber.h>
 #include <libtf/io.h>
 
+#define TF_TIMEOUT_CHANGE_NEEDED			1
+#define TF_TIMEOUT_CHANGE_NEW_VALUE			2
+
 struct tf_fiber {
 	unsigned int		ref_count;
 	int			wakeup_type;
+	unsigned int		timeout_change;
+	tf_mtime_t		timeout;
 	struct tf_list_node	queue_node;
 	struct tf_heap_node	heap_node;
 	char			data[TF_EMPTY_ARRAY];
 };
 
-#include TF_UCTX_H
+#include "uctx-setjmp.h"
 
 /* FIXME: should be in thread local storage */
 struct tf_scheduler *__tf_scheduler;
@@ -108,6 +113,8 @@ static void process_heap(struct tf_scheduler *sched)
 	       tf_mtime_diff(now, tf_heap_get_value(&sched->heap)) >= 0) {
 		node = tf_heap_get_node(&sched->heap);
 		f = container_of(node, struct tf_fiber, heap_node);
+		if (f->wakeup_type == TF_WAKEUP_NONE)
+			f->wakeup_type = TF_WAKEUP_TIMEOUT;
 		run_fiber(sched, f);
 	}
 }
@@ -170,18 +177,58 @@ int tf_main_args(tf_fiber_proc main_fiber, int argc, char **argv)
 	return 0;
 }
 
-int tf_schedule(tf_mtime_diff_t milliseconds)
+void tf_timeout_push(struct tf_timeout *timeout, tf_mtime_diff_t milliseconds)
+{
+	struct tf_fiber *f = tf_get_fiber();
+	tf_mtime_t abs = tf_mtime() + milliseconds;
+	int active;
+
+	if (f->timeout_change)
+		active = (f->timeout_change & TF_TIMEOUT_CHANGE_NEW_VALUE);
+	else
+		active = tf_heap_node_active(&f->heap_node);
+
+	if (!active || tf_mtime_diff(abs, f->timeout) < 0) {
+		/* Save previous timeout */
+		timeout->saved_timeout = f->timeout;
+		timeout->timeout_change = TF_TIMEOUT_CHANGE_NEEDED;
+		if (active)
+			timeout->timeout_change |= TF_TIMEOUT_CHANGE_NEW_VALUE;
+
+		/* Make new timeout pending */
+		f->timeout = abs;
+		f->timeout_change = TF_TIMEOUT_CHANGE_NEEDED
+				  | TF_TIMEOUT_CHANGE_NEW_VALUE;
+	} else {
+		timeout->timeout_change = 0;
+	}
+}
+
+int __tf_timeout_pop(struct tf_timeout *timeout, int err)
+{
+	struct tf_fiber *f = tf_get_fiber();
+
+	f->timeout = timeout->saved_timeout;
+	f->timeout_change = timeout->timeout_change;
+	if (err == TF_WAKEUP_TIMEOUT)
+		err = TF_WAKEUP_THIS_TIMEOUT;
+	return err;
+}
+
+int tf_schedule(void)
 {
 	struct tf_scheduler *sched = tf_get_scheduler();
 	struct tf_fiber *schedf = container_of((void*) sched, struct tf_fiber, data);
 	struct tf_fiber *f = sched->active_fiber;
 
+	if (unlikely(f->timeout_change)) {
+		if (f->timeout_change & TF_TIMEOUT_CHANGE_NEW_VALUE)
+			tf_heap_change(&f->heap_node, &sched->heap, f->timeout);
+		else
+			tf_heap_delete(&f->heap_node, &sched->heap);
+		f->timeout_change = 0;
+	}
 	f->wakeup_type = TF_WAKEUP_NONE;
-	if (milliseconds == TF_NO_TIMEOUT)
-		tf_heap_delete(&f->heap_node, &sched->heap);
-	else if (milliseconds != TF_NO_TIMEOUT_CHANGE)
-		tf_heap_change(&f->heap_node, &sched->heap,
-			       tf_mtime() + milliseconds);
 
 	return tf_uctx_transfer(f, schedf, TF_WAKEUP_IMMEDIATE);
 }
@@ -217,6 +264,6 @@ int tf_yield(void)
 	struct tf_fiber *f = sched->active_fiber;
 
 	tf_list_add_tail(&f->queue_node, &sched->run_q);
-	return tf_schedule(TF_NO_TIMEOUT);
+	return tf_schedule();
 }