feat: implement custom memory management hooks and example usage

sfxfs · sfxfs · commit a8ad7c33cc52 · 2025-10-01T10:05:19.000+08:00
diff --git a/example/CMakeLists.txt b/example/CMakeLists.txt
@@ -6,3 +6,4 @@ add_subdirectory(invoke_a_func)
 add_subdirectory(notif_req)
 add_subdirectory(rpc_with_params)
 add_subdirectory(arg_to_function)
+add_subdirectory(memory_hooks)
diff --git a/example/memory_hooks/CMakeLists.txt b/example/memory_hooks/CMakeLists.txt
@@ -0,0 +1,5 @@
+project(mjsonrpc-example-mem-hooks)
+
+add_executable(${PROJECT_NAME} memory_hooks.c)
+
+target_link_libraries(${PROJECT_NAME} mjsonrpc)
diff --git a/example/memory_hooks/memory_hooks.c b/example/memory_hooks/memory_hooks.c
@@ -0,0 +1,137 @@
+#include "mjsonrpc.h"
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+// Statistics for tracking memory usage
+static size_t malloc_count = 0;
+static size_t free_count = 0;
+static size_t strdup_count = 0;
+
+// Custom memory allocation functions with tracking
+void* custom_malloc(size_t size)
+{
+    malloc_count++;
+    printf("[CUSTOM MALLOC] Allocating %zu bytes (call #%zu)\n", size, malloc_count);
+    return malloc(size);
+}
+
+void custom_free(void* ptr)
+{
+    if (ptr != NULL) {
+        free_count++;
+        printf("[CUSTOM FREE] Freeing memory (call #%zu)\n", free_count);
+    }
+    free(ptr);
+}
+
+char* custom_strdup(const char* str)
+{
+    strdup_count++;
+    printf("[CUSTOM STRDUP] Duplicating string: '%.20s%s' (call #%zu)\n", 
+           str, strlen(str) > 20 ? "..." : "", strdup_count);
+    
+    if (str == NULL) {
+        return NULL;
+    }
+    
+    size_t len = strlen(str) + 1;
+    char* dup = custom_malloc(len);
+    if (dup != NULL) {
+        memcpy(dup, str, len);
+    }
+    return dup;
+}
+
+// Simple RPC method for testing
+cJSON* hello_method(mjrpc_func_ctx_t* context, cJSON* params, cJSON* id)
+{
+    (void) context; // Unused parameter
+    (void) id;      // Unused parameter
+    const char* name = "World";
+    
+    if (params && cJSON_IsObject(params)) {
+        cJSON* name_item = cJSON_GetObjectItem(params, "name");
+        if (cJSON_IsString(name_item)) {
+            name = name_item->valuestring;
+        }
+    }
+    
+    // This will trigger memory allocation for the string
+    char greeting[256];
+    snprintf(greeting, sizeof(greeting), "Hello, %s!", name);
+    
+    return cJSON_CreateString(greeting);
+}
+
+int main(void)
+{
+    printf("=== mjsonrpc Memory Hooks Example ===\n\n");
+    
+    // Step 1: Set custom memory hooks
+    printf("1. Setting custom memory hooks...\n");
+    int result = mjrpc_set_memory_hooks(custom_malloc, custom_free, custom_strdup);
+    if (result != MJRPC_RET_OK) {
+        printf("Failed to set memory hooks!\n");
+        return 1;
+    }
+    printf("Custom memory hooks set successfully!\n\n");
+    
+    // Step 2: Create handle (this will use custom memory functions)
+    printf("2. Creating mjsonrpc handle...\n");
+    mjrpc_handle_t* handle = mjrpc_create_handle(0);
+    if (!handle) {
+        printf("Failed to create handle!\n");
+        return 1;
+    }
+    printf("Handle created successfully!\n\n");
+    
+    // Step 3: Add a method (this will use custom strdup)
+    printf("3. Adding 'hello' method...\n");
+    result = mjrpc_add_method(handle, hello_method, "hello", NULL);
+    if (result != MJRPC_RET_OK) {
+        printf("Failed to add method!\n");
+        mjrpc_destroy_handle(handle);
+        return 1;
+    }
+    printf("Method added successfully!\n\n");
+    
+    // Step 4: Process a request (this may trigger more memory allocations)
+    printf("4. Processing JSON-RPC request...\n");
+    const char* request_str = "{\"jsonrpc\":\"2.0\",\"method\":\"hello\",\"params\":{\"name\":\"Alice\"},\"id\":1}";
+    int ret_code;
+    char* response_str = mjrpc_process_str(handle, request_str, &ret_code);
+    
+    if (response_str) {
+        printf("Response: %s\n", response_str);
+        free(response_str);  // This uses standard free since mjrpc_process_str uses standard memory allocation internally
+    } else {
+        printf("Failed to process request (return code: %d)\n", ret_code);
+    }
+    printf("\n");
+    
+    // Step 5: Clean up (this will use custom free functions)
+    printf("5. Cleaning up...\n");
+    mjrpc_destroy_handle(handle);
+    printf("Cleanup completed!\n\n");
+    
+    // Step 6: Reset to default functions
+    printf("6. Resetting to default memory functions...\n");
+    result = mjrpc_set_memory_hooks(NULL, NULL, NULL);
+    if (result != MJRPC_RET_OK) {
+        printf("Failed to reset memory hooks!\n");
+        return 1;
+    }
+    printf("Memory hooks reset to defaults!\n\n");
+    
+    // Print statistics
+    printf("=== Memory Usage Statistics ===\n");
+    printf("Custom malloc calls: %zu\n", malloc_count);
+    printf("Custom free calls:   %zu\n", free_count);
+    printf("Custom strdup calls: %zu\n", strdup_count);
+    printf("\n");
+    
+    printf("Example completed successfully!\n");
+    return 0;
+}
diff --git a/src/mjsonrpc.c b/src/mjsonrpc.c
@@ -28,6 +28,13 @@
 #include <string.h>
 #include <stdbool.h>
 
