pool.c

#include <stdio.h>
#include <stdlib.h>
#include <stdatomic.h>
#include <pthread.h>

#define THREAD_COUNT 4

atomic_int done=0;

struct Node{
	void (*function) (void *params);
	void *params;
	_Atomic(struct Node *) next;
};
struct Queue{
	_Atomic(struct Node *) head;
	_Atomic(struct Node *) tail;
};
struct Retired_node{
	struct Node *node;
	int epoch;
	struct Retired_node *next;
};
struct Old_age_home{
	struct Retired_node *head;
	struct Retired_node *tail;
};
struct Thread{
	atomic_int local_epoch;
	atomic_int active;
	struct Old_age_home *retired;
};
struct Args{
	struct Queue *queue;
	int id;
	struct Thread **lookup;
	atomic_int *global_epoch;
	struct Args *args;
};

struct Node* queue_init(struct Queue *queue){
	struct Node *sentinel=malloc(sizeof(struct Node));
	atomic_store(&sentinel->next,NULL);
	atomic_store(&queue->head,sentinel);
	atomic_store(&queue->tail,sentinel);
	return sentinel;
}

int push(struct Queue *queue,struct Node *node){
	atomic_store(&node->next,NULL);
	for (;;){
		struct Node *tail=atomic_load(&queue->tail);
		struct Node *next=atomic_load(&tail->next);
		if (tail!=atomic_load(&queue->tail)){
			continue;
		}
		if (next==NULL){
			if (atomic_compare_exchange_weak(&tail->next,&next,node)){
				atomic_compare_exchange_weak(&queue->tail,&tail,node);
				return 0;
			}
		}else{
			atomic_compare_exchange_weak(&queue->tail,&tail,next);
		}
	}
}

struct Node* pop(struct Args *a){
	struct Queue *queue=a->queue;
	for (;;){
		struct Node *head=atomic_load(&queue->head);
		struct Node *tail=atomic_load(&queue->tail);
		struct Node *next=atomic_load(&head->next);
		if (head!=atomic_load(&queue->head)){
			continue; 
		}
		if (head==tail){
			if (next==NULL){
				return NULL;
			}
			atomic_compare_exchange_weak(&queue->tail,&tail,next);
			continue;
		}
		if (atomic_compare_exchange_weak(&queue->head,&head,next)){
			return next;
		}
	}
}

void retire(struct Old_age_home *retired,struct Node *node,int epoch){
	struct Retired_node *new_retiree=malloc(sizeof(struct Retired_node));
	new_retiree->node=node;
	new_retiree->epoch=epoch;
	new_retiree->next=NULL;
	if (retired->head==NULL){
		retired->head=new_retiree;
		retired->tail=new_retiree;
	}else{
		retired->tail->next=new_retiree;
		retired->tail=new_retiree;
	}
}

int advance_epoch(struct Thread **lookup,struct Args *a){
	int global_epoch=atomic_load(a->global_epoch);
	for (int i=0;i<THREAD_COUNT;i++){
		if (atomic_load(&lookup[i]->active) && atomic_load(&lookup[i]->local_epoch)!=global_epoch){
			return 1;
		}
	}
	return !atomic_compare_exchange_weak(a->global_epoch,&global_epoch,global_epoch+1);
}

void grim_reaper(struct Old_age_home *retired,int local_epoch){
	struct Retired_node *curr=retired->head;
	for (;;){
		if (curr==NULL || local_epoch-curr->epoch<=1){
			return;
		}else{
			free(curr->node);
			struct Retired_node *temp=curr;
			curr=curr->next;
			free(temp);
		}
	}
}

int pick_victim(int self_id,unsigned *seed){
	int victim;
	do{
		victim=rand_r(seed)%THREAD_COUNT;
	}while(victim==self_id);
	return victim;
}

void *run(void *args){
	struct Args *a=args;
	unsigned seed=(unsigned)a->id;
	for (;;){
		atomic_store(&a->lookup[a->id]->local_epoch,atomic_load(a->global_epoch));
		atomic_store(&a->lookup[a->id]->active,1);
		struct Node *node=pop(a);
		atomic_store(&a->lookup[a->id]->active,0);
		advance_epoch(a->lookup,a);
	check:
		if (node!=NULL){
			node->function(node->params);
			retire(a->lookup[a->id]->retired,node,atomic_load(a->global_epoch));
			continue;
		}
		if (atomic_load(&done)){
			break;
		}
		grim_reaper(a->lookup[a->id]->retired,a->lookup[a->id]->local_epoch);
		int victim=pick_victim(a->id,&seed);
		atomic_store(&a->lookup[a->id]->local_epoch,atomic_load(a->global_epoch));
		atomic_store(&a->lookup[a->id]->active,1);
		node=pop(&a->args[victim]);
		atomic_store(&a->lookup[a->id]->active,0);
		if (node!=NULL){
			goto check;
		}
		grim_reaper(a->lookup[a->id]->retired,a->lookup[a->id]->local_epoch);
	}
	printf("Done! %i\n",a->id);
	return NULL;
}

void destroy(struct Thread *lookup[THREAD_COUNT]){
	for (int i=0;i<THREAD_COUNT;i++){
		struct Thread *thread=lookup[i];
		struct Retired_node *curr=thread->retired->head;
		for (;;){
			if (curr==NULL){
				free(thread->retired);
				free(thread);
				break;
			}
			struct Retired_node *temp=curr;
			free(curr->node);
			curr=curr->next;
			free(temp);
		}
	}
}

int submit_job(struct Args *args,struct Node *node,int last_index){
	int index=(last_index+1)%THREAD_COUNT;
	push(args[index].queue,node);
	return index;
}

void test_print(void *params){
	int *s=(int *)params;
	printf("%i\n",*s);
}

int main(){
	pthread_t threads[THREAD_COUNT];
	atomic_int global_epoch=0;
	struct Args args[THREAD_COUNT];
	struct Thread* lookup[THREAD_COUNT];
	struct Node* sentinels[THREAD_COUNT];
	int last_index=-1;
	for (int i=0;i<THREAD_COUNT;i++){
		struct Thread *thread=malloc(sizeof(struct Thread));
		atomic_store(&thread->active,0);
		atomic_store(&thread->local_epoch,0);
		struct Old_age_home *retired=malloc(sizeof(struct Old_age_home));
		retired->head=NULL;
		retired->tail=NULL;
		thread->retired=retired;
		lookup[i]=thread;
	}
	for (int i=0;i<THREAD_COUNT;i++){
		struct Queue *queue=malloc(sizeof(struct Queue));
		struct Args arg={.id=i,.queue=queue,.lookup=lookup,.global_epoch=&global_epoch,.args=args};
		args[i]=arg;
		sentinels[i]=queue_init(queue);
	}
	for (int i=0;i<THREAD_COUNT;i++){
		pthread_create(threads+i,NULL,run,&args[i]);
	}
	for (int i=1;i<51;i++){
		struct Node *node=malloc(sizeof(struct Node));
		node->function=test_print;
		node->params=&last_index;
		last_index=submit_job(args,node,last_index);
		
	}
	atomic_store(&done,1);
	for (int i=0;i<THREAD_COUNT;i++){
		pthread_join(threads[i],NULL);
	}
	for (int i=0;i<THREAD_COUNT;i++){
		free(sentinels[i]);
		free(args[i].queue);
	}
	destroy(lookup);
	return 0;
}