Add refcounting and thread-safety to message definitions.

10 files changed, 251 insertions, 31 deletions

diff --git a/gen-deps.sh b/gen-deps.sh
index 6f7b4af..6c0ced3 100755
--- a/gen-deps.sh
+++ b/gen-deps.sh
@@ -11,7 +11,8 @@
 # Since upb_parse.o is actually in src, the dependency information is
 # not used.  To remedy this, we use the -MT flag (see gcc docs).
 
+set -e
 rm -f deps
 for file in $@; do
-  gcc -MM $file -MT ${file%.*}.o -Idescriptor -Isrc -I. >> deps
+  gcc -MM $file -MT ${file%.*}.o -DUPB_THREAD_UNSAFE -Idescriptor -Isrc -I. >> deps
 done
diff --git a/src/upb_atomic.h b/src/upb_atomic.h
new file mode 100644
index 0000000..b3995e4
--- /dev/null
+++ b/src/upb_atomic.h
@@ -0,0 +1,162 @@
+/*
+ * upb - a minimalist implementation of protocol buffers.
+ *
+ * Copyright (c) 2009 Joshua Haberman.  See LICENSE for details.
+ *
+ * Only a very small part of upb is thread-safe.  Notably, individual
+ * messages, arrays, and strings are *not* thread safe for mutating.
+ * However, we do make message *metadata* such as upb_msgdef and
+ * upb_context thread-safe, and their ownership is tracked via atomic
+ * refcounting.  This header implements the small number of atomic
+ * primitives required to support this.  The primitives we implement
+ * are:
+ *
+ * - a reader/writer lock (wrappers around platform-provided mutexes).
+ * - an atomic refcount.
+ */
+
+#ifndef UPB_ATOMIC_H_
+#define UPB_ATOMIC_H_
+
+#include <stdbool.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/* inline if possible, emit standalone code if required. */
+#ifndef INLINE
+#define INLINE static inline
+#endif
+
+#ifdef UPB_THREAD_UNSAFE
+
+/* Non-thread-safe implementations. ******************************************/
+
+typedef struct {
+  int val;
+} upb_atomic_refcount_t;
+
+INLINE void upb_atomic_refcount_init(upb_atomic_refcount_t *a, int val) {
+  a->val = val;
+}
+
+INLINE bool upb_atomic_ref(upb_atomic_refcount_t *a) {
+  return a->val++ == 0;
+}
+
+INLINE bool upb_atomic_unref(upb_atomic_refcount_t *a) {
+  return --a->val == 0;
+}
+
+typedef struct {
+} upb_rwlock_t;
+
+INLINE void upb_rwlock_init(upb_rwlock_t *l) { (void)l; }
+INLINE void upb_rwlock_destroy(upb_rwlock_t *l) { (void)l; }
+INLINE void upb_rwlock_rdlock(upb_rwlock_t *l) { (void)l; }
+INLINE void upb_rwlock_wrlock(upb_rwlock_t *l) { (void)l; }
+INLINE void upb_rwlock_unlock(upb_rwlock_t *l) { (void)l; }
+
+#endif
+
+/* Atomic refcount ************************************************************/
+
+#ifdef UPB_THREAD_UNSAFE
+
+/* Already defined above. */
+
+#elif (__GNUC__ == 4 && __GNUC_MINOR__ >= 1) || __GNUC__ > 4
+
+/* GCC includes atomic primitives. */
+
+typedef struct {
+  volatile int val;
+} upb_atomic_refcount_t;
+
+INLINE void upb_atomic_refcount_init(upb_atomic_refcount_t *a, int val) {
+  a->val = val;
+  __sync_synchronize();   /* Ensure the initialized value is visible. */
+}
+
+INLINE void upb_atomic_ref(upb_atomic_refcount_t *a) {
+  return __sync_fetch_and_add(&a->val) == 0;
+}
+
+INLINE bool upb_atomic_unref(upb_atomic_refcount_t *a) {
+  return __sync_sub_and_fetch(&a->val) == 0;
+}
+
+#elif defined(WIN32)
+
+/* Windows defines atomic increment/decrement. */
+#include <Windows.h>
+
+typedef struct {
+  volatile LONG val;
+} upb_atomic_refcount_t;
+
+INLINE void upb_atomic_refcount_init(upb_atomic_refcount_t *a, int val) {
+  InterlockedExchange(&a->val, val);
+}
+
+INLINE bool upb_atomic_ref(upb_atomic_refcount_t *a) {
+  return InterlockedIncrement(&a->val) == 1;
+}
+
+INLINE bool upb_atomic_unref(upb_atomic_refcount_t *a) {
+  return InterlockedDecrement(&a->val) == 0;
+}
+
+#else
+#error Atomic primitives not defined for your platform/CPU.  \
+       Implement them or compile with UPB_THREAD_UNSAFE.
+#endif
+
+/* Reader/Writer lock. ********************************************************/
+
+#ifdef UPB_THREAD_UNSAFE
+
+/* Already defined. */
+
+#elif defined(_POSIX_THREADS)
+
+typedef struct {
+  pthread_rwlock_t lock;
+} upb_rwlock_t;
+
+INLINE void upb_rwlock_init(upb_rwlock_t *l) {
+  /* TODO: check return value. */
+  pthread_rwlock_init(&l->lock);
+}
+
+INLINE void upb_rwlock_destroy(upb_rwlock_t *l) {
+  /* TODO: check return value. */
+  pthread_rwlock_destroy(&l->lock);
+}
+
+INLINE void upb_rwlock_rdlock(upb_rwlock_t *l) {
+  /* TODO: check return value. */
+  pthread_rwlock_rdlock(&l->lock);
+}
+
+INLINE void upb_rwlock_wrlock(upb_rwlock_t *l) {
+  /* TODO: check return value. */
+  pthread_rwlock_wrlock(&l->lock);
+}
+
+INLINE void upb_rwlock_unlock(upb_rwlock_t *l) {
+  /* TODO: check return value. */
+  pthread_rwlock_unlock(&l->lock);
+}
+
+#else
+#error Reader/writer lock is not defined for your platform/CPU.  \
+       Implement it or compile with UPB_THREAD_UNSAFE.
+#endif
+
+#ifdef __cplusplus
+}  /* extern "C" */
+#endif
+
+#endif  /* UPB_ATOMIC_H_ */
diff --git a/src/upb_context.c b/src/upb_context.c
index 252cf17..fbe3f34 100644
--- a/src/upb_context.c
+++ b/src/upb_context.c
@@ -20,14 +20,19 @@ static int memrchr(char *data, char c, size_t len)
 }
 
 bool addfd(struct upb_strtable *addto, struct upb_strtable *existingdefs,
-           google_protobuf_FileDescriptorProto *fd, bool sort);
+           google_protobuf_FileDescriptorProto *fd, bool sort,
+           struct upb_context *context);
 
 bool upb_context_init(struct upb_context *c)
 {
+  upb_atomic_refcount_init(&c->refcount, 1);
+  upb_rwlock_init(&c->lock);
   upb_strtable_init(&c->symtab, 16, sizeof(struct upb_symtab_entry));
   upb_strtable_init(&c->psymtab, 16, sizeof(struct upb_symtab_entry));
   /* Add all the types in descriptor.proto so we can parse descriptors. */
-  if(!addfd(&c->psymtab, &c->symtab, upb_file_descriptor_set->file->elements[0], false)) {
+  google_protobuf_FileDescriptorProto *fd =
+      upb_file_descriptor_set->file->elements[0]; /* We know there is only 1. */
+  if(!addfd(&c->psymtab, &c->symtab, fd, false, c)) {
     assert(false);
     return false;  /* Indicates that upb is buggy or corrupt. */
   }
@@ -56,7 +61,7 @@ static void free_symtab(struct upb_strtable *t)
   upb_strtable_free(t);
 }
 
-void upb_context_free(struct upb_context *c)
+static void free_context(struct upb_context *c)
 {
   free_symtab(&c->symtab);
   for(size_t i = 0; i < c->fds_len; i++)
@@ -65,6 +70,16 @@ void upb_context_free(struct upb_context *c)
   free(c->fds);
 }
 
+void upb_context_unref(struct upb_context *c)
+{
+  if(upb_atomic_unref(&c->refcount)) {
+    upb_rwlock_wrlock(&c->lock);
+    free_context(c);
+    upb_rwlock_unlock(&c->lock);
+  }
+  upb_rwlock_destroy(&c->lock);
+}
+
 struct upb_symtab_entry *upb_context_lookup(struct upb_context *c,
                                             struct upb_string *symbol)
 {
@@ -146,7 +161,8 @@ static struct upb_string join(struct upb_string *base, struct upb_string *name)
 
 static bool insert_enum(struct upb_strtable *t,
                         google_protobuf_EnumDescriptorProto *ed,
-                        struct upb_string *base)
+                        struct upb_string *base,
+                        struct upb_context *c)
 {
   if(!ed->set_flags.has.name) return false;
 
@@ -163,7 +179,7 @@ static bool insert_enum(struct upb_strtable *t,
   e.e.key = fqname;
   e.type = UPB_SYM_ENUM;
   e.ref._enum = malloc(sizeof(*e.ref._enum));
-  upb_enum_init(e.ref._enum, ed);
+  upb_enum_init(e.ref._enum, ed, c);
   upb_strtable_insert(t, &e.e);
 
   return true;
@@ -171,7 +187,8 @@ static bool insert_enum(struct upb_strtable *t,
 
 static bool insert_message(struct upb_strtable *t,
                            google_protobuf_DescriptorProto *d,
-                           struct upb_string *base, bool sort)
+                           struct upb_string *base, bool sort,
+                           struct upb_context *c)
 {
   if(!d->set_flags.has.name) return false;
 
@@ -188,7 +205,7 @@ static bool insert_message(struct upb_strtable *t,
   e.e.key = fqname;
   e.type = UPB_SYM_MESSAGE;
   e.ref.msg = malloc(sizeof(*e.ref.msg));
-  if(!upb_msgdef_init(e.ref.msg, d, fqname, sort)) {
+  if(!upb_msgdef_init(e.ref.msg, d, fqname, sort, c)) {
     free(fqname.ptr);
     return false;
   }
@@ -197,31 +214,32 @@ static bool insert_message(struct upb_strtable *t,
   /* Add nested messages and enums. */
   if(d->set_flags.has.nested_type)
     for(unsigned int i = 0; i < d->nested_type->len; i++)
-      if(!insert_message(t, d->nested_type->elements[i], &fqname, sort))
+      if(!insert_message(t, d->nested_type->elements[i], &fqname, sort, c))
         return false;
 
   if(d->set_flags.has.enum_type)
     for(unsigned int i = 0; i < d->enum_type->len; i++)
-      if(!insert_enum(t, d->enum_type->elements[i], &fqname))
+      if(!insert_enum(t, d->enum_type->elements[i], &fqname, c))
         return false;
 
   return true;
 }
 
 bool addfd(struct upb_strtable *addto, struct upb_strtable *existingdefs,
-           google_protobuf_FileDescriptorProto *fd, bool sort)
+           google_protobuf_FileDescriptorProto *fd, bool sort,
+           struct upb_context *c)
 {
-  struct upb_string package = {.byte_len=0};
-  if(fd->set_flags.has.package) package = *fd->package;
+  struct upb_string pkg = {.byte_len=0};
+  if(fd->set_flags.has.package) pkg = *fd->package;
 
   if(fd->set_flags.has.message_type)
     for(unsigned int i = 0; i < fd->message_type->len; i++)
-      if(!insert_message(addto, fd->message_type->elements[i], &package, sort))
+      if(!insert_message(addto, fd->message_type->elements[i], &pkg, sort, c))
         return false;
 
   if(fd->set_flags.has.enum_type)
     for(unsigned int i = 0; i < fd->enum_type->len; i++)
-      if(!insert_enum(addto, fd->enum_type->elements[i], &package))
+      if(!insert_enum(addto, fd->enum_type->elements[i], &pkg, c))
         return false;
 
   /* TODO: handle extensions and services. */
@@ -247,7 +265,7 @@ bool addfd(struct upb_strtable *addto, struct upb_strtable *existingdefs,
         else
           continue;  /* No resolving necessary. */
         if(!ref.msg) return false;  /* Ref. to undefined symbol. */
-        upb_msgdef_ref(m, f, ref);
+        upb_msgdef_setref(m, f, ref);
       }
     }
   }
@@ -263,7 +281,7 @@ bool upb_context_addfds(struct upb_context *c,
     struct upb_strtable tmp;
     upb_strtable_init(&tmp, 0, sizeof(struct upb_symtab_entry));
     for(uint32_t i = 0; i < fds->file->len; i++) {
-      if(!addfd(&tmp, &c->symtab, fds->file->elements[i], true)) {
+      if(!addfd(&tmp, &c->symtab, fds->file->elements[i], true, c)) {
         printf("Not added successfully!\n");
         free_symtab(&tmp);
         return false;
diff --git a/src/upb_context.h b/src/upb_context.h
index e802454..a04a875 100644
--- a/src/upb_context.h
+++ b/src/upb_context.h
@@ -14,6 +14,7 @@
 
 #include "upb.h"
 #include "upb_table.h"
+#include "upb_atomic.h"
 
 struct google_protobuf_FileDescriptorProto;
 
@@ -38,6 +39,8 @@ struct upb_symtab_entry {
 };
 
 struct upb_context {
+  upb_atomic_refcount_t refcount;
+  upb_rwlock_t lock;
   struct upb_strtable symtab;   /* The context's symbol table. */
   struct upb_strtable psymtab;  /* Private symbols, for internal use. */
   struct upb_msgdef *fds_msg;   /* In psymtab, ptr here for convenience. */
@@ -48,9 +51,13 @@ struct upb_context {
   struct google_protobuf_FileDescriptorSet **fds;
 };
 
-/* Initializes and frees a upb_context, respectively. */
+/* Initializes a upb_context.  Contexts are not freed explicitly, but unref'd
+ * when the caller is done with them. */
 bool upb_context_init(struct upb_context *c);
-void upb_context_free(struct upb_context *c);
+INLINE void upb_context_ref(struct upb_context *c) {
+  upb_atomic_ref(&c->refcount);
+}
+void upb_context_unref(struct upb_context *c);
 
 /* Looking up symbols. ********************************************************/
 
diff --git a/src/upb_enum.c b/src/upb_enum.c
index b599c9b..4855d89 100644
--- a/src/upb_enum.c
+++ b/src/upb_enum.c
@@ -8,9 +8,12 @@
 #include "upb_enum.h"
 
 void upb_enum_init(struct upb_enum *e,
-                   struct google_protobuf_EnumDescriptorProto *ed) {
+                   struct google_protobuf_EnumDescriptorProto *ed,
+                   struct upb_context *c) {
   int num_values = ed->set_flags.has.value ? ed->value->len : 0;
   e->descriptor = ed;
+  e->context = c;
+  upb_atomic_refcount_init(&e->refcount, 0);
   upb_strtable_init(&e->nametoint, num_values, sizeof(struct upb_enum_ntoi_entry));
   upb_inttable_init(&e->inttoname, num_values, sizeof(struct upb_enum_iton_entry));
 
diff --git a/src/upb_enum.h b/src/upb_enum.h
index cad240a..e43a203 100644
--- a/src/upb_enum.h
+++ b/src/upb_enum.h
@@ -10,10 +10,14 @@
 #define UPB_ENUM_H_
 
 #include <stdint.h>
+#include "upb_atomic.h"
+#include "upb_context.h"
 #include "upb_table.h"
 #include "descriptor.h"
 
 struct upb_enum {
+  upb_atomic_refcount_t refcount;
+  struct upb_context *context;
   struct google_protobuf_EnumDescriptorProto *descriptor;
   struct upb_strtable nametoint;
   struct upb_inttable inttoname;
@@ -29,10 +33,20 @@ struct upb_enum_iton_entry {
   struct upb_string *string;
 };
 
+INLINE void upb_enum_ref(struct upb_enum *e) {
+  if(upb_atomic_ref(&e->refcount)) upb_context_ref(e->context);
+}
+
+INLINE void upb_enum_unref(struct upb_enum *e) {
+  if(upb_atomic_unref(&e->refcount)) upb_context_unref(e->context);
+}
+
+
 /* Initializes and frees an enum, respectively.  Caller retains ownership of
  * ed, but it must outlive e. */
 void upb_enum_init(struct upb_enum *e,
-                   struct google_protobuf_EnumDescriptorProto *ed);
+                   struct google_protobuf_EnumDescriptorProto *ed,
+                   struct upb_context *c);
 void upb_enum_free(struct upb_enum *e);
 
 #endif  /* UPB_ENUM_H_ */
diff --git a/src/upb_msg.c b/src/upb_msg.c
index dd627eb..a6b22e2 100644
--- a/src/upb_msg.c
+++ b/src/upb_msg.c
@@ -42,11 +42,12 @@ void upb_msgdef_sortfds(google_protobuf_FieldDescriptorProto **fds, size_t num)
 }
 
 bool upb_msgdef_init(struct upb_msgdef *m, google_protobuf_DescriptorProto *d,
-                     struct upb_string fqname, bool sort)
+                     struct upb_string fqname, bool sort, struct upb_context *c)
 {
   /* TODO: more complete validation. */
   if(!d->set_flags.has.field) return false;
 
+  upb_atomic_refcount_init(&m->refcount, 0);
   upb_inttable_init(&m->fields_by_num, d->field->len,
                     sizeof(struct upb_fieldsbynum_entry));
   upb_strtable_init(&m->fields_by_name, d->field->len,
@@ -54,6 +55,7 @@ bool upb_msgdef_init(struct upb_msgdef *m, google_protobuf_DescriptorProto *d,
 
   m->descriptor = d;
   m->fqname = fqname;
+  m->context = c;
   m->num_fields = d->field->len;
   m->set_flags_bytes = div_round_up(m->num_fields, 8);
   /* These are incremented in the loop. */
@@ -88,7 +90,7 @@ bool upb_msgdef_init(struct upb_msgdef *m, google_protobuf_DescriptorProto *d,
 
     /* Insert into the tables.  Note that f->ref will be uninitialized, even in
      * the tables' copies of *f, which is why we must update them separately
-     * in upb_msg_ref() below. */
+     * in upb_msg_setref() below. */
     struct upb_fieldsbynum_entry nument = {.e = {.key = fd->number}, .f = *f};
     struct upb_fieldsbyname_entry strent = {.e = {.key = *fd->name}, .f = *f};
     upb_inttable_insert(&m->fields_by_num, &nument.e);
@@ -107,8 +109,8 @@ void upb_msgdef_free(struct upb_msgdef *m)
   free(m->field_descriptors);
 }
 
-void upb_msgdef_ref(struct upb_msgdef *m, struct upb_msg_fielddef *f,
-                    union upb_symbol_ref ref) {
+void upb_msgdef_setref(struct upb_msgdef *m, struct upb_msg_fielddef *f,
+                       union upb_symbol_ref ref) {
   struct google_protobuf_FieldDescriptorProto *d =
       upb_msg_field_descriptor(f, m);
   struct upb_fieldsbynum_entry *int_e = upb_inttable_fast_lookup(
diff --git a/src/upb_msg.h b/src/upb_msg.h
index e2b6903..2ae3f59 100644
--- a/src/upb_msg.h
+++ b/src/upb_msg.h
@@ -54,8 +54,10 @@
 #include <stdint.h>
 
 #include "upb.h"
-#include "upb_table.h"
+#include "upb_atomic.h"
+#include "upb_context.h"
 #include "upb_parse.h"
+#include "upb_table.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -66,6 +68,8 @@ extern "C" {
 struct upb_msg_fielddef;
 /* Structure that describes a single .proto message type. */
 struct upb_msgdef {
+  upb_atomic_refcount_t refcount;
+  struct upb_context *context;
   struct google_protobuf_DescriptorProto *descriptor;
   struct upb_string fqname;      /* Fully qualified. */
   size_t size;
@@ -92,6 +96,14 @@ struct upb_msg_fielddef {
   upb_label_t label;
 };
 
+INLINE void upb_msgdef_ref(struct upb_msgdef *m) {
+  if(upb_atomic_ref(&m->refcount)) upb_context_ref(m->context);
+}
+
+INLINE void upb_msgdef_unref(struct upb_msgdef *m) {
+  if(upb_atomic_unref(&m->refcount)) upb_context_unref(m->context);
+}
+
 INLINE bool upb_issubmsg(struct upb_msg_fielddef *f) {
   return upb_issubmsgtype(f->type);
 }
@@ -379,7 +391,8 @@ void upb_msg_print(struct upb_msg *data, bool single_line, FILE *stream);
  * header for this type that expects the given order. */
 bool upb_msgdef_init(struct upb_msgdef *m,
                      struct google_protobuf_DescriptorProto *d,
-                     struct upb_string fqname, bool sort);
+                     struct upb_string fqname, bool sort,
+                     struct upb_context *c);
 void upb_msgdef_free(struct upb_msgdef *m);
 
 /* Sort the given field descriptors in-place, according to what we think is an
@@ -391,8 +404,8 @@ void upb_msgdef_sortfds(google_protobuf_FieldDescriptorProto **fds, size_t num);
  * the "ref" field in the upb_msg_fielddef.  Since messages can refer to each
  * other in mutually-recursive ways, this step must be separated from
  * initialization. */
-void upb_msgdef_ref(struct upb_msgdef *m, struct upb_msg_fielddef *f,
-                    union upb_symbol_ref ref);
+void upb_msgdef_setref(struct upb_msgdef *m, struct upb_msg_fielddef *f,
+                       union upb_symbol_ref ref);
 
 #ifdef __cplusplus
 }  /* extern "C" */
diff --git a/tests/tests.c b/tests/tests.c
index 2a769a9..5df933c 100644
--- a/tests/tests.c
+++ b/tests/tests.c
@@ -224,7 +224,7 @@ static void test_get_f_uint32_t()
 static void test_upb_context() {
   struct upb_context c;
   ASSERT(upb_context_init(&c));
-  upb_context_free(&c);
+  upb_context_unref(&c);
 }
 
 int main()
diff --git a/tools/upbc.c b/tools/upbc.c
index f9890c6..b3036a0 100644
--- a/tools/upbc.c
+++ b/tools/upbc.c
@@ -686,7 +686,7 @@ int main(int argc, char *argv[])
     fclose(c_file);
   }
   upb_msg_free(fds_msg);
-  upb_context_free(&c);
+  upb_context_unref(&c);
   upb_strfree(descriptor);
   fclose(h_file);