#include <u.h>
#include <libc.h>
#include <thread.h>

typedef struct WaitGroup WaitGroup;
struct WaitGroup {
	Rendez;
	QLock;
	Ref;
};

extern int cas(long *p, long ov, long nv);

void
wginit(WaitGroup *wg, long count)
{
	memset(wg, 0, sizeof(*wg));
	wg->l = &wg->QLock;
	wg->ref = count;
	if(cas(&wg->ref, 0, count) == 0)
		sysfatal("wginit: cas failed");
}

void
wgadd(WaitGroup *wg, long n)
{
	long v;
	v = wg->ref;
	while(cas(&wg->ref, v, v+n) == 0)
		v = wg->ref;
}

void
wgdone(WaitGroup *wg)
{
	if(decref(wg) < 0)
		sysfatal("wgdone: negative WaitGroup count");
	qlock(wg);
	if(rwakeupall(wg) == 0)
		sysfatal("wgdone: zero waiters");
	qunlock(wg);
}

void
wgwait(WaitGroup *wg)
{
	qlock(wg);
	while(!(wg->ref == 0))
		rsleep(wg);
	qunlock(wg);
}

void
proc(void *v)
{
	int id;
	WaitGroup *wg;

	id = threadid();
	wg = v;

	print("proc %d start\n", id);

	sleep(id*250);

	wgdone(wg);

	print("proc %d exit\n", id);
}

void
thr(void *v)
{
	int id;
	Ioproc *io;
	WaitGroup *wg;

	id = threadid();
	wg = v;

	print("thr %d start\n", id);

	io = ioproc();
	iosleep(io, id*250);
	closeioproc(io);

	wgdone(wg);

	print("thr %d exit\n", id);	
}

void
threadmain(int argc, char *argv[])
{
	int rv, nproc;
	WaitGroup wg;

	ARGBEGIN{
	}ARGEND

	srand(truerand());

	nproc = 10;

	wginit(&wg, nproc);

	for(rv = 0; rv < nproc; rv++){
		//wgadd(&wg, 1);
		if((rv % 2) == 0)
			proccreate(proc, &wg, 8192);
		else
			threadcreate(thr, &wg, 8192);
	}

	wgwait(&wg);
	print("all threads exited\n");
	threadexitsall(nil);
}