+/*--- memory management hooks ---*/
+
+// Global memory function pointers, default to standard library functions
+static mjrpc_malloc_func g_mjrpc_malloc = malloc;
+static mjrpc_free_func g_mjrpc_free = free;
+static mjrpc_strdup_func g_mjrpc_strdup = strdup;
+
 /*--- utility ---*/
 
 enum method_state
@@ -54,7 +61,8 @@ static int resize(mjrpc_handle_t* handle)
 
     handle->capacity *= 2;
     handle->size = 0;
-    handle->methods = (struct mjrpc_method*) calloc(handle->capacity, sizeof(struct mjrpc_method));
+    handle->methods =
+        (struct mjrpc_method*) g_mjrpc_malloc(handle->capacity * sizeof(struct mjrpc_method));
     if (handle->methods == NULL)
     {
         handle->capacity = old_capacity;
@@ -67,10 +75,10 @@ static int resize(mjrpc_handle_t* handle)
         if (old_methods[i].state == OCCUPIED)
         {
             mjrpc_add_method(handle, old_methods[i].func, old_methods[i].name, old_methods[i].arg);
-            free(old_methods[i].name);
+            g_mjrpc_free(old_methods[i].name);
         }
     }
-    free(old_methods);
+    g_mjrpc_free(old_methods);
     return MJRPC_RET_OK;
 }
 
