Task restartable ranges

static int
_ranges_cmp(const void *_r1, const void *_r2)
{
	const task_restartable_range_t *r1 = _r1;
	const task_restartable_range_t *r2 = _r2;

	if (r1->location != r2->location) {
		return r1->location < r2->location ? -1 : 1;
	}
	if (r1->length == r2->length) {
		return 0;
	}
	return r1->length < r2->length ? -1 : 1;
}

static kern_return_t
_ranges_validate(task_t task, task_restartable_range_t *ranges, uint32_t count)
{
	qsort(ranges, count, sizeof(task_restartable_range_t), _ranges_cmp);
	uint64_t limit = task_has_64Bit_data(task) ? UINT64_MAX : UINT32_MAX;
	uint64_t end, recovery;

	if (count == 0) {
		return KERN_INVALID_ARGUMENT;
	}

	for (size_t i = 0; i < count; i++) {
		if (ranges[i].length > TASK_RESTARTABLE_OFFSET_MAX ||
		    ranges[i].recovery_offs > TASK_RESTARTABLE_OFFSET_MAX) {
			return KERN_INVALID_ARGUMENT;
		}
		if (ranges[i].flags) {
			return KERN_INVALID_ARGUMENT;
		}
		if (os_add_overflow(ranges[i].location, ranges[i].length, &end)) {
			return KERN_INVALID_ARGUMENT;
		}
		if (os_add_overflow(ranges[i].location, ranges[i].recovery_offs, &recovery)) {
			return KERN_INVALID_ARGUMENT;
		}
		if (ranges[i].location > limit || end > limit || recovery > limit) {
			return KERN_INVALID_ARGUMENT;
		}
		if (i + 1 < count && end > ranges[i + 1].location) {
			return KERN_INVALID_ARGUMENT;
… more in source

__attribute__((always_inline))
static mach_vm_address_t
_ranges_lookup(struct restartable_ranges *rr, mach_vm_address_t pc)
{
	task_restartable_range_t *ranges = rr->rr_ranges;
	uint32_t l = 0, r = rr->rr_count;

	if (pc <= ranges[0].location) {
		return 0;
	}
	if (pc >= ranges[r - 1].location + ranges[r - 1].length) {
		return 0;
	}

	while (l < r) {
		uint32_t i = (r + l) / 2;
		mach_vm_address_t location = ranges[i].location;

		if (pc <= location) {
			/* if the PC is exactly at pc_start, no reset is needed */
			r = i;
		} else if (location + ranges[i].length <= pc) {
			/* if the PC is exactly at the end, it's out of the function */
			l = i + 1;
		} else {
			/* else it's strictly in the range, return the recovery pc */
			return location + ranges[i].recovery_offs;
		}
	}
… more in source

__attribute__((noinline))
static void
_restartable_ranges_dispose(struct restartable_ranges *rr, bool hash_remove)
{
	if (hash_remove) {
		rr_lock();
		remqueue(&rr->rr_link);
		rr_unlock();
	}
	kfree_type(struct restartable_ranges, rr);
}

static bool
_restartable_ranges_equals(
	const struct restartable_ranges *rr1,
	const struct restartable_ranges *rr2)
{
	size_t rr1_size = rr1->rr_count * sizeof(task_restartable_range_t);
	return rr1->rr_hash == rr2->rr_hash &&
	       rr1->rr_count == rr2->rr_count &&
	       memcmp(rr1->rr_ranges, rr2->rr_ranges, rr1_size) == 0;
}

static kern_return_t
_restartable_ranges_create(task_t task, task_restartable_range_t *ranges,
    uint32_t count, struct restartable_ranges **rr_storage)
{
	struct restartable_ranges *rr, *rr_found, *rr_base;
	queue_head_t *head;
	uint32_t base_count, total_count;
	size_t base_size, size;
	kern_return_t kr;

	rr_base = *rr_storage;
	base_count = rr_base ? rr_base->rr_count : 0;
	base_size = sizeof(task_restartable_range_t) * base_count;
	size = sizeof(task_restartable_range_t) * count;

	if (os_add_overflow(base_count, count, &total_count)) {
		return KERN_INVALID_ARGUMENT;
	}
	if (total_count > RR_RANGES_MAX) {
		return KERN_RESOURCE_SHORTAGE;
	}

	rr = kalloc_type(struct restartable_ranges,
	    (zalloc_flags_t) (Z_WAITOK | Z_ZERO | Z_NOFAIL));

	queue_chain_init(rr->rr_link);
	os_ref_init(&rr->rr_ref, NULL);
	rr->rr_count = total_count;
	if (base_size) {
		memcpy(rr->rr_ranges, rr_base->rr_ranges, base_size);
… more in source

__attribute__((always_inline))
void
restartable_ranges_release(struct restartable_ranges *rr)
{
	if (os_ref_release_relaxed(&rr->rr_ref) == 0) {
		_restartable_ranges_dispose(rr, true);
	}
}

__attribute__((always_inline))
void
thread_reset_pcs_will_fault(thread_t thread)
{
	/*
	 * Called in the exception handling code while interrupts
	 * are still disabled.
	 */
	os_atomic_store(&thread->t_rr_state.trr_fault_state,
	    (uint8_t)TRR_FAULT_PENDING, relaxed);
}

__attribute__((always_inline))
void
thread_reset_pcs_done_faulting(struct thread *thread)
{
	thread_rr_state_t state = {
		.trr_ipi_ack_pending = ~0,
	};

	/*
	 * Called by the exception handling code on the way back,
	 * or when the thread is terminated.
	 */
	state.trr_value = os_atomic_and_orig(&thread->t_rr_state.trr_value,
	    state.trr_value, relaxed);

	if (__improbable(state.trr_sync_waiting)) {
		task_t task = get_threadtask(thread);

		task_lock(task);
		wakeup_all_with_inheritor(&thread->t_rr_state, THREAD_AWAKENED);
		task_unlock(task);
	}
}

void
thread_reset_pcs_ack_IPI(struct thread *thread)
{
	thread_rr_state_t trrs;

	/*
	 * Called under the thread lock from IPI or CSwitch context.
	 */
	trrs.trr_value = os_atomic_load(&thread->t_rr_state.trr_value, relaxed);
	if (__improbable(trrs.trr_ipi_ack_pending)) {
		trrs.trr_ipi_ack_pending = false;
		if (trrs.trr_fault_state) {
			assert3u(trrs.trr_fault_state, ==, TRR_FAULT_PENDING);
			trrs.trr_fault_state = TRR_FAULT_OBSERVED;
		}
		os_atomic_store(&thread->t_rr_state.trr_value,
		    trrs.trr_value, relaxed);
	}
}

static bool
thread_rr_wait_if_needed(task_t task, thread_t thread)
{
	thread_rr_state_t state;
	bool did_unlock = false;

	state.trr_value = os_atomic_load(&thread->t_rr_state.trr_value, relaxed);
	if (state.trr_value == 0) {
		return did_unlock;
	}

	assert(state.trr_sync_waiting == 0);

	thread_reference(thread);

	/*
	 * The thread_rr_state state machine is:
	 *
	 *                        ,------------ IPI ack --------------.
	 *                        v                                   |
	 *        .-----> {f:N, w:0, ipi:0} --- IPI sent ---> {f:N, w:0, ipi:1}
	 *        |           |        ^                              |
	 *        |           |        |                              |
	 *      fault       will     fault                          will
	 *      done        fault    done                           fault
	 *        |           |        |                              |
	 *        |           v        |                              v
	 *        |       {f:P, w:0, ipi:0} --- IPI sent ---> {f:P, w:0, ipi:1}
	 *        |               |                                   |
	 *        |               |                                   |
… more in source

bool
thread_reset_pcs_in_range(task_t task, thread_t thread)
{
	return _ranges_lookup(task->t_rr_ranges, machine_thread_pc(thread)) != 0;
}

void
thread_reset_pcs_ast(task_t task, thread_t thread)
{
	struct restartable_ranges *rr;
	mach_vm_address_t pc;

	/*
	 * Because restartable_ranges are set while the task only has on thread
	 * and can't be mutated outside of this, no lock is required to read this.
	 */
	rr = task->t_rr_ranges;
	if (thread->active && rr) {
		pc = _ranges_lookup(rr, machine_thread_pc(thread));

		if (pc) {
			machine_thread_reset_pc(thread, pc);
		}
	}

#if MACH_ASSERT
	thread_rr_state_t state;

	state.trr_value = os_atomic_load(&thread->t_rr_state.trr_value, relaxed);
	assert3u(state.trr_fault_state, ==, TRR_FAULT_NONE);
	assert3u(state.trr_sync_waiting, ==, 0);
#endif
}

void
restartable_init(void)
{
	for (size_t i = 0; i < RR_HASH_SIZE; i++) {
		queue_head_init(rr_hash[i]);
	}
}

kern_return_t
task_restartable_ranges_register(
	task_t                    task,
	task_restartable_range_t *ranges,
	mach_msg_type_number_t    count)
{
	kern_return_t kr;

	if (task != current_task()) {
		return KERN_FAILURE;
	}

#if CONFIG_ROSETTA
	// <rdar://problem/48527888> Obj-C adoption of task_restartable_ranges_register breaks Cambria
	if (task_is_translated(task)) {
		return KERN_RESOURCE_SHORTAGE;
	}
#endif

	kr = _ranges_validate(task, ranges, count);

	if (kr == KERN_SUCCESS) {
		task_lock(task);

		if (task->thread_count > 1) {
			kr = KERN_NOT_SUPPORTED;
#if !DEBUG && !DEVELOPMENT
		} else if (task->t_rr_ranges) {
			/*
			 * For security reasons, on release kernels,
… more in source

kern_return_t
task_restartable_ranges_synchronize(task_t task)
{
	thread_pri_floor_t token;
	thread_t thread;
	bool needs_wait = false;
	kern_return_t kr = KERN_SUCCESS;

	if (task != current_task()) {
		return KERN_FAILURE;
	}

	/*
	 * t_rr_ranges can only be set if the process is single threaded.
	 * As a result, `t_rr_ranges` can _always_ be looked at
	 * from current_thread() without holding a lock:
	 * - either because it's the only thread in the task
	 * - or because the existence of another thread precludes
	 *   modification
	 */
	if (!task->t_rr_ranges) {
		return KERN_SUCCESS;
	}

	/*
	 * When initiating a GC, artificially raise the priority for the
	 * duration of sending ASTs, we want to be preemptible, but this
	 * sequence has to terminate in a timely fashion.
	 */
	token = thread_priority_floor_start();
… more in source

struct restartable_ranges {
	queue_chain_t            rr_link;
	os_refcnt_t              rr_ref;
	uint32_t                 rr_count;
	uint32_t                 rr_hash;
	task_restartable_range_t rr_ranges[RR_RANGES_MAX];
}

typedef int (*cmpfunc_t)(const void *a, const void *b);

extern void qsort(void *a, size_t n, size_t es, cmpfunc_t cmp);