Some more cleanups before sending for PR.

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;
 }

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;
 }

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__

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)) {

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