1 files changed, 124 insertions, 27 deletions

diff --git a/cbits/hook.c b/cbits/hook.c
index 2175be4..d8619ff 100644
--- a/cbits/hook.c
+++ b/cbits/hook.c
@@ -100,8 +100,14 @@ static void shadow_copy(struct ShadowDetails *dst, const struct GCDetails_ *src)
 // GLOBAL VARIABLES
 // --------
 
+static bool constructor_worked = false;
+static bool hook_initialised = false;
+static bool logging_enabled = false;
+static void (*hook_c_delegate)(const struct GCDetails_*) = NULL;
+
+static mtx_t state_mutex;
+
 static void (*old_hook)(const struct GCDetails_ *details) = NULL;
-static mtx_t detlog_mutex;
 static size_t detlog_capacity = 0, detlog_length = 0;
 static struct ShadowDetails *detlog = NULL;
 
@@ -116,13 +122,13 @@ static void hook_callback(const struct GCDetails_ *details) {
 
 	// Do this now already, before waiting on the mutex
 	struct timespec now;
-	if (clock_gettime(CLOCK_MONOTONIC, &now) != 0) {
+	if (logging_enabled && clock_gettime(CLOCK_MONOTONIC, &now) != 0) {
 		perror("clock_gettime");
 		fatal_failure = true;
 		goto cleanup_no_mutex;
 	}
 
-	if (mtx_lock(&detlog_mutex) != thrd_success) {
+	if (mtx_lock(&state_mutex) != thrd_success) {
 		fprintf(stderr, "ghc-gc-hook: ERROR: Mutex lock failed\n");
 		fatal_failure = true;
 		goto cleanup_no_mutex;
@@ -130,62 +136,153 @@ static void hook_callback(const struct GCDetails_ *details) {
 
 	// mutex is locked from here
 
-	if (detlog_length == detlog_capacity) {
-		detlog_capacity = detlog_capacity == 0 ? 128 : 2 * detlog_capacity;
-		detlog = realloc(detlog, detlog_capacity * sizeof(detlog[0]));
-		if (detlog == NULL || detlog_capacity == 0) {  // also check for overflow here
-			fprintf(stderr, "ghc-gc-hook: ERROR: Could not allocate memory for GC log hook\n");
-			fatal_failure = true;
-			goto cleanup;
+	if (logging_enabled) {
+		if (detlog_length == detlog_capacity) {
+			detlog_capacity = detlog_capacity == 0 ? 128 : 2 * detlog_capacity;
+			detlog = realloc(detlog, detlog_capacity * sizeof(detlog[0]));
+			if (detlog == NULL || detlog_capacity == 0) {  // also check for overflow here
+				fprintf(stderr, "ghc-gc-hook: ERROR: Could not allocate memory for GC log hook\n");
+				fatal_failure = true;
+				goto cleanup;
+			}
 		}
+
+		struct ShadowDetails *dst = &detlog[detlog_length];
+		dst->timestamp_sec = now.tv_sec;
+		dst->timestamp_nsec = now.tv_nsec;
+		shadow_copy(dst, details);
+		detlog_length++;
 	}
 
-	struct ShadowDetails *dst = &detlog[detlog_length];
-	dst->timestamp_sec = now.tv_sec;
-	dst->timestamp_nsec = now.tv_nsec;
-	shadow_copy(dst, details);
-	detlog_length++;
+	if (hook_c_delegate) hook_c_delegate(details);
 
 cleanup:
-	mtx_unlock(&detlog_mutex);  // ignore return value
+	mtx_unlock(&state_mutex);  // ignore return value
 
 cleanup_no_mutex:
 	if (old_hook) old_hook(details);
 }
 
+__attribute__((constructor))
+static void constructor(void) {
+	if (mtx_init(&state_mutex, mtx_plain) != thrd_success) {
+		fprintf(stderr, "ghc-gc-hook: ERROR: Mutex initialisation failed\n");
+		return;
+	}
+
+	constructor_worked = true;
+}
+
 // --------
 // EXPORTED FUNCTIONS
 // --------
 
+// Only works if logging is enabled.
 void copy_log_to_buffer(size_t space_available, char *buffer, size_t *unit_size, size_t *num_stored) {
-	if (mtx_lock(&detlog_mutex) != thrd_success) {
+	*unit_size = sizeof(detlog[0]);
+
+	if (mtx_lock(&state_mutex) != thrd_success) {
 		fprintf(stderr, "ghc-gc-hook: ERROR: Mutex lock failed\n");
-		*unit_size = 0;
 		*num_stored = 0;
 		return;
 	}
 
-	const size_t sz = sizeof(detlog[0]);
-	const size_t n = min_sz(space_available / sz, detlog_length);
+	if (detlog_length == 0) {
+		*num_stored = 0;
+		goto unlock_return;
+	}
+
+	const size_t n = min_sz(space_available / sizeof(detlog[0]), detlog_length);
 
 	// First copy over the fitting items
-	memcpy(buffer, detlog, n * sz);
-	*unit_size = sz;
+	memcpy(buffer, detlog, n * sizeof(detlog[0]));
+	*unit_size = sizeof(detlog[0]);
 	*num_stored = n;
 
 	// Then shift back the remaining items
 	memmove(detlog, detlog + n, (detlog_length - n) * sizeof(detlog[0]));
 	detlog_length -= n;
 
-	mtx_unlock(&detlog_mutex);
+unlock_return:
+	mtx_unlock(&state_mutex);
 }
 
-void set_gchook() {
-	if (mtx_init(&detlog_mutex, mtx_plain) != thrd_success) {
-		fprintf(stderr, "ghc-gc-hook: ERROR: Mutex initialisation failed\n");
-		return;
+// Sets the GC hook, logging or C hook delegate not yet enabled. Returns success.
+bool set_gchook(void) {
+	if (mtx_lock(&state_mutex) != thrd_success) {
+		fprintf(stderr, "ghc-gc-hook: ERROR: Mutex lock failed\n");
+		return false;
+	}
+
+	bool retval = false;
+
+	if (!constructor_worked) {
+		fprintf(stderr, "ghc-gc-hook: ERROR: Cannot set hook, system does not allow initialisation\n");
+		goto unlock_return_retval;
+	}
+
+	if (hook_initialised) {
+		fprintf(stderr, "ghc-gc-hook: ERROR: Hook already initialised\n");
+		goto unlock_return_retval;
 	}
 
 	old_hook = rtsConfig.gcDoneHook;
 	rtsConfig.gcDoneHook = hook_callback;
+
+	hook_initialised = true;
+	retval = true;
+
+unlock_return_retval:
+	mtx_unlock(&state_mutex);
+	return retval;
+}
+
+// Enable logging on the GC hook.
+void gchook_enable_logging(bool yes) {
+	if (!hook_initialised) {
+		if (!set_gchook()) exit(1);  // meh
+	}
+
+	if (mtx_lock(&state_mutex) != thrd_success) {
+		fprintf(stderr, "ghc-gc-hook: ERROR: Mutex lock failed\n");
+		return;
+	}
+
+	if (logging_enabled && !yes) {
+		detlog_length = 0;
+		detlog_capacity = 0;
+		free(detlog);
+		detlog = NULL;
+	}
+
+	logging_enabled = yes;
+
+	mtx_unlock(&state_mutex);
+}
+
+// Set a C function to be called after every GC with the GCDetails_ structure
+// from `rts/include/RtsAPI.h`. Returns success.
+bool gchook_set_c_delegate(void (*delegate)(const struct GCDetails_*)) {
+	if (!hook_initialised) {
+		if (!set_gchook()) exit(1);  // meh
+	}
+
+	if (mtx_lock(&state_mutex) != thrd_success) {
+		fprintf(stderr, "ghc-gc-hook: ERROR: Mutex lock failed\n");
+		return false;
+	}
+
+	bool retval = false;
+
+	if (hook_c_delegate != NULL) {
+		fprintf(stderr, "ghc-gc-hook: ERROR: C hook delegate already set\n");
+		goto unlock_return_retval;
+	}
+
+	hook_c_delegate = delegate;
+	retval = true;
+
+unlock_return_retval:
+	mtx_unlock(&state_mutex);
+	return retval;
 }