From 9ece196d350e501d2db7f15938d9b0c2193fc72d Mon Sep 17 00:00:00 2001
From: Vitali Lovich <vlovich@cloudflare.com>
Date: Fri, 11 Mar 2022 07:15:20 -0800
Subject: [PATCH] Add pre_block/post_block/wake_other

Fixes #44
---
 include/coz.h | 63 +++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 46 insertions(+), 17 deletions(-)

diff --git a/include/coz.h b/include/coz.h
index 6f4bf1c..a33d1f1 100644
--- a/include/coz.h
+++ b/include/coz.h
@@ -37,29 +37,55 @@ typedef struct {
   size_t backoff;  // Used to batch updates to the shared counter. Currently unused.
 } coz_counter_t;
 
-// The type of the _coz_get_counter function
-typedef coz_counter_t* (*coz_get_counter_t)(int, const char*);
-
-// Locate and invoke _coz_get_counter
 static coz_counter_t* _call_coz_get_counter(int type, const char* name) {
-  static unsigned char _initialized = 0;
-  static coz_get_counter_t fn; // The pointer to _coz_get_counter
+  // Use memcpy to avoid pedantic GCC complaint about storing function pointer in void*
+  static unsigned char initialized = dlsym == NULL;
+  static coz_counter_t* (*fn)(int, const char*);
+
+  if (!initialized) {
+    void* p = dlsym(RTLD_DEFAULT, "_coz_get_counter");
+    memcpy(&fn, &p, sizeof(p));
+    initialized = true;
+  }
 
-  if(!_initialized) {
-    if(dlsym) {
-      // Locate the _coz_get_counter method
-      void* p = dlsym(RTLD_DEFAULT, "_coz_get_counter");
+  if (fn) return fn(type, name);
+  else return 0;
+}
 
-      // Use memcpy to avoid pedantic GCC complaint about storing function pointer in void*
-      memcpy(&fn, &p, sizeof(p));
-    }
+static void _call_coz_pre_block() {
+  static unsigned char initialized = dlsym == 0;
+  static void (*fn)();
 
-    _initialized = 1;
+  if (!initialized) {
+    void* p = dlsym(RTLD_DEFAULT, "_coz_pre_block");
+    memcpy(&fn, &p, sizeof(p));
+    initialized = true;
   }
+  if (fn) return fn();
+}
 
-  // Call the function, or return null if profiler is not found
-  if(fn) return fn(type, name);
-  else return 0;
+static void _call_coz_post_block() {
+  static unsigned char initialized = dlsym == 0;
+  static void (*fn)();
+
+  if (!initialized) {
+    void* p = dlsym(RTLD_DEFAULT, "_coz_post_block");
+    memcpy(&fn, &p, sizeof(p));
+    initialized = true;
+  }
+  if (fn) return fn();
+}
+
+static void _call_coz_wake_other() {
+  static unsigned char initialized = dlsym == 0;
+  static void (*fn)();
+
+  if (!initialized) {
+    void* p = dlsym(RTLD_DEFAULT, "_coz_wake_other");
+    memcpy(&fn, &p, sizeof(p));
+    initialized = true;
+  }
+  if (fn) return fn();
 }
 
 // Macro to initialize and increment a counter
@@ -85,6 +111,9 @@ static coz_counter_t* _call_coz_get_counter(int type, const char* name) {
 #define COZ_PROGRESS COZ_INCREMENT_COUNTER(COZ_COUNTER_TYPE_THROUGHPUT, __FILE__ ":" STR(__LINE__))
 #define COZ_BEGIN(name) COZ_INCREMENT_COUNTER(COZ_COUNTER_TYPE_BEGIN, name)
 #define COZ_END(name) COZ_INCREMENT_COUNTER(COZ_COUNTER_TYPE_END, name)
+#define COZ_PRE_BLOCK() _call_coz_pre_block()
+#define COZ_POST_BLOCK() _call_coz_post_block()
+#define COZ_WAKE_OTHER() _call_coz_wake_other()
 
 #if defined(__cplusplus)
 }