Add reference counting to token.

socketry · Jul 29, 2023 · 57ba727 · 57ba727
1 parent c506ecf
commit 57ba727
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 7 deletions.
diff --git a/ext/io/event/selector/epoll.c b/ext/io/event/selector/epoll.c
@@ -93,6 +93,7 @@ VALUE IO_Event_Selector_EPoll_allocate(VALUE self) {
 
 	IO_Event_Selector_initialize(&data->backend, Qnil);
 	data->descriptor = -1;
+	data->blocked = 0;
 	data->tokens = NULL;
 
 	return instance;
@@ -308,12 +309,19 @@ VALUE io_wait_ensure(VALUE _arguments) {
 
 	IO_Event_Token_cancel(&arguments->data->tokens, arguments->token);
 
+	int result = 0;
+
 	if (arguments->duplicate >= 0) {
-		epoll_ctl(arguments->data->descriptor, EPOLL_CTL_DEL, arguments->duplicate, NULL);
+		result = epoll_ctl(arguments->data->descriptor, EPOLL_CTL_DEL, arguments->duplicate, NULL);
 
 		close(arguments->duplicate);
 	} else {
-		epoll_ctl(arguments->data->descriptor, EPOLL_CTL_DEL, arguments->descriptor, NULL);
+		result = epoll_ctl(arguments->data->descriptor, EPOLL_CTL_DEL, arguments->descriptor, NULL);
+	}
+
+	if (result == 0) {
+		// It was removed from the epoll set, so we can release it:
+		IO_Event_Token_release(&arguments->data->tokens, arguments->token);
 	}
 
 	return Qnil;

diff --git a/ext/io/event/selector/selector.c b/ext/io/event/selector/selector.c
@@ -341,8 +341,6 @@ void IO_Event_Selector_current_time(struct timespec *time) {
 
 struct IO_Event_Token * IO_Event_Token_acquire(struct IO_Event_Token **head, void * data)
 {
-	fprintf(stderr, "IO_Event_Token_acquire head=%p\n", *head);
-
 	struct IO_Event_Token *token = *head;
 
 	if (token) {
@@ -352,16 +350,21 @@ struct IO_Event_Token * IO_Event_Token_acquire(struct IO_Event_Token **head, voi
 	}
 
 	token->data = data;
+	token->count = 1;
 
 	return token;
 }
 
 void IO_Event_Token_release(struct IO_Event_Token **head, struct IO_Event_Token *token)
 {
-	token->data = NULL;
+	assert(token->count > 0);
+
+	token->count -= 1;
 
-	token->next = *head;
-	*head = token;
+	if (token->count == 0) {
+		token->next = *head;
+		*head = token;
+	}
 }
 
 void IO_Event_Token_free(struct IO_Event_Token **head)
@@ -371,6 +374,7 @@ void IO_Event_Token_free(struct IO_Event_Token **head)
 		*head = token->next;
 
 		assert(token->data == NULL);
+		assert(token->count == 0);
 
 		free(token);
 	}
@@ -380,6 +384,9 @@ void * IO_Event_Token_wrap(struct IO_Event_Token **head, VALUE fiber)
 {
 	struct IO_Event_Token *token = IO_Event_Token_acquire(head, (void*)fiber);
 
+	// We deliberately increment the reference count as ownership here is bifurcated between the fiber and the selector.
+	token->count += 1;
+
 	return (void*)token;
 }
 
@@ -388,13 +395,17 @@ void IO_Event_Token_cancel(struct IO_Event_Token **head, void * data)
 	struct IO_Event_Token *token = (struct IO_Event_Token *)data;
 
 	token->data = NULL;
+
+	IO_Event_Token_release(head, token);
 }
 
 VALUE IO_Event_Token_unwrap(struct IO_Event_Token **head, void * data)
 {
 	struct IO_Event_Token *token = (struct IO_Event_Token *)data;
 
 	VALUE fiber = (VALUE)token->data;
+	token->data = NULL;
+
 	IO_Event_Token_release(head, token);
 
 	return fiber;

diff --git a/ext/io/event/selector/selector.h b/ext/io/event/selector/selector.h
@@ -135,7 +135,14 @@ void IO_Event_Selector_current_time(struct timespec *time);
 
 struct IO_Event_Token {
 	struct IO_Event_Token *next;
+
+	// The fiber that is waiting for the operation to be completed (or cancelled).
 	void * data;
+
+	// The reference count. Both the fiber and the selector may hold a reference to the token.
+	// The fiber holds a reference while it's waiting for the operation to be completed (or cancelled).
+	// The selector holds a reference while the operation is in progress.
+	unsigned count;
 };
 
 // Manage a pool of tokens: