ccan/ccan/lbalance/lbalance.c

#include <ccan/lbalance/lbalance.h>
#include <ccan/tlist/tlist.h>
#include <sys/time.h>
#include <sys/resource.h>
#include <unistd.h>
#include <errno.h>
#include <assert.h>
#include <stdlib.h>

/* Define tlist_lbalance_task */
TLIST_TYPE(lbalance_task, struct lbalance_task);

struct stats {
	/* How many stats of for this value do we have? */
	unsigned int num_stats;
	/* What was our total work rate? */
	float work_rate;
};

struct lbalance {
	struct tlist_lbalance_task tasks;
	unsigned int num_tasks;

	/* We figured out how many we want to run. */
	unsigned int target;
	/* We need to recalc once a report comes in via lbalance_task_free. */
	bool target_uptodate;

	/* Integral of how many tasks were running so far */
	struct timeval prev_tasks_time;
	float tasks_sum;

	/* For differential rusage. */
	struct rusage prev_usage;

	/* How many stats we have collected (we invalidate old ones). */
	unsigned int total_stats;

	/* Array of stats, indexed by number of tasks we were running. */
	unsigned int max_stats;
	struct stats *stats;
};

struct lbalance_task {
	struct lbalance *lb;
	struct list_node list;

	/* The time this task started */
	struct timeval start;
	float tasks_sum_start;
};

struct lbalance *lbalance_new(void)
{
	struct lbalance *lb = malloc(sizeof *lb);
	if (!lb)
		return NULL;

	tlist_init(&lb->tasks);
	lb->num_tasks = 0;
	gettimeofday(&lb->prev_tasks_time, NULL);
	lb->tasks_sum = 0.0;

	getrusage(RUSAGE_CHILDREN, &lb->prev_usage);

	lb->max_stats = 1;
	lb->stats = malloc(sizeof(lb->stats[0]) * lb->max_stats);
	if (!lb->stats) {
		free(lb);
		return NULL;
	}
	lb->stats[0].num_stats = 0;
	lb->stats[0].work_rate = 0.0;
	lb->total_stats = 0;

	/* Start with # CPUS as a guess. */
	lb->target = -1L;
#ifdef _SC_NPROCESSORS_ONLN
	lb->target = sysconf(_SC_NPROCESSORS_ONLN);
#elif defined(_SC_NPROCESSORS_CONF)
	if (lb->target == (unsigned int)-1L)
		lb->target = sysconf(_SC_NPROCESSORS_CONF);
#endif
	/* Otherwise, two is a good number. */
	if (lb->target == (unsigned int)-1L || lb->target < 2)
		lb->target = 2;
	lb->target_uptodate = true;

	return lb;
}

/* Return time differences in usec */
static float timeval_sub(struct timeval recent, struct timeval old)
{
	float diff;

	if (old.tv_usec > recent.tv_usec) {
		diff = 1000000 + recent.tv_usec - old.tv_usec;
		recent.tv_sec--;
	} else
		diff = recent.tv_usec - old.tv_usec;

	diff += (float)(recent.tv_sec - old.tv_sec) * 1000000;
	return diff;
}

/* There were num_tasks running between prev_tasks_time and now. */
static void update_tasks_sum(struct lbalance *lb,
			     const struct timeval *now)
{
	lb->tasks_sum += timeval_sub(*now, lb->prev_tasks_time)
		* lb->num_tasks;
	lb->prev_tasks_time = *now;
}

struct lbalance_task *lbalance_task_new(struct lbalance *lb)
{
	struct lbalance_task *task = malloc(sizeof *task);
	if (!task)
		return NULL;

	if (lb->num_tasks + 1 == lb->max_stats) {
		struct stats *s = realloc(lb->stats,
					  sizeof(*s) * (lb->max_stats + 1));
		if (!s) {
			free(task);
			return NULL;
		}
		lb->stats = s;
		lb->stats[lb->max_stats].num_stats = 0;
		lb->stats[lb->max_stats].work_rate = 0.0;
		lb->max_stats++;
	}

	task->lb = lb;
	gettimeofday(&task->start, NULL);

	/* Record that we ran num_tasks up until now. */
	update_tasks_sum(lb, &task->start);

	task->tasks_sum_start = lb->tasks_sum;
	tlist_add_tail(&lb->tasks, task, list);
	lb->num_tasks++;

	return task;
}

/* We slowly erase old stats, once we have enough. */
static void degrade_stats(struct lbalance *lb)
{
	unsigned int i;

	if (lb->total_stats < lb->max_stats * 16)
		return;

#if 0
	fprintf(stderr, ".");
#endif
	for (i = 0; i < lb->max_stats; i++) {
		struct stats *s = &lb->stats[i];
		unsigned int stats_lost = (s->num_stats + 1) / 2;
		s->work_rate *= (float)(s->num_stats - stats_lost)
			/ s->num_stats;
		s->num_stats -= stats_lost;
		lb->total_stats -= stats_lost;
		if (s->num_stats == 0)
			s->work_rate = 0.0;
	}
}

