lib: embedded ATQ

lib/atq.bpf.c

@@ -17,7 +17,7 @@ u64 scx_atq_create_internal(bool fifo, size_t capacity)
 	if (!atq)
 		return (u64)NULL;
 
-	atq->tree = rb_create();
+	atq->tree = rb_create(RB_NOALLOC, RB_DUPLICATE);
 	if (!atq->tree)
 		return (u64)NULL;
 
@@ -28,99 +28,65 @@ u64 scx_atq_create_internal(bool fifo, size_t capacity)
 	return (u64)atq;
 }
 
+/* 
+ * XXXETSAL: We are using the __hidden antipattern for API functions because some
+ * older kernels do not allow function calls with preemption disabled. We will replace
+ * these annotations with the proper ones (__weak) at some point in the future.
+ */
+
 __hidden
-int scx_atq_insert(scx_atq_t *atq, u64 taskc_ptr)
+int scx_atq_insert_vtime(scx_atq_t __arg_arena *atq, rbnode_t __arg_arena *node, u64 taskc_ptr, u64 vtime)
 {
-	rbnode_t *node;
 	int ret;
 
-	if (!atq->fifo)
-		return -EINVAL;
-
-	/*
-	 * Use dummy sequence number because we're
-	 * outside of the critical section.
-	 */
-	node = rb_node_alloc(atq->tree, 0, taskc_ptr);
-	if (!node)
-		return -ENOMEM;
-
 	ret = arena_spin_lock(&atq->lock);
-	if (ret) {
-		rb_node_free(atq->tree, node);
+	if (ret)
 		return ret;
-	}
 
 	if (unlikely(atq->size == atq->capacity)) {
 		ret = -ENOSPC;
-		goto error;
+		goto done;
+	}
+
+	if ((vtime == SCX_ATQ_FIFO) != atq->fifo) {
+		ret = -EINVAL;
+		goto done;
 	}
 
 	/*
-	 * "Leak" the seq on error. We only want
+	 * For FIFO, "Leak" the seq on error. We only want
 	 * sequence numbers to be monotonic, not
 	 * consecutive.
 	 */
-	node->key = atq->seq++;
+	node->key = (vtime == SCX_ATQ_FIFO) ? atq->seq++ : vtime;
+	node->value = taskc_ptr;
 
-	ret = rb_insert_node(atq->tree, node, RB_DUPLICATE);
+	ret = rb_insert_node(atq->tree, node);
 	if (ret)
-		goto error;
+		goto done;
 
 	atq->size += 1;
 
+done:
 	arena_spin_unlock(&atq->lock);
 
-	return 0;
-
-error:
-	arena_spin_unlock(&atq->lock);
-	rb_node_free(atq->tree, node);
-
 	return ret;
 }
 
 __hidden
-int scx_atq_insert_vtime(scx_atq_t *atq, u64 taskc_ptr, u64 vtime)
+int scx_atq_insert(scx_atq_t *atq, rbnode_t __arg_arena *node, u64 taskc_ptr)
 {
-	rbnode_t *node;
-	int ret;
-
-	if (atq->fifo)
-		return -EINVAL;
-
-	node = rb_node_alloc(atq->tree, vtime, taskc_ptr);
-	if (!node)
-		return -ENOMEM;
-
-	ret = arena_spin_lock(&atq->lock);
-	if (ret) {
-		rb_node_free(atq->tree, node);
-		return ret;
-	}
-
-	if (unlikely(atq->size == atq->capacity)) {
-		ret = -ENOSPC;
-		goto error;
-	}
-
-	ret = rb_insert_node(atq->tree, node, RB_DUPLICATE);
-	if (ret)
-		goto error;
-
-	atq->size += 1;
-
-	arena_spin_unlock(&atq->lock);
-
-	return 0;
-
-error:
-	arena_spin_unlock(&atq->lock);
-	rb_node_free(atq->tree, node);
-
-	return ret;
+	return scx_atq_insert_vtime(atq, node, taskc_ptr, SCX_ATQ_FIFO);
 }
 
+/*
+ * XXXETSAL: There is a mismatch between insert and pop here: We are inserting
+ * rbnodes, but returning a key/value pair. This is deliberate: We can use CO:RE
+ * to find the rbnode from any scheduler's task_ctx in a generic way, but there is
+ * no container_of equivalent that lets us go rbnode -> task_ctx (especially since
+ * the actual layout of task_ctx varies by scheduler. For now, pass the task_ctx 
+ * as a value to the node and use it to find the original rbonde.
+ */
 __hidden
 u64 scx_atq_pop(scx_atq_t *atq)
 {