python/free_threading/weak_map.c

// Protocol Buffers - Google's data interchange format
// Copyright 2023 Google LLC.  All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file or at
// https://developers.google.com/open-source/licenses/bsd

#include "python/free_threading/weak_map.h"

#include <assert.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>

#include "python/free_threading/mutex.h"
#include "upb/hash/common.h"
#include "upb/hash/int_table.h"
#include "upb/mem/alloc.h"
#include "upb/mem/arena.h"

// Must be last.
#include "upb/port/def.inc"

struct PyUpb_WeakMap {
  PyUpb_Mutex mutex;
  upb_inttable table;
  upb_Arena* arena;
};

struct PyUpb_WeakMapEntry {
  const void* key;
  PyObject* obj;
};

PyUpb_WeakMap* PyUpb_WeakMap_New(void) {
  upb_Arena* arena = upb_Arena_New();
  if (!arena) {
    PyErr_SetNone(PyExc_MemoryError);
    return NULL;
  }
  PyUpb_WeakMap* map = upb_Arena_Malloc(arena, sizeof(*map));
  if (!map) {
    PyErr_SetNone(PyExc_MemoryError);
    upb_Arena_Free(arena);
    return NULL;
  }
  map->arena = arena;
  PyUpb_Mutex_Init(&map->mutex);
  if (!upb_inttable_init(&map->table, map->arena)) {
    PyErr_SetNone(PyExc_MemoryError);
    PyUpb_Mutex_Destroy(&map->mutex);
    upb_Arena_Free(arena);
    return NULL;
  }
  return map;
}

void PyUpb_WeakMap_Free(PyUpb_WeakMap* map) {
  if (!map) return;
  PyUpb_Mutex_Destroy(&map->mutex);
  upb_Arena_Free(map->arena);
}

size_t PyUpb_WeakMap_Count(PyUpb_WeakMap* map) {
  PyUpb_Mutex_Lock(&map->mutex);
  size_t count = map->table.t.count;
  PyUpb_Mutex_Unlock(&map->mutex);
  return count;
}

// To give better entropy in the table key, we shift away low bits that are
// always zero.
static const int PyUpb_PtrShift = (sizeof(void*) == 4) ? 2 : 3;

uintptr_t PyUpb_WeakMap_GetKey(const void* key) {
  uintptr_t n = (uintptr_t)key;
  assert((n & ((1 << PyUpb_PtrShift) - 1)) == 0);
  return n >> PyUpb_PtrShift;
}

static bool PyUpb_WeakMap_TryIncRef(PyUpb_WeakMap* map, uintptr_t key,
                                    upb_value val) {
  PyObject* obj = upb_value_getptr(val);
#ifdef Py_GIL_DISABLED
  if (!PyUnstable_TryIncRef(obj)) {
    upb_inttable_remove(&map->table, key, NULL);
    return false;
  }
#else
  Py_INCREF(obj);
#endif
  return true;
}

static PyObject* PyUpb_WeakMap_GetLocked(PyUpb_WeakMap* map, const void* key) {
  upb_value val;
  if (upb_inttable_lookup(&map->table, PyUpb_WeakMap_GetKey(key), &val) &&
      PyUpb_WeakMap_TryIncRef(map, PyUpb_WeakMap_GetKey(key), val)) {
    return upb_value_getptr(val);
  } else {
    return NULL;
  }
}

PyObject* PyUpb_WeakMap_Get(PyUpb_WeakMap* map, const void* key) {
  PyUpb_Mutex_Lock(&map->mutex);
  PyObject* obj = PyUpb_WeakMap_GetLocked(map, key);
  PyUpb_Mutex_Unlock(&map->mutex);
  return obj;
}

bool PyUpb_WeakMap_Add(PyUpb_WeakMap* map, const void* key, PyObject** obj) {
#ifdef Py_GIL_DISABLED
  PyUnstable_EnableTryIncRef(*obj);
#endif

  PyObject* to_decref = NULL;
  PyObject* existing = NULL;
  PyUpb_Mutex_Lock(&map->mutex);

  bool ok = true;
  if ((existing = PyUpb_WeakMap_GetLocked(map, key)) != NULL) {
    to_decref = *obj;
    *obj = existing;
  } else {
    const uintptr_t k = PyUpb_WeakMap_GetKey(key);
    ok = upb_inttable_insert(&map->table, k, upb_value_ptr(*obj), map->arena);
  }

  PyUpb_Mutex_Unlock(&map->mutex);

  // Replacing an exception can retire user objects and reenter this map.
  if (!ok) PyErr_SetNone(PyExc_MemoryError);
  if (to_decref) {
    // This can trigger a dealloc which calls back into this WeakMap, so it must
    // be after the unlock.
    Py_DECREF(to_decref);
  }
  return ok;
}

