Determine offset AFTER kern_return_t check, not before

Author: AZero13

src/allocator.c

@@ -542,15 +542,14 @@ _dispatch_alloc_maybe_madvise_page(dispatch_continuation_t c)
 	}
 	// They are all unallocated, so we could madvise the page. Try to
 	// take ownership of them all.
-	int last_locked = 0;
-	do {
-		if (!os_atomic_cmpxchg(&page_bitmaps[last_locked], BITMAP_C(0),
+	for (i = 0; i < BITMAPS_PER_PAGE; i++) {
+		if (!os_atomic_cmpxchg(&page_bitmaps[i], BITMAP_C(0),
 				BITMAP_ALL_ONES, relaxed)) {
 			// We didn't get one; since there is a cont allocated in
 			// the page, we can't madvise. Give up and unlock all.
 			goto unlock;
 		}
-	} while (++last_locked < (signed)BITMAPS_PER_PAGE);
+	}
 #if DISPATCH_DEBUG
 	//fprintf(stderr, "%s: madvised page %p for cont %p (next = %p), "
 	//		"[%u+1]=%u bitmaps at %p\n", __func__, page, c, c->do_next,
@@ -654,27 +653,37 @@ _dispatch_allocator_enumerate(task_t remote_task,
 		vm_address_t zone_address, memory_reader_t reader,
 		void (^recorder)(vm_address_t, void *, size_t, bool *stop))
 {
-	const size_t heap_size = remote_dal->dal_magazine_size;
-	const size_t dc_size = remote_dal->dal_allocation_size;
-	const size_t dc_flags_offset = remote_dal->dal_allocation_isa_offset;
-	bool stop = false;
-	void *heap;
-
-	while (zone_address) {
-		// FIXME: improve this by not faulting everything and driving it through
-		//        the bitmap.
-		kern_return_t kr = reader(remote_task, zone_address, heap_size, &heap);
-		size_t offs = remote_dal->dal_first_allocation_offset;
-		if (kr) return kr;
-		while (offs < heap_size) {
-			void *isa = *(void **)(heap + offs + dc_flags_offset);
-			if (isa && isa != remote_dal->dal_deferred_free_isa) {
-				recorder(zone_address + offs, heap + offs, dc_size, &stop);
-				if (stop) return KERN_SUCCESS;
+	if (zone_address)
+	{
+		const size_t heap_size = remote_dal->dal_magazine_size;
+		const size_t dc_size = remote_dal->dal_allocation_size;
+		const size_t dc_flags_offset = remote_dal->dal_allocation_isa_offset;
+		bool stop = false;
+		void *heap = NULL;
+
+		do
+		{
+			// FIXME: improve this by not faulting everything and driving it through
+			//        the bitmap.
+			kern_return_t kr;
+			size_t offs;
+
+			kr = reader(remote_task, zone_address, heap_size, &heap);
+			if (kr)
+				return kr;
+
+			for (offs = remote_dal->dal_first_allocation_offset; offs < heap_size; offs += dc_size)
+			{
+				void *isa = *(void **)(heap + offs + dc_flags_offset);
+				if (isa && isa != remote_dal->dal_deferred_free_isa)
+				{
+					recorder(zone_address + offs, heap + offs, dc_size, &stop);
+					if (stop)
+						return KERN_SUCCESS;
+				}
 			}
-			offs += dc_size;
-		}
-		zone_address = (vm_address_t)((dispatch_heap_t)heap)->header.dh_next;
+			zone_address = (vm_address_t)((dispatch_heap_t)heap)->header.dh_next;
+		} while (zone_address);
 	}
 
 	return KERN_SUCCESS;