Add Arena tracing handlers.

PiperOrigin-RevId: 615779322

upb/mem/arena.c

@@ -7,6 +7,10 @@
 
 #include "upb/mem/arena.h"
 
+#ifdef UPB_TRACING_ENABLED
+#include <stdatomic.h>
+#endif
+
 #include <stddef.h>
 #include <stdint.h>
 
@@ -117,6 +121,38 @@ static bool _upb_ArenaInternal_HasInitialBlock(upb_ArenaInternal* ai) {
   return ai->block_alloc & 0x1;
 }
 
+#ifdef UPB_TRACING_ENABLED
+static void (*_init_arena_trace_handler)(const upb_Arena*, size_t size) = NULL;
+static void (*_fuse_arena_trace_handler)(const upb_Arena*,
+                                         const upb_Arena*) = NULL;
+static void (*_free_arena_trace_handler)(const upb_Arena*) = NULL;
+
+void upb_Arena_SetTraceHandler(
+    void (*initArenaTraceHandler)(const upb_Arena*, size_t size),
+    void (*fuseArenaTraceHandler)(const upb_Arena*, const upb_Arena*),
+    void (*freeArenaTraceHandler)(const upb_Arena*)) {
+  _init_arena_trace_handler = initArenaTraceHandler;
+  _fuse_arena_trace_handler = fuseArenaTraceHandler;
+  _free_arena_trace_handler = freeArenaTraceHandler;
+}
+
+void upb_Arena_LogInit(const upb_Arena* arena, size_t size) {
+  if (_init_arena_trace_handler) {
+    _init_arena_trace_handler(arena, size);
+  }
+}
+void upb_Arena_LogFuse(const upb_Arena* arena1, const upb_Arena* arena2) {
+  if (_fuse_arena_trace_handler) {
+    _fuse_arena_trace_handler(arena1, arena2);
+  }
+}
+void upb_Arena_LogFree(const upb_Arena* arena) {
+  if (_free_arena_trace_handler) {
+    _free_arena_trace_handler(arena);
+  }
+}
+#endif  // UPB_TRACING_ENABLED
+
 static upb_ArenaRoot _upb_Arena_FindRoot(upb_Arena* a) {
   upb_ArenaInternal* ai = upb_Arena_Internal(a);
   uintptr_t poc = upb_Atomic_Load(&ai->parent_or_count, memory_order_acquire);
@@ -280,7 +316,13 @@ upb_Arena* upb_Arena_Init(void* mem, size_t n, upb_alloc* alloc) {
   n = UPB_ALIGN_DOWN(n, UPB_ALIGN_OF(upb_ArenaState));
 
   if (UPB_UNLIKELY(n < sizeof(upb_ArenaState))) {
+#ifdef UPB_TRACING_ENABLED
+    upb_Arena* ret = _upb_Arena_InitSlow(alloc);
+    upb_Arena_LogInit(ret, n);
+    return ret;
+#else
     return _upb_Arena_InitSlow(alloc);
+#endif
   }
 
   a = UPB_PTR_AT(mem, n - sizeof(upb_ArenaState), upb_ArenaState);
@@ -293,13 +335,14 @@ upb_Arena* upb_Arena_Init(void* mem, size_t n, upb_alloc* alloc) {
   a->body.block_alloc = _upb_Arena_MakeBlockAlloc(alloc, 1);
   a->head.UPB_PRIVATE(ptr) = mem;
   a->head.UPB_PRIVATE(end) = UPB_PTR_AT(mem, n - sizeof(upb_ArenaState), char);
-
+#ifdef UPB_TRACING_ENABLED
+  upb_Arena_LogInit(&a->head, n);
+#endif
   return &a->head;
 }
 
 static void _upb_Arena_DoFree(upb_ArenaInternal* ai) {
   UPB_ASSERT(_upb_Arena_RefCountFromTagged(ai->parent_or_count) == 1);
-
   while (ai != NULL) {
     // Load first since arena itself is likely from one of its blocks.
     upb_ArenaInternal* next_arena =
@@ -330,6 +373,9 @@ retry:
   // expensive then direct loads.  As an optimization, we only do RMW ops
   // when we need to update things for other threads to see.
   if (poc == _upb_Arena_TaggedFromRefcount(1)) {
+#ifdef UPB_TRACING_ENABLED
+    upb_Arena_LogFree(a);
+#endif
     _upb_Arena_DoFree(ai);
     return;
   }
@@ -449,6 +495,10 @@ static bool _upb_Arena_FixupRefs(upb_ArenaInternal* new_root,
 bool upb_Arena_Fuse(upb_Arena* a1, upb_Arena* a2) {
   if (a1 == a2) return true;  // trivial fuse
 
+#ifdef UPB_TRACING_ENABLED
+  upb_Arena_LogFuse(a1, a2);
+#endif
+
   upb_ArenaInternal* ai1 = upb_Arena_Internal(a1);
   upb_ArenaInternal* ai2 = upb_Arena_Internal(a2);
 

upb/mem/arena.h

@@ -65,6 +65,14 @@ UPB_API_INLINE void* upb_Arena_Realloc(upb_Arena* a, void* ptr, size_t oldsize,
 UPB_API_INLINE void upb_Arena_ShrinkLast(upb_Arena* a, void* ptr,
                                          size_t oldsize, size_t size);
 
+#ifdef UPB_TRACING_ENABLED
+void upb_Arena_SetTraceHandler(void (*initArenaTraceHandler)(const upb_Arena*,
+                                                             size_t size),
+                               void (*fuseArenaTraceHandler)(const upb_Arena*,
+                                                             const upb_Arena*),
+                               void (*freeArenaTraceHandler)(const upb_Arena*));
+#endif
+
 #ifdef __cplusplus
 } /* extern "C" */
 #endif

upb/message/internal/message.c

@@ -62,13 +62,13 @@ bool UPB_PRIVATE(_upb_Message_Realloc)(struct upb_Message* msg, size_t need,
 static void (*_new_message_trace_handler)(const upb_MiniTable*,
                                           const upb_Arena*);
 
-void upb_Message_SetNewMessageTraceHandler(
+void UPB_PRIVATE(upb_Message_SetNewMessageTraceHandler)(
     void (*new_message_trace_handler)(const upb_MiniTable*, const upb_Arena*)) {
   _new_message_trace_handler = new_message_trace_handler;
 }
 
-void upb_Message_LogNewMessage(const upb_MiniTable* mini_table,
+void UPB_PRIVATE(upb_Message_LogNewMessage)(const upb_MiniTable* mini_table,
-                               const upb_Arena* arena) {
+                                            const upb_Arena* arena) {
   if (_new_message_trace_handler) {
     _new_message_trace_handler(mini_table, arena);
   }

upb/message/internal/message.h

@@ -60,17 +60,17 @@ typedef struct upb_Message_Internal {
 } upb_Message_Internal;
 
 #ifdef UPB_TRACING_ENABLED
-void upb_Message_SetNewMessageTraceHandler(
+void UPB_PRIVATE(upb_Message_SetNewMessageTraceHandler)(
     void (*newMessageTraceHandler)(const upb_MiniTable*, const upb_Arena*));
-void upb_Message_LogNewMessage(const upb_MiniTable* mini_table,
+void UPB_PRIVATE(upb_Message_LogNewMessage)(const upb_MiniTable* mini_table,
-                               const upb_Arena* arena);
+                                            const upb_Arena* arena);
 #endif
 
 // Inline version upb_Message_New(), for internal use.
 UPB_INLINE struct upb_Message* _upb_Message_New(const upb_MiniTable* m,
                                                 upb_Arena* a) {
 #ifdef UPB_TRACING_ENABLED
-  upb_Message_LogNewMessage(m, a);
+  UPB_PRIVATE(upb_Message_LogNewMessage)(m, a);
 #endif
   const int size = m->UPB_PRIVATE(size);
   struct upb_Message* msg = (struct upb_Message*)upb_Arena_Malloc(a, size);

upb/message/message.h

@@ -15,6 +15,7 @@
 #include <stddef.h>
 
 #include "upb/mem/arena.h"
+#include "upb/message/internal/message.h"
 #include "upb/message/internal/types.h"
 #include "upb/mini_table/message.h"
 
@@ -45,6 +46,18 @@ UPB_API void upb_Message_Freeze(upb_Message* msg, const upb_MiniTable* m);
 // Returns whether a message has been frozen.
 UPB_API_INLINE bool upb_Message_IsFrozen(const upb_Message* msg);
 
+#ifdef UPB_TRACING_ENABLED
+UPB_INLINE void upb_Message_SetNewMessageTraceHandler(
+    void (*newMessageTraceHandler)(const upb_MiniTable* mini_table,
+                                   const upb_Arena* arena)) {
+  UPB_PRIVATE(upb_Message_SetNewMessageTraceHandler)(newMessageTraceHandler);
+}
+UPB_INLINE void upb_Message_LogNewMessage(const upb_MiniTable* mini_table,
+                                          const upb_Arena* arena) {
+  UPB_PRIVATE(upb_Message_LogNewMessage)(mini_table, arena);
+}
+#endif
+
 #ifdef __cplusplus
 } /* extern "C" */
 #endif