@@ -107,8 +115,8 @@ static cJSON* invoke_callback(const mjrpc_handle_t* handle, const char* method_n
     ctx.error_message = NULL;
     if (!method_get((mjrpc_handle_t*) handle, method_name, &func, &arg) || !func)
     {
-        return mjrpc_response_error(JSON_RPC_CODE_METHOD_NOT_FOUND, strdup("Method not found."),
-                                    id);
+        return mjrpc_response_error(JSON_RPC_CODE_METHOD_NOT_FOUND,
+                                    g_mjrpc_strdup("Method not found."), id);
     }
     ctx.data = arg;
     returned = func(&ctx, params, id);
@@ -144,9 +152,9 @@ static cJSON* rpc_handle_obj_req(const mjrpc_handle_t* handle, const cJSON* requ
         const cJSON* version = cJSON_GetObjectItem(request, "jsonrpc");
         if (version == NULL || version->type != cJSON_String ||
             strcmp("2.0", version->valuestring) != 0)
-            return mjrpc_response_error(JSON_RPC_CODE_INVALID_REQUEST,
-                                        strdup("Invalid request received: JSONRPC version error."),
-                                        id_copy);
+            return mjrpc_response_error(
+                JSON_RPC_CODE_INVALID_REQUEST,
+                g_mjrpc_strdup("Invalid request received: JSONRPC version error."), id_copy);
 
         const cJSON* method = cJSON_GetObjectItem(request, "method");
         if (method != NULL && method->type == cJSON_String)
@@ -156,12 +164,12 @@ static cJSON* rpc_handle_obj_req(const mjrpc_handle_t* handle, const cJSON* requ
             return invoke_callback(handle, method->valuestring, params, id_copy);
         }
         return mjrpc_response_error(JSON_RPC_CODE_INVALID_REQUEST,
-                                    strdup("Invalid request received: No 'method' member."),
+                                    g_mjrpc_strdup("Invalid request received: No 'method' member."),
                                     id_copy);
     }
     // Invalid id type
     return mjrpc_response_error(JSON_RPC_CODE_INVALID_REQUEST,
-                                strdup("Invalid request received: 'id' member type error."),
+                                g_mjrpc_strdup("Invalid request received: 'id' member type error."),
                                 cJSON_CreateNull());
 }
 
@@ -243,7 +251,7 @@ cJSON* mjrpc_response_error(int code, char* message, cJSON* id)
     if (id == NULL)
     {
         if (message)
-            free(message);
+            g_mjrpc_free(message);
         return NULL;
     }
 
