Change how the default zone is found

On OSX 10.12, malloc_default_zone returns a special zone that is not
present in the list of registered zones. That zone uses a "lite zone"
if one is present (apparently enabled when malloc stack logging is
enabled), or the first registered zone otherwise. In practice this
means unless malloc stack logging is enabled, the first registered
zone is the default.

So get the list of zones to get the first one, instead of relying on
malloc_default_zone.

1 files changed, 31 insertions, 2 deletions

diff --git a/src/zone.c b/src/zone.c
index 2c17123..bf5c930 100644
--- a/src/zone.c
+++ b/src/zone.c
@@ -168,6 +168,33 @@ zone_force_unlock(malloc_zone_t *zone)
 		jemalloc_postfork_parent();
 }
 
+static malloc_zone_t *get_default_zone()
+{
+	malloc_zone_t **zones = NULL;
+	unsigned int num_zones = 0;
+
+	/*
+	 * On OSX 10.12, malloc_default_zone returns a special zone that is not
+	 * present in the list of registered zones. That zone uses a "lite zone"
+	 * if one is present (apparently enabled when malloc stack logging is
+	 * enabled), or the first registered zone otherwise. In practice this
+	 * means unless malloc stack logging is enabled, the first registered
+	 * zone is the default.
+	 * So get the list of zones to get the first one, instead of relying on
+	 * malloc_default_zone.
+	 */
+        if (KERN_SUCCESS != malloc_get_all_zones(0, NULL, (vm_address_t**) &zones,
+	                                         &num_zones)) {
+		/* Reset the value in case the failure happened after it was set. */
+		num_zones = 0;
+	}
+
+	if (num_zones)
+		return zones[0];
+
+	return malloc_default_zone();
+}
+
 JEMALLOC_ATTR(constructor)
 void
 register_zone(void)
@@ -177,7 +204,7 @@ register_zone(void)
 	 * If something else replaced the system default zone allocator, don't
 	 * register jemalloc's.
 	 */
-	malloc_zone_t *default_zone = malloc_default_zone();
+	malloc_zone_t *default_zone = get_default_zone();
 	malloc_zone_t *purgeable_zone = NULL;
 	if (!default_zone->zone_name ||
 	    strcmp(default_zone->zone_name, "DefaultMallocZone") != 0) {
@@ -272,5 +299,7 @@ register_zone(void)
 			malloc_zone_unregister(purgeable_zone);
 			malloc_zone_register(purgeable_zone);
 		}
-	} while (malloc_default_zone() != &zone);
+
+		default_zone = get_default_zone();
+	} while (default_zone != &zone);
 }