#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <pthread.h>
#include <time.h>

#define THREAD_COUNT 4

/* custom barrier implementation with mutexes 
 * forced variable time execution for threads via sleep 
   gcc barrier_demo.c -o bd -pthread */

typedef struct {
    pthread_mutex_t mutex;
    pthread_cond_t proceed_ok;
    int threads_needed_at_barrier;
    int threads_not_yet_at_barrier;
    int threads_at_barrier;
    unsigned int counter;
} pthread_barrier_t_cst;

pthread_barrier_t_cst mybarrier;

int pthread_barrier_init_cstm(pthread_barrier_t_cst *barrier, void *args, int nthreads) {
    barrier->threads_needed_at_barrier = nthreads;
    barrier->threads_not_yet_at_barrier = nthreads;
    barrier->threads_at_barrier = 0;
    barrier->counter = 0;
    pthread_mutex_init(&barrier->mutex, NULL);
    pthread_cond_init(&barrier->proceed_ok, NULL);
    return 0;
}

int pthread_barrier_wait_cstm(pthread_barrier_t_cst *barrier) {
    // prevent other threads access to mutex
    pthread_mutex_lock(&barrier->mutex);
    unsigned int counter;

    barrier->threads_at_barrier++;
    if (barrier->threads_at_barrier == barrier->threads_needed_at_barrier) {
        barrier->counter++;
        barrier->threads_not_yet_at_barrier = barrier->threads_needed_at_barrier;
        barrier->threads_at_barrier = 0;
        pthread_cond_broadcast(&barrier->proceed_ok);
        pthread_mutex_unlock(&barrier->mutex);
        return 1;
    } else {
        counter = barrier->counter;
        // wait until counter of barrier changes
        while (counter == barrier->counter)
            pthread_cond_wait(&barrier->proceed_ok, &barrier->mutex);

        pthread_mutex_unlock(&barrier->mutex);
        return 0;
    }
}

void* threadFn(void *id_ptr) {
  pthread_t pt;
  pt = pthread_self();
  printf("thread %x running\n", (int)pt);
  int thread_id = *(int*)id_ptr;
  int wait_sec = 1 + rand() % 3;
  printf("thread %d: waiting for %d seconds.\n", thread_id, wait_sec);
  sleep(wait_sec);
  printf("thread %d: done with sleep block.\n", thread_id);

  printf("thread %d is calling barrier_wait()\n", thread_id);
  pthread_barrier_wait_cstm(&mybarrier);

/* When a thread calls pthread_barrier_wait(), it will block until the number of threads specified initially in the pthread_barrier_init() have called pthread_barrier_wait() (and blocked too). When the correct number of threads have called pthread_barrier_wait(), all those threads will unblock at once.*/ 

  printf("thread %d: after pthread_barrier_wait()\n", thread_id);
  return NULL;
}


int main() {
  int i;
  pthread_t ids[THREAD_COUNT];
  int short_ids[THREAD_COUNT];

  srand(time(NULL));

  pthread_barrier_init_cstm(&mybarrier, NULL, THREAD_COUNT + 1);

  for (i=0; i < THREAD_COUNT; i++) {
    short_ids[i] = i;
    pthread_create(&ids[i], NULL, threadFn, &short_ids[i]);
  }

  printf("main() is calling barrier wait.\n");

  pthread_barrier_wait_cstm(&mybarrier);

  printf("main() is continuing after all threads have unblocked.\n");

  for (i=0; i < THREAD_COUNT; i++) {
    pthread_join(ids[i], NULL);
  }

  return 0;
}