@@ -266,7 +274,7 @@ cJSON* mjrpc_response_error(int code, char* message, cJSON* id)
     if (message)
     {
         cJSON_AddStringToObject(error_root, "message", message);
-        free(message);
+        g_mjrpc_free(message);
     }
     else
     {
@@ -286,15 +294,16 @@ mjrpc_handle_t* mjrpc_create_handle(size_t initial_capacity)
 {
     if (initial_capacity == 0)
         initial_capacity = DEFAULT_INITIAL_CAPACITY;
-    mjrpc_handle_t* handle = malloc(sizeof(mjrpc_handle_t));
+    mjrpc_handle_t* handle = g_mjrpc_malloc(sizeof(mjrpc_handle_t));
     if (handle == NULL)
         return NULL;
     handle->capacity = initial_capacity;
     handle->size = 0;
-    handle->methods = (struct mjrpc_method*) calloc(handle->capacity, sizeof(struct mjrpc_method));
+    handle->methods =
+        (struct mjrpc_method*) g_mjrpc_malloc(handle->capacity * sizeof(struct mjrpc_method));
     if (handle->methods == NULL)
     {
-        free(handle);
+        g_mjrpc_free(handle);
         return NULL;
     }
     return handle;
@@ -308,13 +317,13 @@ int mjrpc_destroy_handle(mjrpc_handle_t* handle)
     {
         if (handle->methods[i].state == OCCUPIED)
         {
-            free(handle->methods[i].name);
+            g_mjrpc_free(handle->methods[i].name);
             if (handle->methods[i].arg != NULL)
-                free(handle->methods[i].arg);
+                g_mjrpc_free(handle->methods[i].arg);
         }
     }
-    free(handle->methods);
-    free(handle);
+    g_mjrpc_free(handle->methods);
+    g_mjrpc_free(handle);
     return MJRPC_RET_OK;
 }
 
@@ -345,7 +354,7 @@ int mjrpc_add_method(mjrpc_handle_t* handle, mjrpc_func function_pointer, const
         index = (index + probe_count * probe_count) % handle->capacity;
     }
 
-    handle->methods[index].name = strdup(method_name);
+    handle->methods[index].name = g_mjrpc_strdup(method_name);
     handle->methods[index].func = function_pointer;
     handle->methods[index].arg = arg2func;
     handle->methods[index].state = OCCUPIED;
@@ -365,9 +374,9 @@ int mjrpc_del_method(mjrpc_handle_t* handle, const char* name)
         if (handle->methods[index].state == OCCUPIED &&
             strcmp(handle->methods[index].name, name) == 0)
         {
-            free(handle->methods[index].name);
+            g_mjrpc_free(handle->methods[index].name);
             if (handle->methods[index].arg != NULL)
-                free(handle->methods[index].arg);
+                g_mjrpc_free(handle->methods[index].arg);
             handle->methods[index].state = DELETED;
             handle->size--;
             return MJRPC_RET_OK;
@@ -410,7 +419,8 @@ cJSON* mjrpc_process_cjson(mjrpc_handle_t* handle, const cJSON* request_cjson, i
             *ret_code = ret;
         return mjrpc_response_error(
             JSON_RPC_CODE_PARSE_ERROR,
-            strdup("Invalid request received: Not a JSON formatted request."), cJSON_CreateNull());
+            g_mjrpc_strdup("Invalid request received: Not a JSON formatted request."),
+            cJSON_CreateNull());
     }
 
     cJSON* cjson_return = NULL;
@@ -422,7 +432,7 @@ cJSON* mjrpc_process_cjson(mjrpc_handle_t* handle, const cJSON* request_cjson, i
             ret = MJRPC_RET_ERROR_EMPTY_REQUEST;
             cjson_return = mjrpc_response_error(
                 JSON_RPC_CODE_PARSE_ERROR,
-                strdup("Invalid request received: Empty JSON array."), cJSON_CreateNull());
+                g_mjrpc_strdup("Invalid request received: Empty JSON array."), cJSON_CreateNull());
         }
         else
         {
@@ -441,7 +451,7 @@ cJSON* mjrpc_process_cjson(mjrpc_handle_t* handle, const cJSON* request_cjson, i
             ret = MJRPC_RET_ERROR_EMPTY_REQUEST;
             cjson_return = mjrpc_response_error(
                 JSON_RPC_CODE_PARSE_ERROR,
-                strdup("Invalid request received: Empty JSON object."), cJSON_CreateNull());
+                g_mjrpc_strdup("Invalid request received: Empty JSON object."), cJSON_CreateNull());
         }
         else
         {
@@ -456,10 +466,37 @@ cJSON* mjrpc_process_cjson(mjrpc_handle_t* handle, const cJSON* request_cjson, i
     {
         cjson_return = mjrpc_response_error(
             JSON_RPC_CODE_PARSE_ERROR,
-            strdup("Invalid request received: Not a JSON object or array."), cJSON_CreateNull());
+            g_mjrpc_strdup("Invalid request received: Not a JSON object or array."),
+            cJSON_CreateNull());
         ret = MJRPC_RET_ERROR_NOT_OBJ_ARY;
     }
     if (ret_code)
         *ret_code = ret;
     return cjson_return;
 }
+
+int mjrpc_set_memory_hooks(mjrpc_malloc_func malloc_func, mjrpc_free_func free_func,
+                           mjrpc_strdup_func strdup_func)
+{
+    // If all parameters are NULL, reset to default functions
+    if (malloc_func == NULL && free_func == NULL && strdup_func == NULL)
+    {
+        g_mjrpc_malloc = malloc;
+        g_mjrpc_free = free;
+        g_mjrpc_strdup = strdup;
+        return MJRPC_RET_OK;
+    }
+
+    // If any parameter is not NULL, all must be provided
+    if (malloc_func == NULL || free_func == NULL || strdup_func == NULL)
+    {
+        return MJRPC_RET_ERROR_INVALID_PARAM;
+    }
+
+    // Set custom functions
+    g_mjrpc_malloc = malloc_func;
+    g_mjrpc_free = free_func;
+    g_mjrpc_strdup = strdup_func;
+
+    return MJRPC_RET_OK;
+}
diff --git a/src/mjsonrpc.h b/src/mjsonrpc.h
