From d8361d53acbf6514b352bb5ee21102a7d2110ced Mon Sep 17 00:00:00 2001
From: Jaremy Creechley <creechley@gmail.com>
Date: Tue, 28 Nov 2023 22:35:35 -0700
Subject: [PATCH] move to event based enum for api

---
 chronos/profiler.nim         | 12 ++-----
 chronos/profiler/events.nim  | 62 ------------------------------------
 chronos/profiler/metrics.nim | 40 ++++++++++++++++++-----
 3 files changed, 35 insertions(+), 79 deletions(-)
 delete mode 100644 chronos/profiler/events.nim

diff --git a/chronos/profiler.nim b/chronos/profiler.nim
index 1f43f11..0953bc3 100644
--- a/chronos/profiler.nim
+++ b/chronos/profiler.nim
@@ -15,14 +15,8 @@ when chronosProfiling:
     ## Returns metrics for the current event loop.
     result = futureMetrics
 
-  proc enableEventCallbacks*(): void =
-    onFutureEvent = handleFutureEventCB
-    onFutureExecEvent = handleFutureExecEventCB
-    
+  proc setFutureEventCallback*(): void =
+
   proc enableProfiling*() =
     ## Enables profiling on the current event loop.
-    if not isNil(handleFutureEvent): return
-
-    enableEventCallbacks()
-    handleFutureEvent = proc (e: Event) {.nimcall.} = 
-      futureMetrics.processEvent(e)
+    onFutureEvent = metrics.handleFutureEvent
diff --git a/chronos/profiler/events.nim b/chronos/profiler/events.nim
deleted file mode 100644
index c9e7346..0000000
--- a/chronos/profiler/events.nim
+++ /dev/null
@@ -1,62 +0,0 @@
-## This module defines the lower-level callback implementations that hook into
-## the Chronos scheduler when profiling is enabled. The main goal is to provide
-## timestamped events changes for futures while allowing a simpler implementation
-## (only one event object type) for the remainder of the profiler.
-
-import ".."/timer
-import ".."/futures
-import ".."/srcloc
-
-type
-  ExtendedFutureState* {.pure.} = enum
-    Pending,
-    Running,
-    Paused,
-    Completed,
-    Cancelled,
-    Failed,
-
-  Event* = object
-    future: FutureBase
-    newState*: ExtendedFutureState
-    timestamp*: Moment
-  
-var handleFutureEvent* {.threadvar.}: proc (event: Event) {.nimcall, gcsafe, raises: [].}
-
-proc `location`*(self: Event): SrcLoc =
-  self.future.internalLocation[Create][]
-
-proc `futureId`*(self: Event): uint =
-  self.future.id
-
-proc mkEvent(future: FutureBase, state: ExtendedFutureState): Event =
-  Event(
-    future: future,
-    newState: state,
-    timestamp: Moment.now(),
-  )
-  
-proc handleFutureEventCB*(future: FutureBase,
-    state: FutureState): void {.nimcall.} =
-  {.cast(gcsafe).}:
-    let extendedState = case state:
-      of FutureState.Pending: ExtendedFutureState.Pending
-      of FutureState.Completed: ExtendedFutureState.Completed
-      of FutureState.Cancelled: ExtendedFutureState.Cancelled
-      of FutureState.Failed: ExtendedFutureState.Failed
-
-    if not isNil(handleFutureEvent):
-      handleFutureEvent(mkEvent(future, extendedState))
-
-proc handleFutureExecEventCB*(future: FutureBase,
-    state: FutureExecutionState): void {.nimcall.} =
-  {.cast(gcsafe).}:
-    let extendedState = case state:
-      of FutureExecutionState.Running: ExtendedFutureState.Running
-      of FutureExecutionState.Paused: ExtendedFutureState.Paused
-
-    if not isNil(handleFutureEvent):
-      handleFutureEvent(mkEvent(future, extendedState))
-
-
-
diff --git a/chronos/profiler/metrics.nim b/chronos/profiler/metrics.nim
index 4c22bd1..69347ef 100644
--- a/chronos/profiler/metrics.nim
+++ b/chronos/profiler/metrics.nim
@@ -6,6 +6,33 @@ import ../[timer, srcloc]
 
 export timer, tables, srcloc
 
+type
+  Event* = object
+    futureId: uint
+    loc: SrcLoc
+    state: FutureState
+    event*: FutureExecutionEvent
+    timestamp*: Moment
+
+proc processEvent*(self: var ProfilerMetrics, event: Event): void
+
+proc handleFutureEvent*(future: FutureBase,
+                        event: FutureExecutionEvent): void {.nimcall.} =
+  {.cast(gcsafe).}:
+    let extendedState = case state:
+        of Init: ExtendedFutureState.Pending
+        of Run: ExtendedFutureState.Running
+        of Pause: ExtendedFutureState.Paused
+        of Finish: future.internalState.toExtendedEvent()
+
+    processEvent Event(
+        future: future,
+        loc: future.internalLocation[Create][],
+        state: future.internalState,
+        event: event,
+        timestamp: Moment.now()
+    )
+
 type
   AggregateFutureMetrics* = object
     execTime*: Duration
@@ -119,14 +146,11 @@ proc futureCompleted(self: var ProfilerMetrics, event: Event): void =
     self.partials.del(event.futureId)
 
 proc processEvent*(self: var ProfilerMetrics, event: Event): void =
-  case event.newState:
-  of Pending: self.futureCreated(event)
-  of Running: self.futureRunning(event)
-  of Paused: self.futurePaused(event)
-  # Completion, failure and cancellation are currently handled the same way.
-  of Completed: self.futureCompleted(event)
-  of Failed: self.futureCompleted(event)
-  of Cancelled: self.futureCompleted(event)
+  case event.event:
+  of Init: self.futureCreated(event)
+  of Run: self.futureRunning(event)
+  of Pause: self.futurePaused(event)
+  of Finish: self.futureCompleted(event)
 
 proc processAllEvents*(self: var ProfilerMetrics, events: seq[Event]): void =
   for event in events: