From ec068f7785819492cd555ea71b608e58efe01d6d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20S=C3=A1nchez=20Ram=C3=ADrez?=
 <sergio.sanchez.ramirez+git@bsc.es>
Date: Wed, 26 Jun 2024 11:53:54 +0200
Subject: [PATCH] Fix initialization of `Distributed`

---
 Project.toml                |  6 ++++++
 ext/ExtraeDistributedExt.jl | 21 +++++++++++++++++++++
 2 files changed, 27 insertions(+)
 create mode 100644 ext/ExtraeDistributedExt.jl

diff --git a/Project.toml b/Project.toml
index 3eeaca3..0936b4d 100644
--- a/Project.toml
+++ b/Project.toml
@@ -6,6 +6,12 @@ version = "0.1.0"
 [deps]
 Extrae_jll = "2b2c4be0-e38c-5918-b8b4-9a308845a1e9"
 
+[weakdeps]
+Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
+
+[extension]
+ExtraeDistributedExt = "Distributed"
+
 [compat]
 Extrae_jll = "4.1.2"
 julia = "1.9"
diff --git a/ext/ExtraeDistributedExt.jl b/ext/ExtraeDistributedExt.jl
new file mode 100644
index 0000000..429edae
--- /dev/null
+++ b/ext/ExtraeDistributedExt.jl
@@ -0,0 +1,21 @@
+module ExtraeDistributedExt
+
+using Extrae
+using Distributed
+
+taskid()::Cuint = Distributed.myid() - 1
+ntasks()::Cuint = Distributed.nworkers() + 1
+
+function init(::Val{:Distributed})
+    FFI.Extrae_set_numtasks_function(ntasks)
+    FFI.Extrae_set_taskid_function(taskid)
+
+    ENV["EXTRAE_PROGRAM_NAME"] = "JULIATRACE-$(Distributed.myid())"
+
+    Extrae.init()
+    Libc.flush_cstdio()
+
+    return nothing
+end
+
+end