hakmem/core/pool_tls_registry.c

#include "pool_tls_registry.h"
#include <pthread.h>
#include <stdlib.h>
#include <string.h>

typedef struct RegEntry {
  void* base;
  void* end;
  pid_t tid;
  int class_idx;
  struct RegEntry* next;
} RegEntry;

#define REG_BUCKETS 1024
static RegEntry* g_buckets[REG_BUCKETS];
static pthread_mutex_t g_locks[REG_BUCKETS];
static pthread_once_t g_init_once = PTHREAD_ONCE_INIT;

static void reg_init(void){
  for (int i=0;i<REG_BUCKETS;i++) pthread_mutex_init(&g_locks[i], NULL);
}

static inline uint64_t hash_ptr(void* p){
  uintptr_t x=(uintptr_t)p; x ^= x>>33; x*=0xff51afd7ed558ccdULL; x ^= x>>33; x*=0xc4ceb9fe1a85ec53ULL; x ^= x>>33; return x;
}

void pool_reg_register(void* base, size_t size, pid_t tid, int class_idx){
  pthread_once(&g_init_once, reg_init);
  void* end = (void*)((char*)base + size);
  uint64_t h = hash_ptr(base) & (REG_BUCKETS-1);
  pthread_mutex_lock(&g_locks[h]);
  RegEntry* e = (RegEntry*)malloc(sizeof(RegEntry));
  e->base = base; e->end = end; e->tid = tid; e->class_idx = class_idx; e->next = g_buckets[h];
  g_buckets[h] = e;
  pthread_mutex_unlock(&g_locks[h]);
}

void pool_reg_unregister(void* base, size_t size, pid_t tid){
  pthread_once(&g_init_once, reg_init);
  uint64_t h = hash_ptr(base) & (REG_BUCKETS-1);
  pthread_mutex_lock(&g_locks[h]);
  RegEntry** pp = &g_buckets[h];
  while (*pp){
    RegEntry* e = *pp;
    if (e->base == base && e->tid == tid){
      *pp = e->next; free(e); break;
    }
    pp = &e->next;
  }
  pthread_mutex_unlock(&g_locks[h]);
}

int pool_reg_lookup(void* ptr, pid_t* tid_out, int* class_idx_out){
  pthread_once(&g_init_once, reg_init);
  uint64_t h = hash_ptr(ptr) & (REG_BUCKETS-1);
  pthread_mutex_lock(&g_locks[h]);
  for (RegEntry* e = g_buckets[h]; e; e=e->next){
    if (ptr >= e->base && ptr < e->end){
      if (tid_out) *tid_out = e->tid;
      if (class_idx_out) *class_idx_out = e->class_idx;
      pthread_mutex_unlock(&g_locks[h]);
      return 1;
    }
  }
  pthread_mutex_unlock(&g_locks[h]);
  return 0;
}