From 8c9088e0182de25cc6503e73928a430a36c17ec1 Mon Sep 17 00:00:00 2001
From: JinZhou5042 <142265839+JinZhou5042@users.noreply.github.com>
Date: Wed, 5 Feb 2025 12:35:08 -0500
Subject: [PATCH] vine: track running libraries on the worker (#4047)

* init

* delete and clean current_libraries
---
 taskvine/src/manager/vine_manager.c     | 9 +++++++++
 taskvine/src/manager/vine_worker_info.c | 2 ++
 taskvine/src/manager/vine_worker_info.h | 1 +
 3 files changed, 12 insertions(+)

diff --git a/taskvine/src/manager/vine_manager.c b/taskvine/src/manager/vine_manager.c
index 77a55a9e88..ea99eb81eb 100644
--- a/taskvine/src/manager/vine_manager.c
+++ b/taskvine/src/manager/vine_manager.c
@@ -930,6 +930,7 @@ static void cleanup_worker(struct vine_manager *q, struct vine_worker_info *w)
 	}
 
 	itable_clear(w->current_tasks, 0);
+	itable_clear(w->current_libraries, 0);
 
 	w->finished_tasks = 0;
 
@@ -2972,6 +2973,10 @@ static vine_result_code_t commit_task_to_worker(struct vine_manager *q, struct v
 		/* If start_one_task_fails, this will be decremented in handle_failure below. */
 		t->library_task->function_slots_inuse++;
 	}
+	/* If this is a library task, bookkeep it on the worker's side */
+	if (t->provides_library) {
+		itable_insert(w->current_libraries, t->task_id, t);
+	}
 
 	t->hostname = xxstrdup(w->hostname);
 	t->addrport = xxstrdup(w->addrport);
@@ -3156,6 +3161,10 @@ static void reap_task_from_worker(struct vine_manager *q, struct vine_worker_inf
 
 	itable_remove(w->current_tasks, t->task_id);
 
+	if (t->provides_library) {
+		itable_remove(w->current_libraries, t->task_id);
+	}
+
 	/*
 	If this was a function call assigned to a library,
 	then decrease the count of functions assigned,
diff --git a/taskvine/src/manager/vine_worker_info.c b/taskvine/src/manager/vine_worker_info.c
index ce9158959f..fef94e8b97 100644
--- a/taskvine/src/manager/vine_worker_info.c
+++ b/taskvine/src/manager/vine_worker_info.c
@@ -34,6 +34,7 @@ struct vine_worker_info *vine_worker_create(struct link *lnk)
 
 	w->current_files = hash_table_create(0, 0);
 	w->current_tasks = itable_create(0);
+	w->current_libraries = itable_create(0);
 
 	w->start_time = timestamp_get();
 	w->end_time = -1;
@@ -69,6 +70,7 @@ void vine_worker_delete(struct vine_worker_info *w)
 	hash_table_clear(w->current_files, (void *)vine_file_replica_delete);
 	hash_table_delete(w->current_files);
 	itable_delete(w->current_tasks);
+	itable_delete(w->current_libraries);
 
 	free(w);
 
diff --git a/taskvine/src/manager/vine_worker_info.h b/taskvine/src/manager/vine_worker_info.h
index c52fa91103..e7e67122ad 100644
--- a/taskvine/src/manager/vine_worker_info.h
+++ b/taskvine/src/manager/vine_worker_info.h
@@ -62,6 +62,7 @@ struct vine_worker_info {
 	/* Current files and tasks that have been transfered to this worker */
 	struct hash_table   *current_files;
 	struct itable       *current_tasks;
+	struct itable		*current_libraries;
 
 	/* The number of tasks running last reported by the worker */
 	int         dynamic_tasks_running;