From eb38a2a608183f28ac5e296b3e1cec5a2ad3168f Mon Sep 17 00:00:00 2001 From: Joshua Haberman Date: Sun, 5 Dec 2021 23:54:34 -0800 Subject: [PATCH] Some more cleanups before sending for PR. --- python/descriptor_pool.c | 25 ++++-- python/message.c | 188 ++++++++++++++++++++++----------------- python/message.h | 35 +++++++- python/protobuf.c | 119 ++++++++++++++++++++----- python/protobuf.h | 20 ++--- 5 files changed, 258 insertions(+), 129 deletions(-) diff --git a/python/descriptor_pool.c b/python/descriptor_pool.c index b506b404a4..8e95bdf526 100644 --- a/python/descriptor_pool.c +++ b/python/descriptor_pool.c @@ -46,17 +46,23 @@ PyObject* PyUpb_DescriptorPool_GetDefaultPool() { return s->default_pool; } -static PyObject* PyUpb_DescriptorPool_DoCreate(PyTypeObject* type, - PyObject* db) { - PyUpb_DescriptorPool* pool = PyObject_GC_New(PyUpb_DescriptorPool, type); +static PyObject* PyUpb_DescriptorPool_DoCreateWithCache(PyTypeObject* type, + PyObject* db, + PyUpb_WeakMap *obj_cache) { + PyUpb_DescriptorPool* pool = (void*)PyType_GenericAlloc(type, 0); pool->symtab = upb_symtab_new(); pool->db = db; Py_XINCREF(pool->db); - PyObject_GC_Track(&pool->ob_base); - PyUpb_ObjCache_Add(pool->symtab, &pool->ob_base); + PyUpb_WeakMap_Add(obj_cache, pool->symtab, &pool->ob_base); return &pool->ob_base; } +static PyObject* PyUpb_DescriptorPool_DoCreate(PyTypeObject* type, + PyObject* db) { + return PyUpb_DescriptorPool_DoCreateWithCache(type, db, + PyUpb_ObjCache_Instance()); +} + upb_symtab* PyUpb_DescriptorPool_GetSymtab(PyObject* pool) { return ((PyUpb_DescriptorPool*)pool)->symtab; } @@ -235,13 +241,14 @@ static PyType_Spec PyUpb_DescriptorPool_Spec = { // ----------------------------------------------------------------------------- bool PyUpb_InitDescriptorPool(PyObject* m) { + PyUpb_ModuleState *state = PyUpb_ModuleState_GetFromModule(m); PyTypeObject* descriptor_pool_type = AddObject(m, "DescriptorPool", &PyUpb_DescriptorPool_Spec); if (!descriptor_pool_type) return false; - PyObject* default_pool = - PyUpb_DescriptorPool_DoCreate(descriptor_pool_type, NULL); - return default_pool && - PyModule_AddObject(m, "default_pool", default_pool) == 0; + state->default_pool = PyUpb_DescriptorPool_DoCreateWithCache( + descriptor_pool_type, NULL, state->obj_cache); + return state->default_pool && + PyModule_AddObject(m, "default_pool", state->default_pool) == 0; } diff --git a/python/message.c b/python/message.c index 2db3f6a6f1..87c699dbdb 100644 --- a/python/message.c +++ b/python/message.c @@ -34,7 +34,8 @@ #include "upb/text_encode.h" #include "upb/util/required_fields.h" -const upb_msgdef* PyUpb_MessageMeta_GetMsgdef(PyObject* cls); +static const upb_msgdef* PyUpb_MessageMeta_GetMsgdef(PyObject* cls); +static PyObject* PyUpb_MessageMeta_GetAttr(PyObject* self, PyObject* name); // ----------------------------------------------------------------------------- // CPythonBits @@ -60,8 +61,6 @@ typedef struct { // than the version we are dynamically linked against. Here we want the // version that is actually running in this process. long python_version_hex; // PY_VERSION_HEX - - PyObject* descriptor; } PyUpb_CPythonBits; // A global containing the values for this process. @@ -214,6 +213,15 @@ static PyObject* PyUpb_CMessage_New(PyObject* cls, PyObject* unused_args, return ret; } +/* + * PyUpb_CMessage_LookupName() + * + * Tries to find a field or oneof named `py_name` in the message object `self`. + * The user can pass `f` and/or `o` to indicate whether a field or a oneof name + * is expected. If the name is found and is has an expected type, the function + * sets `*f` or `*o` respectively and returns true. Otherwise returns false + * and sets an exception of type `exc_type` if provided. + */ static bool PyUpb_CMessage_LookupName(PyUpb_CMessage* self, PyObject* py_name, const upb_fielddef** f, const upb_oneofdef** o, @@ -247,8 +255,23 @@ static bool PyUpb_CMessage_LookupName(PyUpb_CMessage* self, PyObject* py_name, return true; } -int PyUpb_CMessage_InitMapAttributes(PyObject* map, PyObject* value, - const upb_fielddef* f) { +static bool PyUpb_CMessage_InitMessageMapEntry(PyObject* dst, PyObject* src) { + if (!src || !dst) return false; + + // Currently we are doing Clear()+MergeFrom(). Replace with CopyFrom() once that + // is implemented. + PyObject* ok = PyObject_CallMethod(dst, "Clear", NULL); + if (!ok) return false; + Py_DECREF(ok); + ok = PyObject_CallMethod(dst, "MergeFrom", "O", src); + if (!ok) return false; + Py_DECREF(ok); + + return true; +} + +static int PyUpb_CMessage_InitMapAttributes(PyObject* map, PyObject* value, + const upb_fielddef* f) { const upb_msgdef* entry_m = upb_fielddef_msgsubdef(f); const upb_fielddef* val_f = upb_msgdef_field(entry_m, 1); if (upb_fielddef_issubmsg(val_f)) { @@ -263,22 +286,10 @@ int PyUpb_CMessage_InitMapAttributes(PyObject* map, PyObject* value, PyObject* src = PyObject_GetItem(value, item); PyObject* dst = PyObject_GetItem(map, item); Py_DECREF(item); - - if (!src || !dst) { - Py_XDECREF(src); - Py_XDECREF(dst); - Py_DECREF(iter); - return -1; - } - - PyObject* ok = PyObject_CallMethod(dst, "Clear", NULL); - assert(ok); - Py_DECREF(ok); - PyObject* ok2 = PyObject_CallMethod(dst, "MergeFrom", "O", src); - Py_DECREF(src); - Py_DECREF(dst); - Py_XDECREF(ok2); - if (!ok2) { + bool ok = PyUpb_CMessage_InitMessageMapEntry(dst, src); + Py_XDECREF(src); + Py_XDECREF(dst); + if (!ok) { Py_DECREF(iter); return -1; } @@ -473,6 +484,21 @@ void PyUpb_CMessage_AssureWritable(PyUpb_CMessage* self) { self->version++; } +/* + * PyUpb_CMessage_SyncSubobjs() + * + * This operation must be invoked whenever the underlying upb_msg has been + * mutated directly in C. This will attach any newly-present field data + * to previously returned "empty" wrapper objects. + * + * For example: + * foo = FooMessage() + * sub = foo.submsg # Empty, unset sub-message + * + * # SyncSubobjs() is required to connect our existing 'sub' wrapper to the + * # newly created foo.submsg data in C. + * foo.MergeFrom(FooMessage(submsg={})) + */ static void PyUpb_CMessage_SyncSubobjs(PyUpb_CMessage* self) { PyUpb_WeakMap* subobj_map = self->unset_subobj_map; upb_msg* msg = PyUpb_CMessage_GetMsg(self); @@ -511,6 +537,9 @@ static void PyUpb_CMessage_SyncSubobjs(PyUpb_CMessage* self) { } Py_DECREF(&self->ob_base); + + // TODO(haberman): present fields need to be iterated too if they can reach + // a WeakMap. } static PyObject* PyUpb_CMessage_ToString(PyUpb_CMessage* self) { @@ -596,7 +625,8 @@ PyObject* PyUpb_CMessage_Get(upb_msg* u_msg, const upb_msgdef* m, if (!ret) { PyObject* cls = PyUpb_Descriptor_GetClass(m); - // https://bugs.python.org/issue35810 + // It is not safe to use PyObject_{,GC}_New() due to: + // https://bugs.python.org/issue35810 PyUpb_CMessage* py_msg = (void*)PyType_GenericAlloc((PyTypeObject*)cls, 0); py_msg->arena = arena; py_msg->def = (uintptr_t)m; @@ -613,6 +643,17 @@ PyObject* PyUpb_CMessage_Get(upb_msg* u_msg, const upb_msgdef* m, return ret; } +/* + * PyUpb_CMessage_GetFieldValue() + * + * Implements the equivalent of getattr(msg, field), once `field` has + * already been resolved to a `upb_fielddef*`. + * + * This may involve constructing a wrapper object for the given field, or + * returning one that was previously constructed. If the field is not actually + * set, the wrapper object will be an "unset" object that is not actually + * connected to any C data. + */ PyObject* PyUpb_CMessage_GetFieldValue(PyObject* _self, const upb_fielddef* field) { PyUpb_CMessage* self = (PyUpb_CMessage*)_self; @@ -709,8 +750,6 @@ int PyUpb_CMessage_GetVersion(PyObject* _self) { return self->version; } -static PyObject* PyUpb_MessageMeta_GetAttr(PyObject* self, PyObject* name); - /* * PyUpb_CMessage_GetAttr() * @@ -946,7 +985,8 @@ static PyObject* PyUpb_CMessage_ClearField(PyUpb_CMessage* self, obj = PyUpb_ObjCache_Get(msgval.map_val); } if (obj) { - PyUpb_MapContainer_Invalidate(obj); + // TODO(haberman) add when maps are checked in. + // PyUpb_MapContainer_Invalidate(obj); Py_DECREF(obj); } } @@ -1117,6 +1157,7 @@ static PyGetSetDef PyUpb_CMessage_Getters[] = { {NULL}}; static PyMethodDef PyUpb_CMessage_Methods[] = { + // TODO(https://github.com/protocolbuffers/upb/issues/459) //{ "__deepcopy__", (PyCFunction)DeepCopy, METH_VARARGS, // "Makes a deep copy of the class." }, //{ "__unicode__", (PyCFunction)ToUnicode, METH_NOARGS, @@ -1129,6 +1170,7 @@ static PyMethodDef PyUpb_CMessage_Methods[] = { "Clears a message field."}, {"ClearField", (PyCFunction)PyUpb_CMessage_ClearField, METH_O, "Clears a message field."}, + // TODO(https://github.com/protocolbuffers/upb/issues/459) //{ "CopyFrom", (PyCFunction)CopyFrom, METH_O, // "Copies a protocol message into the current message." }, {"DiscardUnknownFields", (PyCFunction)PyUpb_CMessage_DiscardUnknownFields, @@ -1142,6 +1184,7 @@ static PyMethodDef PyUpb_CMessage_Methods[] = { "Checks if a message field is set."}, {"HasField", PyUpb_CMessage_HasField, METH_O, "Checks if a message field is set."}, + // TODO(https://github.com/protocolbuffers/upb/issues/459) //{ "IsInitialized", (PyCFunction)IsInitialized, METH_VARARGS, // "Checks if all required fields of a protocol message are set." }, {"ListFields", PyUpb_CMessage_ListFields, METH_NOARGS, @@ -1152,6 +1195,7 @@ static PyMethodDef PyUpb_CMessage_Methods[] = { "Merges a serialized message into the current message."}, {"ParseFromString", PyUpb_CMessage_ParseFromString, METH_O, "Parses a serialized message into the current message."}, + // TODO(https://github.com/protocolbuffers/upb/issues/459) //{ "RegisterExtension", (PyCFunction)RegisterExtension, METH_O | // METH_CLASS, // "Registers an extension with the current message." }, @@ -1169,6 +1213,7 @@ static PyMethodDef PyUpb_CMessage_Methods[] = { {"WhichOneof", PyUpb_CMessage_WhichOneof, METH_O, "Returns the name of the field set inside a oneof, " "or None if no field is set."}, + // TODO(https://github.com/protocolbuffers/upb/issues/459) //{ "_CheckCalledFromGeneratedFile", //(PyCFunction)_CheckCalledFromGeneratedFile, // METH_NOARGS | METH_STATIC, @@ -1202,12 +1247,24 @@ PyType_Spec PyUpb_CMessage_Spec = { // MessageMeta // ----------------------------------------------------------------------------- +// MessageMeta is the metaclass for message objects. The generated code uses it +// to construct message classes, ie. +// +// FooMessage = _message.MessageMeta('FooMessage', (_message.Message), {...}) +// +// (This is not quite true: at the moment the Python library subclasses +// MessageMeta, and uses that subclass as the metaclass. There is a TODO below +// to simplify this, so that the illustration above is indeed accurate). + typedef struct { const upb_msglayout* layout; PyObject* py_message_descriptor; } PyUpb_MessageMeta; -PyUpb_MessageMeta* PyUpb_GetMessageMeta(PyObject* cls) { +// The PyUpb_MessageMeta struct is trailing data tacked onto the end of +// MessageMeta instances. This means that we get our instances of this struct +// by adding the appropriate number of bytes. +static PyUpb_MessageMeta* PyUpb_GetMessageMeta(PyObject* cls) { #ifndef NDEBUG PyUpb_ModuleState* state = PyUpb_ModuleState_MaybeGet(); assert(!state || cls->ob_type == state->message_meta_type); @@ -1215,6 +1272,11 @@ PyUpb_MessageMeta* PyUpb_GetMessageMeta(PyObject* cls) { return (PyUpb_MessageMeta*)((char*)cls + cpython_bits.type_basicsize); } +static const upb_msgdef* PyUpb_MessageMeta_GetMsgdef(PyObject* cls) { + PyUpb_MessageMeta* self = PyUpb_GetMessageMeta(cls); + return PyUpb_Descriptor_GetDef(self->py_message_descriptor); +} + PyObject* PyUpb_MessageMeta_ModuleQualifiedName(const upb_msgdef* m) { const upb_filedef* file = upb_msgdef_file(m); const char* filename = upb_filedef_name(file); @@ -1326,46 +1388,6 @@ static void PyUpb_MessageMeta_Dealloc(PyObject* self) { PyUpb_Dealloc(self); } -PyObject* PyUpb_MessageMeta_CreateClass(PyObject* py_descriptor) { - const upb_msgdef* m = PyUpb_Descriptor_GetDef(py_descriptor); - PyObject* dict = PyDict_New(); - PyObject* ret = PyUpb_MessageMeta_DoCreateClass(py_descriptor, - upb_msgdef_fullname(m), dict); - Py_DECREF(dict); - return ret; -} - -/* -// Compute some class attributes on the fly: -// - All the _FIELD_NUMBER attributes, for all fields and nested extensions. -// Returns a new reference, or NULL with an exception set. -static PyObject* GetClassAttribute(CMessageClass *self, PyObject* name) { - char* attr; - Py_ssize_t attr_size; - static const char kSuffix[] = "_FIELD_NUMBER"; - if (PyString_AsStringAndSize(name, &attr, &attr_size) >= 0 && - HasSuffixString(StringPiece(attr, attr_size), kSuffix)) { - std::string field_name(attr, attr_size - sizeof(kSuffix) + 1); - LowerString(&field_name); - - // Try to find a field with the given name, without the suffix. - const FieldDescriptor* field = - self->message_descriptor->FindFieldByLowercaseName(field_name); - if (!field) { - // Search nested extensions as well. - field = - self->message_descriptor->FindExtensionByLowercaseName(field_name); - } - if (field) { - return PyInt_FromLong(field->number()); - } - } - PyErr_SetObject(PyExc_AttributeError, name); - return NULL; -} - -*/ - static PyObject* PyUpb_MessageMeta_GetDynamicAttr(PyObject* self, PyObject* name) { const char* name_buf = PyUpb_GetStrData(name); @@ -1395,6 +1417,9 @@ static PyObject* PyUpb_MessageMeta_GetDynamicAttr(PyObject* self, ret = PyUpb_FieldDescriptor_Get(ext); } + // TODO(haberman): *_FIELD_NUMBER attributes? Haven't seen any failing + // tests for these yet. + Py_DECREF(py_key); return ret; @@ -1437,29 +1462,26 @@ static PyType_Spec PyUpb_MessageMeta_Spec = { PyUpb_MessageMeta_Slots, }; -bool PyUpb_GetTypeFuncs(void) { +static PyObject* PyUpb_MessageMeta_CreateType() { + PyObject* bases = Py_BuildValue("(O)", &PyType_Type); + if (!bases) return NULL; + PyUpb_MessageMeta_Spec.basicsize = + cpython_bits.type_basicsize + sizeof(PyUpb_MessageMeta); + PyObject* type = PyType_FromSpecWithBases(&PyUpb_MessageMeta_Spec, bases); + Py_DECREF(bases); + return type; } bool PyUpb_InitMessage(PyObject* m) { - if (!PyUpb_GetTypeFuncs()) return false; + if (!PyUpb_CPythonBits_Init(&cpython_bits)) return false; + PyObject* message_meta_type = PyUpb_MessageMeta_CreateType(); PyUpb_ModuleState* state = PyUpb_ModuleState_GetFromModule(m); state->cmessage_type = PyUpb_AddClass(m, &PyUpb_CMessage_Spec); + state->message_meta_type = (PyTypeObject*)message_meta_type; - if (!state->cmessage_type) return false; + if (!state->cmessage_type || !state->message_meta_type) return false; + if (PyModule_AddObject(m, "MessageMeta", message_meta_type)) return false; - PyObject* bases = Py_BuildValue("(O)", &PyType_Type); - PyUpb_MessageMeta_Spec.basicsize = - cpython_bits.type_basicsize + sizeof(PyUpb_MessageMeta); - PyObject* type = PyType_FromSpecWithBases(&PyUpb_MessageMeta_Spec, bases); - Py_DECREF(bases); - - if (!type) return false; - - if (PyModule_AddObject(m, "MessageMeta", type) == 0) { - state->message_meta_type = (PyTypeObject*)type; - return true; - } else { - return false; - } + return true; } diff --git a/python/message.h b/python/message.h index e28f24f4f0..e10f83251c 100644 --- a/python/message.h +++ b/python/message.h @@ -33,28 +33,61 @@ #include "python/protobuf.h" #include "upb/reflection.h" -PyObject* PyUpb_MessageMeta_CreateClass(PyObject* py_descriptor); +// Removes the wrapper object for this field from the unset subobject cache. void PyUpb_CMessage_CacheDelete(PyObject* _self, const upb_fielddef* f); + +// Sets the field value for `f` to `subobj`, evicting the wrapper object from +// the "unset subobject" cache now that real data exists for it. The caller +// must also update the wrapper associated with `f` to point to `subobj` also. void PyUpb_CMessage_SetConcreteSubobj(PyObject* _self, const upb_fielddef* f, upb_msgval subobj); + +// Gets a Python wrapper object for message `u_msg` of type `m`, returning a +// cached wrapper if one was previously created. If a new object is created, +// it will reference `arena`, which must own `u_msg`. PyObject* PyUpb_CMessage_Get(upb_msg* u_msg, const upb_msgdef* m, PyObject* arena); + +// Call to check whether the given Python object is a message. If not, returns +// false and sets a TypeError exception. bool PyUpb_CMessage_Check(PyObject* self); + +// Gets the upb_msg* for this message object if the message is set/writable. +// Otherwise returns NULL. upb_msg* PyUpb_CMessage_GetIfWritable(PyObject* _self); + +// Returns the `upb_msgdef` for a given CMessage. const upb_msgdef* PyUpb_CMessage_GetMsgdef(PyObject* self); + +// Functions that match the corresponding methods on the message object. PyObject* PyUpb_CMessage_MergeFrom(PyObject* self, PyObject* arg); PyObject* PyUpb_CMessage_MergeFromString(PyObject* self, PyObject* arg); PyObject* PyUpb_CMessage_SerializeToString(PyObject* self, PyObject* args, PyObject* kwargs); + +// Sets fields of the message according to the attribuges in `kwargs`. int PyUpb_CMessage_InitAttributes(PyObject* _self, PyObject* args, PyObject* kwargs); + +// Clears the ExtensionDict from the message. The message must have an +// ExtensionDict set. void PyUpb_CMessage_ClearExtensionDict(PyObject* _self); + +// Implements the equivalent of getattr(msg, field), once `field` has +// already been resolved to a `upb_fielddef*`. PyObject* PyUpb_CMessage_GetFieldValue(PyObject* _self, const upb_fielddef* field); + +// Implements the equivalent of setattr(msg, field, value), once `field` has +// already been resolved to a `upb_fielddef*`. int PyUpb_CMessage_SetFieldValue(PyObject* _self, const upb_fielddef* field, PyObject* value); + +// Returns the version associated with this message. The version will be +// incremented when the message changes. int PyUpb_CMessage_GetVersion(PyObject* _self); +// Module-level init. bool PyUpb_InitMessage(PyObject* m); #endif // PYPB_MESSAGE_H__ diff --git a/python/protobuf.c b/python/protobuf.c index 367f70d796..2c4d5857d8 100644 --- a/python/protobuf.c +++ b/python/protobuf.c @@ -33,7 +33,7 @@ static void PyUpb_ModuleDealloc(void *module) { PyUpb_ModuleState *s = PyModule_GetState(module); - upb_arena_free(s->obj_cache_arena); + PyUpb_WeakMap_Free(s->obj_cache); } static struct PyModuleDef module_def = {PyModuleDef_HEAD_INIT, @@ -88,33 +88,59 @@ PyObject *PyUpb_GetWktBases(PyUpb_ModuleState *state) { } // ----------------------------------------------------------------------------- -// ObjectCache +// WeakMap // ----------------------------------------------------------------------------- -void PyUpb_ObjCache_Add(const void *key, PyObject *py_obj) { - PyUpb_ModuleState *s = PyUpb_ModuleState_Get(); - upb_inttable_insert(&s->obj_cache, (uintptr_t)key, upb_value_ptr(py_obj), - s->obj_cache_arena); +struct PyUpb_WeakMap { + upb_inttable table; + upb_arena *arena; +}; + +PyUpb_WeakMap *PyUpb_WeakMap_New(void) { + upb_arena *arena = upb_arena_new(); + PyUpb_WeakMap *map = upb_arena_malloc(arena, sizeof(*map)); + map->arena = arena; + upb_inttable_init(&map->table, map->arena); + return map; } -void PyUpb_ObjCache_Delete(const void *key) { - PyUpb_ModuleState *s = PyUpb_ModuleState_MaybeGet(); - if (!s) { - // During the shutdown sequence, our object's Dealloc() methods can be - // called *after* our module Dealloc() method has been called. At that - // point our state will be NULL and there is nothing to delete out of the - // map. - return; +void PyUpb_WeakMap_Free(PyUpb_WeakMap *map) { + upb_arena_free(map->arena); +} + +uintptr_t PyUpb_WeakMap_GetKey(const void *key) { + uintptr_t n = (uintptr_t)key; + assert((n & 7) == 0); + return n >> 3; +} + +void PyUpb_WeakMap_Add(PyUpb_WeakMap *map, const void *key, PyObject *py_obj) { + /* + uintptr_t k = (uintptr_t)key; + for (int i = 0; i < 64; i++) { + counts[i] += ((k & (1ULL << i)) != 0); } + total++; + */ + upb_inttable_insert(&map->table, PyUpb_WeakMap_GetKey(key), + upb_value_ptr(py_obj), map->arena); +} + +void PyUpb_WeakMap_Delete(PyUpb_WeakMap *map, const void *key) { upb_value val; - upb_inttable_remove(&s->obj_cache, (uintptr_t)key, &val); - assert(upb_value_getptr(val)); + bool removed = + upb_inttable_remove(&map->table, PyUpb_WeakMap_GetKey(key), &val); + (void)removed; + assert(removed); } -PyObject *PyUpb_ObjCache_Get(const void *key) { - PyUpb_ModuleState *s = PyUpb_ModuleState_Get(); +void PyUpb_WeakMap_TryDelete(PyUpb_WeakMap *map, const void *key) { + upb_inttable_remove(&map->table, PyUpb_WeakMap_GetKey(key), NULL); +} + +PyObject *PyUpb_WeakMap_Get(PyUpb_WeakMap *map, const void *key) { upb_value val; - if (upb_inttable_lookup(&s->obj_cache, (uintptr_t)key, &val)) { + if (upb_inttable_lookup(&map->table, PyUpb_WeakMap_GetKey(key), &val)) { PyObject *ret = upb_value_getptr(val); Py_INCREF(ret); return ret; @@ -123,6 +149,52 @@ PyObject *PyUpb_ObjCache_Get(const void *key) { } } +bool PyUpb_WeakMap_Next(PyUpb_WeakMap *map, const void **key, PyObject **obj, + intptr_t *iter) { + uintptr_t u_key; + upb_value val; + if (upb_inttable_next2(&map->table, &u_key, &val, iter)) { + *key = (void *)(u_key << 3); + *obj = upb_value_getptr(val); + return true; + } else { + return false; + } +} + +void PyUpb_WeakMap_DeleteIter(PyUpb_WeakMap *map, intptr_t *iter) { + upb_inttable_removeiter(&map->table, iter); +} + +// ----------------------------------------------------------------------------- +// ObjCache +// ----------------------------------------------------------------------------- + +PyUpb_WeakMap *PyUpb_ObjCache_Instance(void) { + PyUpb_ModuleState *state = PyUpb_ModuleState_Get(); + return state->obj_cache; +} + +void PyUpb_ObjCache_Add(const void *key, PyObject *py_obj) { + PyUpb_WeakMap_Add(PyUpb_ObjCache_Instance(), key, py_obj); +} + +void PyUpb_ObjCache_Delete(const void *key) { + PyUpb_ModuleState *state = PyUpb_ModuleState_MaybeGet(); + if (!state) { + // During the shutdown sequence, our object's Dealloc() methods can be + // called *after* our module Dealloc() method has been called. At that + // point our state will be NULL and there is nothing to delete out of the + // map. + return; + } + PyUpb_WeakMap_Delete(state->obj_cache, key); +} + +PyObject *PyUpb_ObjCache_Get(const void *key) { + return PyUpb_WeakMap_Get(PyUpb_ObjCache_Instance(), key); +} + // ----------------------------------------------------------------------------- // Arena // ----------------------------------------------------------------------------- @@ -218,11 +290,12 @@ PyObject *PyUpb_Forbidden_New(PyObject *cls, PyObject *args, PyObject *kwds) { PyMODINIT_FUNC PyInit__message(void) { PyObject *m = PyModule_Create(&module_def); - PyState_AddModule(m, &module_def); - PyUpb_ModuleState *state = PyUpb_ModuleState_Get(); + if (!m) return NULL; + + PyUpb_ModuleState *state = PyUpb_ModuleState_GetFromModule(m); - state->obj_cache_arena = upb_arena_new(); - upb_inttable_init(&state->obj_cache, state->obj_cache_arena); + state->wkt_bases = NULL; + state->obj_cache = PyUpb_WeakMap_New(); if (!PyUpb_InitDescriptorContainers(m) || !PyUpb_InitDescriptorPool(m) || !PyUpb_InitDescriptor(m) || !PyUpb_InitArena(m)) { diff --git a/python/protobuf.h b/python/protobuf.h index 161b26b887..0bc13f610c 100644 --- a/python/protobuf.h +++ b/python/protobuf.h @@ -38,6 +38,9 @@ #define PYUPB_RETURN_OOM return PyErr_SetNone(PyExc_MemoryError), NULL +struct PyUpb_WeakMap; +typedef struct PyUpb_WeakMap PyUpb_WeakMap; + // ----------------------------------------------------------------------------- // ModuleState // ----------------------------------------------------------------------------- @@ -73,8 +76,7 @@ typedef struct { // From protobuf.c PyObject *wkt_bases; PyTypeObject *arena_type; - upb_arena *obj_cache_arena; - upb_inttable obj_cache; + PyUpb_WeakMap *obj_cache; } PyUpb_ModuleState; // Returns the global state object from the current interpreter. The current @@ -105,9 +107,6 @@ PyObject *PyUpb_GetWktBases(PyUpb_ModuleState *state); // remove itself from the map when it is destroyed. The map is weak so it does // not take references to the cached objects. -struct PyUpb_WeakMap; -typedef struct PyUpb_WeakMap PyUpb_WeakMap; - PyUpb_WeakMap *PyUpb_WeakMap_New(void); void PyUpb_WeakMap_Free(PyUpb_WeakMap *map); @@ -128,20 +127,15 @@ bool PyUpb_WeakMap_Next(PyUpb_WeakMap *map, const void **key, PyObject **obj, void PyUpb_WeakMap_DeleteIter(PyUpb_WeakMap *map, intptr_t *iter); // ----------------------------------------------------------------------------- -// ObjectCache +// ObjCache // ----------------------------------------------------------------------------- // The object cache is a global WeakMap for mapping upb objects to the // corresponding wrapper. - -// Adds the given object to the cache, indexed by the given key. void PyUpb_ObjCache_Add(const void *key, PyObject *py_obj); - -// Removes the given key from the cache. It must exist in the cache currently. void PyUpb_ObjCache_Delete(const void *key); - -// Returns a new reference to an object if it exists, otherwise returns NULL. -PyObject *PyUpb_ObjCache_Get(const void *key); +PyObject *PyUpb_ObjCache_Get(const void *key); // returns NULL if not present. +PyUpb_WeakMap *PyUpb_ObjCache_Instance(void); // ----------------------------------------------------------------------------- // Arena