bool PyUpb_WeakMap_UniqueAdd(PyUpb_WeakMap* map, const void* key,
                             PyObject* obj) {
  PyObject* dst = obj;
  bool ret = PyUpb_WeakMap_Add(map, key, &dst);
  // Function precondition violated: key was already in map.
  assert(dst == obj);
  return ret;
}

void PyUpb_WeakMap_Erase(PyUpb_WeakMap* map, const void* key, PyObject* obj) {
  PyUpb_WeakMap_EraseIfEqual(map, key, obj);
}

bool PyUpb_WeakMap_EraseIfEqual(PyUpb_WeakMap* map, const void* key,
                                PyObject* obj) {
  PyUpb_Mutex_Lock(&map->mutex);
  const uintptr_t k = PyUpb_WeakMap_GetKey(key);
  upb_value val;
  bool ret = false;
  if (upb_inttable_lookup(&map->table, k, &val) &&
      upb_value_getptr(val) == obj) {
    upb_inttable_remove(&map->table, k, NULL);
    ret = true;
  }
  PyUpb_Mutex_Unlock(&map->mutex);
  return ret;
}

bool PyUpb_WeakMapIter_Begin(PyUpb_WeakMapIter* iter, PyUpb_WeakMap* map) {
  iter->map = map;
  iter->entries = NULL;
  iter->count = 0;
  iter->index = 0;
  PyUpb_Mutex_Lock(&map->mutex);
  size_t cap = map->table.t.count;
  if (cap > SIZE_MAX / sizeof(*iter->entries)) {
    PyUpb_Mutex_Unlock(&map->mutex);
    PyErr_SetNone(PyExc_MemoryError);
    return false;
  }
  if (cap) {
    iter->entries = upb_gmalloc(cap * sizeof(*iter->entries));
    if (!iter->entries) {
      PyUpb_Mutex_Unlock(&map->mutex);
      PyErr_SetNone(PyExc_MemoryError);
      return false;
    }
  }
  intptr_t cursor = UPB_INTTABLE_BEGIN;
  uintptr_t key;
  upb_value val;
  while (upb_inttable_next(&map->table, &key, &val, &cursor)) {
    PyObject* obj = upb_value_getptr(val);
#ifdef Py_GIL_DISABLED
    if (!PyUnstable_TryIncRef(obj)) {
      // Object is being destroyed, remove it from the map and proceed to the
      // next entry.
      upb_inttable_removeiter(&map->table, &cursor);
      continue;
    }
#else
    Py_INCREF(obj);
#endif
    assert(iter->count < cap);
    iter->entries[iter->count++] =
        (PyUpb_WeakMapEntry){(const void*)(key << PyUpb_PtrShift), obj};
  }
  // Reification, exception construction and reference retirement can reenter
  // a cache. Never hold its non-recursive mutex while running the consumer.
  PyUpb_Mutex_Unlock(&map->mutex);
  return true;
}

void PyUpb_WeakMapIter_End(PyUpb_WeakMapIter* iter) {
  PyUpb_WeakMapEntry* entries = iter->entries;
  size_t count = iter->count;
  size_t start = iter->index > 0 ? iter->index - 1 : 0;
  iter->entries = NULL;
  iter->count = 0;
  iter->index = 0;
  for (size_t i = start; i < count; i++) {
    Py_XDECREF(entries[i].obj);
  }
  upb_gfree(entries);
}

bool PyUpb_WeakMapIter_Next(PyUpb_WeakMapIter* iter, const void** key,
                            PyObject** obj) {
  if (iter->index > 0) {
    Py_CLEAR(iter->entries[iter->index - 1].obj);
  }
  if (iter->index == iter->count) return false;
  const PyUpb_WeakMapEntry* entry = &iter->entries[iter->index++];
  *key = entry->key;
  *obj = entry->obj;
  return true;
}

void PyUpb_WeakMapIter_Delete(PyUpb_WeakMapIter* iter) {
  assert(iter->index > 0 && iter->index <= iter->count);
  const PyUpb_WeakMapEntry* entry = &iter->entries[iter->index - 1];
  // A concurrent or reentrant replacement must not be erased by a stale
  // snapshot entry. The snapshot owns obj until End().
  PyUpb_WeakMap_EraseIfEqual(iter->map, entry->key, entry->obj);
}