Implement memory accounting infrastructure

pclsync/pmem.c

@@ -1,10 +1,19 @@
 #include <sys/mman.h>
 #include <stdlib.h>
 #include <stdint.h>
+#include <string.h>
 
 #include "pcompiler.h"
 #include "psql.h"
 #include "pdbg.h"
+#include "pmem.h"
+
+static size_t subsystem_stats[6] = {0};
+
+typedef struct {
+  size_t size;
+  pmem_subsystem_t subsystem;
+} pmem_header_t;
 
 void *pmem_mmap(size_t size) {
   void *ret = mmap(NULL, size, PROT_READ | PROT_WRITE,
@@ -69,3 +78,61 @@ void *pmem_calloc_safe(size_t nmemb, size_t size) {
   }
   return calloc(nmemb, size);
 }
+
+void *pmem_malloc(pmem_subsystem_t subsystem, size_t size) {
+  pmem_header_t *hdr;
+  size_t total_size = sizeof(pmem_header_t) + size;
+  
+  hdr = (pmem_header_t *)malloc(total_size);
+  if (!hdr) {
+    return NULL;
+  }
+  
+  hdr->size = size;
+  hdr->subsystem = subsystem;
+  __atomic_add_fetch(&subsystem_stats[subsystem], size, __ATOMIC_RELAXED);
+  
+  return (void *)(hdr + 1);
+}
+
+void pmem_free(pmem_subsystem_t subsystem, void *ptr) {
+  pmem_header_t *hdr;
+  
+  if (!ptr) {
+    return;
+  }
+  
+  hdr = ((pmem_header_t *)ptr) - 1;
+  __atomic_sub_fetch(&subsystem_stats[hdr->subsystem], hdr->size, __ATOMIC_RELAXED);
+  free(hdr);
+}
+
+void *pmem_realloc(pmem_subsystem_t subsystem, void *ptr, size_t size) {
+  pmem_header_t *hdr, *new_hdr;
+  size_t old_size;
+  size_t total_size = sizeof(pmem_header_t) + size;
+  
+  if (!ptr) {
+    return pmem_malloc(subsystem, size);
+  }
+  
+  hdr = ((pmem_header_t *)ptr) - 1;
+  old_size = hdr->size;
+  
+  new_hdr = (pmem_header_t *)realloc(hdr, total_size);
+  if (!new_hdr) {
+    return NULL;
+  }
+  
+  __atomic_sub_fetch(&subsystem_stats[hdr->subsystem], old_size, __ATOMIC_RELAXED);
+  __atomic_add_fetch(&subsystem_stats[subsystem], size, __ATOMIC_RELAXED);
+  
+  new_hdr->size = size;
+  new_hdr->subsystem = subsystem;
+  
+  return (void *)(new_hdr + 1);
+}
+
+size_t pmem_get_stats(pmem_subsystem_t subsystem) {
+  return __atomic_load_n(&subsystem_stats[subsystem], __ATOMIC_RELAXED);
+}

pclsync/pmem.h

@@ -4,6 +4,15 @@
 #include <stddef.h>
 #include <stdint.h>
 
+typedef enum {
+  PMEM_SUBSYS_API,
+  PMEM_SUBSYS_CACHE,
+  PMEM_SUBSYS_UPLOAD,
+  PMEM_SUBSYS_DOWNLOAD,
+  PMEM_SUBSYS_SYNC,
+  PMEM_SUBSYS_OTHER
+} pmem_subsystem_t;
+
 void *pmem_mmap(size_t size);
 void *pmem_mmap_safe(size_t size);
 int pmem_munmap(void *ptr, size_t size);
@@ -11,5 +20,9 @@ int pmem_mlock(void *ptr, size_t size);
 int pmem_munlock(void *ptr, size_t size);
 void pmem_reset(void *ptr, size_t size);
 void *pmem_calloc_safe(size_t nmemb, size_t size);
+void *pmem_malloc(pmem_subsystem_t subsystem, size_t size);
+void pmem_free(pmem_subsystem_t subsystem, void *ptr);
+void *pmem_realloc(pmem_subsystem_t subsystem, void *ptr, size_t size);
+size_t pmem_get_stats(pmem_subsystem_t subsystem);
 
 #endif