static void add_to_stats(struct lbalance *lb,
			 unsigned int num_tasks,
			 float work_rate)
{
#if 0
	fprintf(stderr, "With %.2f running, work rate was %.5f\n",
		num_tasks, work_rate);
#endif
	assert(num_tasks >= 1);
	assert(num_tasks < lb->max_stats);

	lb->stats[num_tasks].num_stats++;
	lb->stats[num_tasks].work_rate += work_rate;
	lb->total_stats++;
	lb->target_uptodate = false;
}

void lbalance_task_free(struct lbalance_task *task,
			const struct rusage *usage)
{
	float work_done, duration;
	unsigned int num_tasks;
	struct timeval now;
	struct rusage ru;

	gettimeofday(&now, NULL);
	duration = timeval_sub(now, task->start);
	
	getrusage(RUSAGE_CHILDREN, &ru);
	if (usage) {
		work_done = usage->ru_utime.tv_usec + usage->ru_stime.tv_usec
			+ (usage->ru_utime.tv_sec + usage->ru_stime.tv_sec)
			* 1000000;
	} else {
		/* Take difference in rusage as rusage of that task. */
		work_done = timeval_sub(ru.ru_utime,
					task->lb->prev_usage.ru_utime)
			+ timeval_sub(ru.ru_stime,
				      task->lb->prev_usage.ru_utime);
	}
	/* Update previous usage. */
	task->lb->prev_usage = ru;

	/* Record that we ran num_tasks up until now. */
	update_tasks_sum(task->lb, &now);

	/* So, on average, how many tasks were running during this time? */
	num_tasks = (task->lb->tasks_sum - task->tasks_sum_start)
		/ duration + 0.5;

	/* Record the work rate for that many tasks. */
	add_to_stats(task->lb, num_tasks, work_done / duration);

	/* We throw away old stats. */
	degrade_stats(task->lb);

	/* We need to recalculate the target. */
	task->lb->target_uptodate = false;

	/* Remove this task. */
	tlist_del_from(&task->lb->tasks, task, list);
	task->lb->num_tasks--;
	free(task);
}

/* We look for the point where the work rate starts to drop.  Say you have
 * 4 cpus, we'd expect the work rate for 5 processes to drop 20%.
 *
 * If we're within 1/4 of that ideal ratio, we assume it's still
 * optimal.  Any drop of more than 1/2 is interpreted as the point we
 * are overloaded. */
static unsigned int best_target(const struct lbalance *lb)
{
	unsigned int i, best = 0, found_drop = 0;
	float best_f_max = -1.0, cliff = -1.0;

#if 0
	for (i = 1; i < lb->max_stats; i++) {
		printf("%u: %f (%u)\n", i,
		       lb->stats[i].work_rate / lb->stats[i].num_stats,
		       lb->stats[i].num_stats);
	}
#endif

	for (i = 1; i < lb->max_stats; i++) {
		float f;

		if (!lb->stats[i].num_stats)
			f = 0;
		else
			f = lb->stats[i].work_rate / lb->stats[i].num_stats;

		if (f > best_f_max) {
#if 0
			printf("Best is %i\n", i);
#endif
			best_f_max = f - (f / (i + 1)) / 4;
			cliff = f - (f / (i + 1)) / 2;
			best = i;
			found_drop = 0;
		} else if (!found_drop && f < cliff) {
#if 0
			printf("Found drop at %i\n", i);
#endif
			found_drop = i;
		}
	}

	if (found_drop) {
		return found_drop - 1;
	}
	return i - 1;
}

static unsigned int calculate_target(struct lbalance *lb)
{
	unsigned int target;

	target = best_target(lb);

	/* Jitter if the adjacent ones are unknown. */
	if (target >= lb->max_stats || lb->stats[target].num_stats == 0)
		return target;

	if (target + 1 == lb->max_stats || lb->stats[target+1].num_stats == 0)
		return target + 1;

	if (target > 1 && lb->stats[target-1].num_stats == 0)
		return target - 1;

	return target;
}

unsigned lbalance_target(struct lbalance *lb)
{
	if (!lb->target_uptodate) {
		lb->target = calculate_target(lb);
		lb->target_uptodate = true;
	}
	return lb->target;
}
	
void lbalance_free(struct lbalance *lb)
{
	struct lbalance_task *task;

	while ((task = tlist_top(&lb->tasks, struct lbalance_task, list))) {
		assert(task->lb == lb);
		tlist_del_from(&lb->tasks, task, list);
		lb->num_tasks--;
		free(task);
	}
	assert(lb->num_tasks == 0);
	free(lb->stats);
	free(lb);
}