/*
 * reproducer v3 for:
 * [BUG] scheduler doesn't balance thread to idle cpu for 3 seconds
 *
 * Based on LTP's pthread_cond_wait_1.c
 *
 */

#define _GNU_SOURCE
#include <pthread.h>
#include <sched.h>
#include <stdio.h>
#include <stdlib.h>
#include <signal.h>
#include <inttypes.h>
#include <unistd.h>
#include <time.h>
#include <sys/time.h>
#include <sys/resource.h>

#define ERROR_PREFIX "unexpected error: "

#define HIGH_PRIORITY 10
#define LOW_PRIORITY  5
#define RUNTIME       5
#define POLICY        SCHED_RR

#define PTS_PASS 0
#define PTS_FAIL 1
#define PTS_UNRESOLVED 2

pthread_mutex_t mutex = PTHREAD_MUTEX_INITIALIZER;
pthread_cond_t cond = PTHREAD_COND_INITIALIZER;

/* Flags that the threads use to indicate events */
volatile int woken_up = 0;
volatile int low_done = 0;

/* Signal handler that handle the ALRM and wakes up
 * the high priority thread
 */
void signal_handler(int sig)
{
	(void) sig;
	if (pthread_cond_signal(&cond) != 0) {
		printf(ERROR_PREFIX "pthread_cond_signal\n");
		exit(PTS_UNRESOLVED);
	}
}

/* Utility function to find difference between two time values */
float timediff(struct timespec t2, struct timespec t1)
{
	float diff = t2.tv_sec - t1.tv_sec;
	diff += (t2.tv_nsec - t1.tv_nsec) / 1000000000.0;
	return diff;
}

void *hi_priority_thread(void *tmp)
{
	struct sched_param param;
	int policy;
	int rc = 0;

	(void) tmp;
	param.sched_priority = HIGH_PRIORITY;

	rc = pthread_setschedparam(pthread_self(), POLICY, &param);
	if (rc != 0) {
		printf(ERROR_PREFIX "pthread_setschedparam\n");
		exit(PTS_UNRESOLVED);
	}
	rc = pthread_getschedparam(pthread_self(), &policy, &param);
	if (rc != 0) {
		printf(ERROR_PREFIX "pthread_getschedparam\n");
		exit(PTS_UNRESOLVED);
	}
	if ((policy != POLICY) || (param.sched_priority != HIGH_PRIORITY)) {
		printf("Error: the policy or priority not correct\n");
		exit(PTS_UNRESOLVED);
	}

	/* Install a signal handler for ALRM */
	if (signal(SIGALRM, signal_handler) != 0) {
		perror(ERROR_PREFIX "signal:");
		exit(PTS_UNRESOLVED);
	}

	/* acquire the mutex */
	rc = pthread_mutex_lock(&mutex);
	if (rc != 0) {
		printf(ERROR_PREFIX "pthread_mutex_lock\n");
		exit(PTS_UNRESOLVED);
	}

	/* Setup an alarm to go off in 2 seconds */
	alarm(2);

	/* Block, to be woken up by the signal handler */
	rc = pthread_cond_wait(&cond, &mutex);
	if (rc != 0) {
		printf(ERROR_PREFIX "pthread_cond_wait\n");
		exit(PTS_UNRESOLVED);
	}

	/* This variable is unprotected because the scheduling removes
	 * the contention
	 */
	if (low_done != 1)
		woken_up = 1;

	rc = pthread_mutex_unlock(&mutex);
	if (rc != 0) {
		printf(ERROR_PREFIX "pthread_mutex_unlock\n");
		exit(PTS_UNRESOLVED);
	}
	return NULL;
}

void *low_priority_thread(void *tmp)
{
	struct timespec start_time, current_time;
	struct sched_param param;
	int policy;
	cpu_set_t cpuset;
	int rc = 0, slept_times = 0;
	float slept_for = 0;
	uintptr_t tnum = (uintptr_t)tmp;

	param.sched_priority = LOW_PRIORITY;

	rc = pthread_setschedparam(pthread_self(), POLICY, &param);
	if (rc != 0) {
		printf(ERROR_PREFIX "pthread_setschedparam\n");
		exit(PTS_UNRESOLVED);
	}
	rc = pthread_getschedparam(pthread_self(), &policy, &param);
	if (rc != 0) {
		printf(ERROR_PREFIX "pthread_getschedparam\n");
		exit(PTS_UNRESOLVED);
	}
	if ((policy != POLICY) || (param.sched_priority != LOW_PRIORITY)) {
		printf("Error: the policy or priority not correct\n");
		exit(PTS_UNRESOLVED);
	}

	CPU_ZERO(&cpuset);
	CPU_SET(tnum, &cpuset);

	rc = pthread_setaffinity_np(pthread_self(), sizeof(cpu_set_t), &cpuset);
	if (rc != 0) {
		printf(ERROR_PREFIX "pthread_setaffinity_np\n");
		exit(PTS_UNRESOLVED);
	}

	/* grab the start time and busy loop for 5 seconds */
	clock_gettime(CLOCK_REALTIME, &start_time);
	while (!woken_up && !low_done) {
		clock_gettime(CLOCK_REALTIME, &current_time);
		if (timediff(current_time, start_time) > RUNTIME)
			break;
	}

	low_done = 1;
	return NULL;
}

int main()
{
	pthread_t high_id, *low_id, paused_id;
	struct sched_param param;
	int rc = 0;
	int i, ncpus = sysconf(_SC_NPROCESSORS_ONLN);

	low_id = malloc(ncpus * sizeof(pthread_t));

	/* high prio thread */
	rc = pthread_create(&high_id, NULL, hi_priority_thread, NULL);
	if (rc != 0) {
		printf(ERROR_PREFIX "pthread_create\n");
		exit(PTS_UNRESOLVED);
	}

	/* low prio thread on each cpu except last one */
	for (i = 0; i < ncpus - 1; i++) {
		uintptr_t tnum = i;
		rc = pthread_create(&low_id[i], NULL, low_priority_thread, (void *)tnum);
		if (rc != 0) {
			printf(ERROR_PREFIX "pthread_create\n");
			exit(PTS_UNRESOLVED);
		}
	}

	param.sched_priority = 0;
	rc = pthread_setschedparam(pthread_self(), SCHED_OTHER, &param);
	if (rc != 0) {
		printf(ERROR_PREFIX "pthread_setschedparam\n");
		exit(PTS_UNRESOLVED);
	}

	/* Wait for the threads to exit */
	rc = pthread_join(high_id, NULL);
	if (rc != 0) {
		printf(ERROR_PREFIX "pthread_join\n");
		exit(PTS_UNRESOLVED);
	}

	for (i = 0; i < ncpus - 1; i++) {
		rc = pthread_join(low_id[i], NULL);
		if (rc != 0) {
			printf(ERROR_PREFIX "pthread_join\n");
			exit(PTS_UNRESOLVED);
		}
	}

	if (woken_up == 0) {
		printf("Test FAILED: high priority was not woken up\n");
		exit(PTS_FAIL);
	}

	printf("Test PASSED\n");
	exit(PTS_PASS);
}