diff --git a/tests/test_auxiliary_procs.nim b/tests/test_auxiliary_procs.nim new file mode 100644 index 0000000..55b632c --- /dev/null +++ b/tests/test_auxiliary_procs.nim @@ -0,0 +1,18 @@ +import + std/[atomics, os], + ../weave + +var count: Atomic[int] + +proc increment() = count += 1 +proc decrement() = count -= 1 + +const ThreadsToLaunch = 8 # Arbitrary value + +putEnv("WEAVE_NUM_THREADS", $ThreadsToLaunch) + +init(Weave, increment) +syncRoot(Weave) +doAssert(count.load == ThreadsToLaunch - 1) +exit(Weave, decrement) +doAssert(count.load == 0) \ No newline at end of file diff --git a/weave.nimble b/weave.nimble index 40cb216..4f28e0a 100644 --- a/weave.nimble +++ b/weave.nimble @@ -47,6 +47,7 @@ task test, "Run Weave tests": test "", "weave/parallel_reduce.nim" test "--debugger:native", "tests/test_background_jobs.nim" + test "--debugger:native", "tests/test_auxiliary_procs.nim" test "-d:WV_LazyFlowvar", "weave/parallel_tasks.nim" test "-d:WV_LazyFlowvar", "weave/parallel_for.nim" @@ -117,6 +118,7 @@ task test_gc_arc, "Run Weave tests with --gc:arc": test "--gc:arc", "weave/parallel_reduce.nim" test "--gc:arc", "tests/test_background_jobs.nim" + test "--gc:arc", "tests/test_auxiliary_procs.nim" test "--gc:arc -d:WV_LazyFlowvar", "weave/parallel_tasks.nim" test "--gc:arc -d:WV_LazyFlowvar", "weave/parallel_for.nim" diff --git a/weave/datatypes/context_global.nim b/weave/datatypes/context_global.nim index 184e961..58dc2f6 100644 --- a/weave/datatypes/context_global.nim +++ b/weave/datatypes/context_global.nim @@ -50,6 +50,7 @@ type barrier*: SyncBarrier ## Barrier for initialization and teardown manager*: ManagerContext + auxiliaryInit*, auxiliaryExit*: proc() {.nimcall, gcsafe.} ManagerContext* = object ## Manager context diff --git a/weave/runtime.nim b/weave/runtime.nim index dc63715..a543f8f 100644 --- a/weave/runtime.nim +++ b/weave/runtime.nim @@ -29,7 +29,10 @@ else: # Runtime public routines # ---------------------------------------------------------------------------------- -proc init*(_: type Weave) = +type AuxiliaryProc* = proc() {.nimcall, gcsafe.} | proc() {.cdecl, gcsafe.} +const WV_NoAuxiliary: proc() {.nimcall, gcsafe.} = nil # for type resolution + +proc init*(_: type Weave, auxiliary: AuxiliaryProc = WV_NoAuxiliary) = # TODO detect Hyper-Threading and NUMA domain manager.acceptsJobs.store(false, moRelaxed) @@ -42,6 +45,11 @@ proc init*(_: type Weave) = else: workforce() = int32 countProcessors() + when auxiliary is proc() {.cdecl, gcsafe.}: + globalCtx.auxiliaryInit = proc() {.nimcall, gcsafe.} = auxiliary + else: + globalCtx.auxiliaryInit = auxiliary + ## Allocation of the global context. globalCtx.mempools = wv_alloc(TLPoolAllocator, workforce()) globalCtx.threadpool = wv_alloc(Thread[WorkerID], workforce()) @@ -167,7 +175,12 @@ proc globalCleanup() = metrics: log("+========================================+\n") -proc exit*(_: type Weave) = +proc exit*(_: type Weave, auxiliary: AuxiliaryProc = WV_NoAuxiliary) = + when auxiliary is proc() {.cdecl, gcsafe.}: + globalCtx.auxiliaryExit = proc() {.nimcall, gcsafe.} = auxiliary + else: + globalCtx.auxiliaryExit = auxiliary + syncRoot(_) signalTerminate(nil) workerContext.signaledTerminate = true diff --git a/weave/scheduler.nim b/weave/scheduler.nim index 92f8089..6fd5a16 100644 --- a/weave/scheduler.nim +++ b/weave/scheduler.nim @@ -171,6 +171,11 @@ proc worker_entry_fn*(id: WorkerID) = myID() = id # If this crashes, you need --tlsemulation:off myMemPool().initialize() setupWorker() + + # Unstealable user-definable procedure to be executed on all threads + {.gcsafe.}: + if not globalCtx.auxiliaryInit.isNil: globalCtx.auxiliaryInit() + discard globalCtx.barrier.wait() eventLoop() @@ -181,6 +186,10 @@ proc worker_entry_fn*(id: WorkerID) = # 1 matching barrier in init(Runtime) for lead thread workerMetrics() + # Same as auxiliaryInit, but for cleanup + {.gcsafe.}: + if not globalCtx.auxiliaryExit.isNil: globalCtx.auxiliaryExit() + teardownWorker() postCondition: localThreadKind == Unknown