cudadl/controller.cc

extern "C" {
#include <event2/listener.h>
#include <event2/bufferevent.h>
#include <event2/buffer.h>
}

#include <NTL/vec_ZZ.h>
#include <NTL/ZZ.h>
#include <NTL/ZZ_p.h>

#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#include <sys/time.h>
#include <unistd.h>

#include <fstream>
#include <vector>
#include <set>
#include <map>
#include <string.h>

#include "evutils.h"
#include "subproblem.h"
#include "controller.h"

NTL_CLIENT

// #undef VERBOSE

// Used to prefix output from this process
std::string output_prefix;

struct SubproblemProgress;

typedef std::set<struct bufferevent *> BESet;
typedef std::map<struct bufferevent *, SubproblemProgress *> BEMap;

static void besetdump(const BESet &bes, ostream &os)
{
    BESet::const_iterator besit;

    os << hex << "    ";
    for (besit = bes.begin(); besit != bes.end(); ++besit) {
	os << *besit << " ";
    }
    os << dec << "\n";
}

static void besetfree(BESet &bes)
{
    BESet::iterator besit;

    for (besit = bes.begin(); besit != bes.end(); ++besit) {
	bufferevent_free(*besit);
    }

    bes.clear();
}

static void bemapdump(const BEMap &bem, ostream &os)
{
    BEMap::const_iterator bemit;

    os << hex << "    ";
    for (bemit = bem.begin(); bemit != bem.end(); ++bemit) {
	os << bemit->first << "->" << bemit->second << " ";
    }
    os << dec << "\n";
}

static void bemapfree(BEMap &bem)
{
    BEMap::iterator bemit;

    for (bemit = bem.begin(); bemit != bem.end(); ++bemit) {
	bufferevent_free(bemit->first);
    }

    bem.clear();
}

struct Statuses {
    BESet idle;
    BEMap working;

    // Dump the state for debug purposes
    void dump(ostream &os) const {
	os << "  idle (" << idle.size() << "):\n";
	besetdump(idle, os);
	os << "  working (" << working.size() << "):\n";
	bemapdump(working, os);
    }

    void free(void) {
	besetfree(idle);
	bemapfree(working);
    }
};

struct FactorDecomp {
    ZZ factor;
    vec_ZZ fvec;
};

static void vsppdump(const vector<SubproblemProgress> &spv, ostream &os);

static unsigned short curproblemid = 0;

static struct ControllerState {
    ZZ rho;
    FactorDecomp p, q;
    ZZ base, target;
    bool working;
    struct timeval started_working;
    vector<SubproblemProgress> subproblems_p, subproblems_q;
    Statuses dpnodes, workers;
    unsigned int num_unsolved_subproblems;
    Worklist worklist;
    struct evconnlistener *listener;
    unsigned int problemid;
    unsigned int dpnode_count;

    ControllerState() : working(false), num_unsolved_subproblems(0),
	listener(NULL), problemid(0), dpnode_count(0) {}

    // Reset the state for a new problem with the same modulus
    void reset(void) {
	base = 0;
	target = 0;
	working = false;
	subproblems_p.clear();
	subproblems_q.clear();
	num_unsolved_subproblems = 0;
	started_working.tv_sec = 0;
	started_working.tv_usec = 0;
	problemid = curproblemid++;
    }

    // Dump the state for debug purposes
    void dump(ostream &os) const {
	if (!working) {
	    os << "Not working\n";
	    return;
	}
	os << "P:\n";
	vsppdump(subproblems_p, os);
	os << "Q:\n";
	vsppdump(subproblems_q, os);
	os << "dpnodes:\n";
	dpnodes.dump(os);
	os << "workers:\n";
	workers.dump(os);
    }
} ctrlstate;

struct IPPort {
    unsigned char ipport[6];

    IPPort(unsigned char *ipp) {
	memmove(ipport, ipp, 6);
    }

    void dump(ostream &os) const {
	os << int(ipport[0]) << "." << int(ipport[1]) << "." <<
	    int(ipport[2]) << "." << int(ipport[3]) << ":" <<
	    ((ipport[4] << 8) + ipport[5]) << " ";
    }
};

typedef vector<IPPort> IPPortSet;

static void ipportsetdump(const IPPortSet &ipps, ostream &os)
{
    IPPortSet::const_iterator ippsit;

    os << "    ";
    for (ippsit = ipps.begin(); ippsit != ipps.end(); ++ippsit) {
	ippsit->dump(os);
    }
    os << "\n";
}

void custom_desired_resources(const ZZ &order,
    unsigned short &desired_dpnodes, unsigned int &max_workers,
    unsigned int &dpfreq);

struct SubproblemProgress : Subproblem {

    // The sets of dpnodes and workers currently working on this subproblem
    BESet dpnodes, workers;
    // The dpnode IPPorts registered for this subproblem
    IPPortSet ipports;

    // The desired number of DPnodes for this subproblem
    unsigned short desired_dpnodes;
    // The maximum number of workers useful for this subproblem
    unsigned int max_workers;

    // Have we found a solution?
    bool solved;
    // The solution, if found.
    ZZ solution;

    unsigned int kernel_launch_count;
    unsigned long long worker_time_ms;
    unsigned int num_workers_replied;
    unsigned int worker_id_counter;

    SubproblemProgress(unsigned short id, const ZZ &b, const ZZ &t,
	    const ZZ &m, const ZZ &o) :
	    // By default, 1 in 1000 points are distinguihed points.  The
	    // number in the next line is 2^32/1000
	    Subproblem(id, b, t, m, o, 4294967), solved(false), kernel_launch_count(0),
		    worker_time_ms(0), num_workers_replied(0), worker_id_counter(0) {

	custom_desired_resources(order, desired_dpnodes, max_workers, dpfreq);

	// We need to update the binary description of the subproblem,
	// as we may have just changed dpfreq.
	updatedesc();
    }

    // Stop all dpnodes and workers
    void stop(void) {
	BESet::iterator iter;
	unsigned char stopcmd[1] = { 'S' };

	for (BESet::iterator iter = dpnodes.begin(); iter != dpnodes.end();
		++iter) {
	    bufferevent_write(*iter, stopcmd, 1);
	}
	for (BESet::iterator iter = workers.begin(); iter != workers.end();
		++iter) {
	    bufferevent_write(*iter, stopcmd, 1);
	}
    }

    void reset(void) {
	BESet::iterator iter;

	for (BESet::iterator iter = dpnodes.begin(); iter != dpnodes.end();
		++iter) {
	    ctrlstate.dpnodes.working.erase(*iter);
	    ctrlstate.dpnodes.idle.insert(*iter);
	}
	for (BESet::iterator iter = workers.begin(); iter != workers.end();
		++iter) {
	    ctrlstate.workers.working.erase(*iter);
	    ctrlstate.workers.idle.insert(*iter);
	}
	dpnodes.clear();
	workers.clear();
	ipports.clear();
    }

    // Dump for debugging purposes
    void dump(ostream &os) const {
	os << "    Subproblem " << problemid << "\n";
	os << "    dpnodes (" << dpnodes.size() << "):\n";
	besetdump(dpnodes, os);
	os << "    workers (" << workers.size() << "):\n";
	besetdump(workers, os);
	os << "    ipports (" << ipports.size() << "):\n";
	ipportsetdump(ipports, os);
	if (solved) {
	    os << "    solution: " << solution << "\n\n";
	}
    }

    void worker_write(struct bufferevent *bev) {
	bev_write(bev);

	bufferevent_write(bev, &worker_id_counter, sizeof(unsigned int));
	worker_id_counter++;

	unsigned short num_ipports = ipports.size();
	bufferevent_write(bev, &num_ipports, 2);

	for (unsigned short i = 0; i < num_ipports; ++i) {
	    bufferevent_write(bev, ipports[i].ipport, 6);
	}
#ifdef VERBOSE
	cerr << "Added worker " << bev << " to subproblem "
		<< problemid << "\n";
#endif
    }
};

// Dump the state for debug purposes
static void vsppdump(const vector<SubproblemProgress> &spv, ostream &os)
{
    vector<SubproblemProgress>::const_iterator spiter;
    int count = 0;

    for (spiter = spv.begin(); spiter != spv.end(); ++spiter) {
	++count;
	os << "  " << count << ":\n";
	spiter->dump(os);
    }
    os << "\n";
}

// Find a subproblem in the given vector that could use another DPnode,
// and give it one of the idle ones.  Only allocate it to a subproblem
// with no current DPnodes if consider_empty is true.
static void find_subproblem_for_dpnode(vector<SubproblemProgress> &spv,
	bool consider_empty)
{
    vector<SubproblemProgress>::iterator spiter;

    for (spiter = spv.begin(); spiter != spv.end(); ++spiter) {
	if (spiter->solved) continue;
	if (spiter->dpnodes.size() == 0 && consider_empty == false) continue;
	// How many DPnodes would we like to have for this subproblem?
	while (spiter->dpnodes.size() < spiter->desired_dpnodes &&
		ctrlstate.dpnodes.idle.size() > 0) {
	    // Get the first idle DPnode
	    BESet::iterator beviter = ctrlstate.dpnodes.idle.begin();
	    struct bufferevent *firstbev = *beviter;

	    // Allocate it to the subproblem
	    spiter->dpnodes.insert(firstbev);
	    ctrlstate.dpnodes.working[firstbev] = &(*spiter);
	    ctrlstate.dpnodes.idle.erase(firstbev);

	    // Tell it to start listening for DPs
	    spiter->bev_write(firstbev);
	}
    }
}

// Find a subproblem in the given vector that has all of its DPnodes and
// could use another worker, and give it one of the idle ones.
static void find_subproblem_for_worker(vector<SubproblemProgress> &spv)
{
    vector<SubproblemProgress>::iterator spiter;

    for (spiter = spv.begin(); spiter != spv.end(); ++spiter) {
	if (spiter->solved) continue;

	while (spiter->ipports.size() == spiter->desired_dpnodes &&
		spiter->workers.size() < spiter->max_workers &&
		ctrlstate.workers.idle.size() > 0) {
	    // Get the first idle worker
	    BESet::iterator beviter = ctrlstate.workers.idle.begin();
	    struct bufferevent *firstbev = *beviter;

	    // Allocate it to the subproblem
	    spiter->workers.insert(firstbev);
	    ctrlstate.workers.working[firstbev] = &(*spiter);
	    ctrlstate.workers.idle.erase(firstbev);

	    // Tell it to start working on the subproblem
	    spiter->worker_write(firstbev);
	}
    }
}

// See if there are any idle DPnodes or workers we can put to use
static void schedule(void)
{
    // Check the DPnodes

    // Iterate through the subproblems, looking for one that can use
    // another DPnode.  First look for subproblems that already have
    // some, but not all, of their DPnodes
    if (ctrlstate.dpnodes.idle.size() > 0) {
	find_subproblem_for_dpnode(ctrlstate.subproblems_p, false);
    }
    if (ctrlstate.dpnodes.idle.size() > 0) {
	find_subproblem_for_dpnode(ctrlstate.subproblems_q, false);
    }
    // If there are still more dpnodes to place, start assigning them to
    // subproblems with no current dpnodes
    if (ctrlstate.dpnodes.idle.size() > 0) {
	find_subproblem_for_dpnode(ctrlstate.subproblems_p, true);
    }
    if (ctrlstate.dpnodes.idle.size() > 0) {
	find_subproblem_for_dpnode(ctrlstate.subproblems_q, true);
    }

    // Check the workers

    // Iterate through the subproblems, looking for one that can use
    // another worker.
    if (ctrlstate.workers.idle.size() > 0) {
	find_subproblem_for_worker(ctrlstate.subproblems_p);
    }
    if (ctrlstate.workers.idle.size() > 0) {
	find_subproblem_for_worker(ctrlstate.subproblems_q);
    }

    // cerr << "After schedule:\n"; ctrlstate.dump(cerr);
}

static ZZ computation_complete_p(const vector<SubproblemProgress> &v)
{
    ZZ curmodulus, curexp;
    curmodulus = 2;
    curexp = 0;

    vector<SubproblemProgress>::const_iterator vit;
    for(vit = v.begin(); vit != v.end(); ++vit) {
	CRT(curexp, curmodulus, vit->solution, vit->order);
    }

    if (curexp < 0) {
	curexp += curmodulus;
    }

    return curexp;
}

// All subproblems are solved.  Combine the results.
static void computation_complete(void)
{
    ZZ exp_p = computation_complete_p(ctrlstate.subproblems_p);
    ZZ exp_q = computation_complete_p(ctrlstate.subproblems_q);
    ZZ pm1 = (ctrlstate.p.factor - 1)/2;
    ZZ qm1 = (ctrlstate.q.factor - 1)/2;
    CRT(exp_p, pm1, exp_q, qm1);
    if (exp_p < 0) {
	exp_p += pm1;
    }
    ZZ& expon = exp_p;

    struct timeval ended_working;
    gettimeofday(&ended_working, NULL);
    unsigned long long computation_length_ms =
	(ended_working.tv_sec - ctrlstate.started_working.tv_sec) * 1000 +
	(ended_working.tv_usec - ctrlstate.started_working.tv_usec) / 1000;
    char length_buf[50];
    sprintf(length_buf, "%lld.%03lld s", computation_length_ms / 1000,
					    computation_length_ms % 1000);

    char timestamp[20];
    sprintf(timestamp, "%d.%06d",
	    (int)ended_working.tv_sec, (int)ended_working.tv_usec);

    cout << timestamp << ":" << output_prefix << ": " <<
	    "expon = " << expon << "\n";
    cout << timestamp << ":" << output_prefix << ": " <<
	    ctrlstate.worklist[0].first << " ";
    ZZ base_exp = PowerMod(ctrlstate.base, expon, ctrlstate.rho);
    if (base_exp == ctrlstate.target) {
	cout << "CORRECT in " << length_buf << "\n";
    } else {
	cout << "INCORRECT in " << length_buf << ":\n";
	cout << "base^exp = " << base_exp << "\n";
	cout << "target   = " << ctrlstate.target << "\n";
    }
    cout.flush();

    ctrlstate.reset();
}

static unsigned short cursubproblemid = 0;

// Take base and target mod f.factor, then decompose that into small
// subproblems given our knowledge of the factors of phi(f.factor)
static vector<SubproblemProgress> decomp(const ZZ_p &base, const ZZ_p &target,
    const FactorDecomp &f)
{
    vector<SubproblemProgress> ret;

    // Compute phi(factor)
    const int fveclen = f.fvec.length();
    ZZ phi = to_ZZ(2);
    for (int i = 0; i < fveclen; ++i) {
	phi *= f.fvec[i];
    }

    ZZ_pPush push; // This will backup and restore the ZZ_p at the end of the function
    ZZ_p::init(f.factor);
    for (int i = 0; i < fveclen; ++i) {
	const ZZ& order = f.fvec[i];
	ZZ quotient = phi / order;
	ZZ_p subgroup_base = to_ZZ_p(rep(base));
	subgroup_base = power(subgroup_base, quotient);
	ZZ_p subgroup_target = to_ZZ_p(rep(target));
	subgroup_target = power(subgroup_target, quotient);

	if (subgroup_base == 1) {
	    // The original base wasn't a generator of the whole group
	    if (subgroup_target == 1) {
		// But the target is in the subgroup.  Lucky us.
		continue;
	    } else {
		ret.clear();
		return ret;
	    }
	}

	++cursubproblemid;
	ret.push_back(SubproblemProgress(cursubproblemid-1, rep(subgroup_base),
				    rep(subgroup_target), f.factor, order));
    }

    return ret;
}

// Read the modulus (and the factorization of the modulus and its
// totient) from the given file.  "-" means cin.  Returns true if
// successful.
static bool read_modulus(const char *filename)
{
    if (strcmp(filename, "-")) {
	ifstream ins(filename);
	if (!ins.good()) return false;
	ins >> ctrlstate.rho >> ctrlstate.p.factor >> ctrlstate.p.fvec >>
	    ctrlstate.q.factor >> ctrlstate.q.fvec;
	ins.close();
    } else {
	cin >> ctrlstate.rho >> ctrlstate.p.factor >> ctrlstate.p.fvec >>
	    ctrlstate.q.factor >> ctrlstate.q.fvec;
    }
    return true;
}

static int generate_problem()
{
    // If there's already a problem on the go, don't generate another one
    if (ctrlstate.working == true) {
	return -1;
    }

    while (ctrlstate.worklist[0].second == 0) {
	ctrlstate.worklist.erase(ctrlstate.worklist.begin());
	if (ctrlstate.worklist.size() > 0) {
	    read_modulus(ctrlstate.worklist[0].first);
	} else {
	    break;
	}
    }

    // If there are no more problems to generate, close the listener
    if (ctrlstate.worklist.size() == 0) {
	evconnlistener_free(ctrlstate.listener);
	ctrlstate.dpnodes.free();
	ctrlstate.workers.free();
	return -1;
    }

    ctrlstate.working = true;

    int num_subproblems_p = 0;
    int num_subproblems_q = 0;

    // Generate a DLP mod rho (in the large odd-order subgroup)
    ZZ_p::init(ctrlstate.rho);
    unsigned int attempt = 0;
    do {
	ZZ_p base;
	ZZ_p target;
	{
#ifdef DERANDOMIZE
	    RandomStreamPush push_seed;
	    // the seed will be reset to its original value
	    // once we exit this scope
	    SetSeed(ctrlstate.rho*ctrlstate.p.factor*ctrlstate.q.factor+attempt);
#endif
	    base = power(random_ZZ_p(), 2);
	    target = power(random_ZZ_p(), 2);
	}
	gettimeofday(&ctrlstate.started_working, NULL);
	ctrlstate.base = rep(base);
	ctrlstate.target = rep(target);

	// Decompose it mod p and mod q
	ctrlstate.subproblems_p = decomp(base, target, ctrlstate.p);
	ctrlstate.subproblems_q = decomp(base, target, ctrlstate.q);
	num_subproblems_p = ctrlstate.subproblems_p.size();
	num_subproblems_q = ctrlstate.subproblems_q.size();

	attempt++;
    } while (num_subproblems_p == 0 || num_subproblems_q == 0);

    ctrlstate.num_unsolved_subproblems =
	    num_subproblems_p + num_subproblems_q;

    schedule();

    --(ctrlstate.worklist[0].second);

    return 0;
}

typedef enum {
    CCSTATE_START,
    CCSTATE_DPWAITRESP,
    CCSTATE_DPLISTENING,
    CCSTATE_DPEXPON,
    CCSTATE_END
} CCState;

struct ControllerConnInfo {
    CCState state;

    ControllerConnInfo() : state(CCSTATE_DPWAITRESP) {}
};

static void controller_dpnode_reader(struct bufferevent *bev, void *ctx)
{
    struct evbuffer *input = bufferevent_get_input(bev);
    ControllerConnInfo *info = (ControllerConnInfo *)ctx;
    unsigned char cmd[1];
    ZZ expon;

    while(1) {
	size_t len = evbuffer_get_length(input);
	switch (info->state) {
	    case CCSTATE_START:
	    case CCSTATE_DPWAITRESP:
		    if (len < 1) return;
		    bufferevent_read(bev, cmd, 1);
		    switch (cmd[0]) {
			case 'L':
			    info->state = CCSTATE_DPLISTENING;
			    break;
			case 'E':
			    info->state = CCSTATE_DPEXPON;
			    break;
			default:
			    /* Unknown DPnode command received */
			    fprintf(stderr, "Unknown command in "
				    "controller_dpnode_reader: "
				    "%c\n", cmd[0]);
			    info->state = CCSTATE_END;
			    break;
		    }
		    break;

	    case CCSTATE_DPLISTENING:
		// Read 6 bytes
		if (len < 6) return;
		unsigned char ipport[6];
		unsigned int DPip;
		unsigned short DPport;
		bufferevent_read(bev, ipport, 6);
		memmove(&DPip, ipport, 4);
		memmove(&DPport, ipport+4, 2);
		{
#ifdef VERBOSE
		    struct in_addr DPaddr = { DPip };
		    fprintf(stderr, "DP node at %s:%d\n", inet_ntoa(DPaddr), ntohs(DPport));
#endif
		    if (ctrlstate.dpnodes.working.count(bev) > 0) {
			ctrlstate.dpnodes.working[bev]->ipports.push_back(
			    IPPort(ipport));
			schedule();
		    }
		}
		info->state = CCSTATE_DPWAITRESP;
		break;

	    case CCSTATE_DPEXPON:
		// Read the subproblemid and the answer to the subproblem
		if (len < 2 + 3*sizeof(unsigned int)) return;
		unsigned char exponbytes[2 + 3*sizeof(unsigned int)];
		unsigned short problemid;
		bufferevent_read(bev, exponbytes, 2 + 3*sizeof(unsigned int));
		memmove(&problemid, exponbytes, 2);
		ZZFromBytes(expon, exponbytes+2, 3*sizeof(unsigned int));
		// Find the subproblem and check the answer
		if (ctrlstate.dpnodes.working.count(bev) > 0) {
		    SubproblemProgress *spp = ctrlstate.dpnodes.working[bev];
		    if (spp->problemid == problemid &&
			    spp->solved == false &&
			    spp->target ==
				PowerMod(spp->base, expon, spp->modulus)) {
			// Subproblem solved!
			ctrlstate.num_unsolved_subproblems--;
			spp->solution = expon;
			spp->solved = true;
			spp->stop();
			schedule();
		    }
		}
		info->state = CCSTATE_DPWAITRESP;
		break;

	    case CCSTATE_END:
		// Shut down the connection
		delete info;
		bufferevent_free(bev);
		return;
	}
    }

}


static void controller_worker_reader(struct bufferevent *bev, void *ctx)
{
    struct evbuffer *input = bufferevent_get_input(bev);
    ControllerConnInfo *info = (ControllerConnInfo *)ctx;

    while(1) {
	unsigned int kernel_launch_count = 0;
	unsigned long long worker_time_ms = 0;

	size_t len = evbuffer_get_length(input);
	if (len < sizeof(kernel_launch_count)+sizeof(worker_time_ms)) return;

	bufferevent_read(bev, &kernel_launch_count, sizeof(kernel_launch_count));
	bufferevent_read(bev, &worker_time_ms, sizeof(worker_time_ms));

	// this assumes that workers never leave a subproblem (never crash or get re-assigned)
	// otherwise we'll miss out on the kernel_launch_count for some workers

	SubproblemProgress *spp = ctrlstate.workers.working[bev];

	spp->kernel_launch_count += kernel_launch_count;
	spp->worker_time_ms += worker_time_ms;
	spp->num_workers_replied += 1;

	if (spp->num_workers_replied == spp->workers.size()) {

	    char worker_time_sec_buf[50];
	    sprintf(worker_time_sec_buf, "%lld.%03lld",
			spp->worker_time_ms / 1000,
			spp->worker_time_ms % 1000);

	    cout << "Timing (name, problemid, subproblemid, launches, worker_time): "
	         << ctrlstate.worklist[0].first << ", "
	         << ctrlstate.problemid << ", "
	         << spp->problemid << ", "
	         << spp->kernel_launch_count << ", "
	         << worker_time_sec_buf << "\n";
	    cout.flush();

	    spp->reset();
	    schedule();
	}

	if (ctrlstate.num_unsolved_subproblems == 0 && ctrlstate.workers.working.size() == 0) {
	    computation_complete();
	    generate_problem();
	}
    }
}

static void controller_dpnode_event_cb(struct bufferevent *bev, short events,
    void *ctx)
{
    if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR)) {
	ControllerConnInfo *info = (ControllerConnInfo*)ctx;
	fprintf(stderr, "Closing dpnode connection\n");
	if (ctrlstate.dpnodes.working.count(bev)) {
	    // If we lose a dpnode from an active computation, the
	    // computation is useless.
	    SubproblemProgress *spp = ctrlstate.dpnodes.working[bev];
	    ctrlstate.dpnodes.working.erase(bev);
	    spp->dpnodes.erase(bev);
	    spp->stop();
	    spp->reset();
	} else {
	    ctrlstate.dpnodes.idle.erase(bev);
	}
	delete info;
	bufferevent_free(bev);
	schedule();
    }
}

static void controller_worker_event_cb(struct bufferevent *bev, short events,
    void *ctx)
{
    if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR)) {
	ControllerConnInfo *info = (ControllerConnInfo*)ctx;
	fprintf(stderr, "Closing worker connection\n");
	if (ctrlstate.workers.working.count(bev)) {
	    SubproblemProgress *spp = ctrlstate.workers.working[bev];
	    ctrlstate.workers.working.erase(bev);
	    spp->workers.erase(bev);
	} else {
	    ctrlstate.workers.idle.erase(bev);
	}
	delete info;
	bufferevent_free(bev);
	schedule();
    }
}

static void controller_event_cb(struct bufferevent *bev, short events,
    void *ctx)
{
    if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR)) {
	fprintf(stderr, "Closing connection\n");
	ControllerConnInfo *info = (ControllerConnInfo*)ctx;
	delete info;
	bufferevent_free(bev);
    }
}

// We're just going to read a single byte that will tell us whether the
// peer is a DPnode or a Worker
static void controller_master_reader(struct bufferevent *bev, void *ctx)
{
    struct evbuffer *input = bufferevent_get_input(bev);

    size_t len = evbuffer_get_length(input);

    if (len < 1) return;

    char msg_id = 'I';
    char indata[1];
    bufferevent_read(bev, indata, 1);
    switch(indata[0]) {
	case 'D':
	    /* Add this DPnode to the list of available ones */
	    ctrlstate.dpnodes.idle.insert(bev);
	    bufferevent_setcb(bev, controller_dpnode_reader, NULL,
		    controller_dpnode_event_cb, ctx);
	    controller_dpnode_reader(bev, ctx);
	    bufferevent_write(bev, &msg_id, sizeof(msg_id));
	    bufferevent_write(bev, &ctrlstate.dpnode_count, sizeof(ctrlstate.dpnode_count));
	    ctrlstate.dpnode_count++;
	    schedule();
	    return;
	case 'W':
	    ctrlstate.workers.idle.insert(bev);
	    bufferevent_enable(bev, EV_WRITE);
	    bufferevent_setcb(bev, controller_worker_reader, NULL,
		    controller_worker_event_cb, ctx);
	    controller_worker_reader(bev, ctx);
	    schedule();
	    return;
	default:
	    fprintf(stderr, "Unknown command in controller_master_reader: "
		    "%c\n", indata[0]);

	    ControllerConnInfo *info = (ControllerConnInfo*)ctx;
	    delete info;
	    bufferevent_free(bev);
	    return;
    }
}

static void controller_accept_cb(struct evconnlistener *listener,
    evutil_socket_t fd, struct sockaddr *address, int socklen,
    void *ctx)
{
    // Create the state of the new connection
    ControllerConnInfo *info = new ControllerConnInfo();

    // Create a bufferevent for the new connection
    struct event_base *base = evconnlistener_get_base(listener);
    struct bufferevent *bev = bufferevent_socket_new(
	    base, fd, BEV_OPT_CLOSE_ON_FREE);

    bufferevent_setcb(bev, controller_master_reader, NULL,
	    controller_event_cb, info);

    bufferevent_enable(bev, EV_READ|EV_WRITE);
}

// Create a new controller socket.  bindport is the port to bind to (in
// host byte order), or 0 if any port will do.  ip and boundport are set
// to the IP and port of the socket, in network byte order.
static struct evconnlistener *controller_create(struct event_base *evbase,
    unsigned short bindport, unsigned int *ip, unsigned short *boundport)
{
    return listener_create(evbase, bindport, controller_accept_cb, NULL,
	ip, boundport, false);
}

int controller_parse_args(int argc, char **argv, unsigned short &bindport,
    Worklist &worklist, unsigned short &total_nodes,
    unsigned short &GB_mem_per_node)
{
    bindport = 0;
    total_nodes = 1;
    GB_mem_per_node = 1;
    int reps = 1;

    while (argc > 2 && argv[1][0] == '-') {
	if (!strncmp(argv[1], "-p", 2)) {
	    // A port number was specified
	    bindport = strtoul(argv[2], NULL, 10);
	    argc -= 2;
	    argv += 2;
	} else if (!strncmp(argv[1], "-n", 2)) {
	    total_nodes = strtoul(argv[2], NULL, 10);
	    argc -= 2;
	    argv += 2;
	} else if (!strncmp(argv[1], "-m", 2)) {
	    GB_mem_per_node = strtoul(argv[2], NULL, 10);
	    argc -= 2;
	    argv += 2;
	} else if (!strncmp(argv[1], "-r", 2)) {
	    reps = strtoul(argv[2], NULL, 10);
	    argc -= 2;
	    argv += 2;
	}
    }

    if (argc < 3 || (argc % 2 == 0)) {
	return 1;
    }

    for (int r=0;r<reps;++r) {
	for (int i=1;i<argc;i+=2) {
	    worklist.push_back(
		Workentry(argv[i], strtoul(argv[i+1], NULL, 10)));
	}
    }

    return 0;
}

int controller_main(const Worklist &worklist, unsigned short bindport,
    void (*boundcb)(const char *boundaddr, unsigned short boundport))
{
    // Initialize the prng with some randomness from the kernel
    unsigned char randbuf[1024];
    ifstream urand("/dev/urandom");
    urand.read((char *)randbuf, sizeof(randbuf));
    urand.close();
    ZZ randzz = ZZFromBytes(randbuf, sizeof(randbuf));
    SetSeed(randzz);

    if (worklist.size() == 0) {
	return 0;
    }

    // Initialize the output prefix
    char prefix[300];
    gethostname(prefix, 256);
    strcat(prefix, "-C");
    output_prefix = std::string(prefix);

    ctrlstate.worklist = worklist;

    // Read the modulus and the factorization of its totient from the
    // specified file
    if (!read_modulus(worklist[0].first)) {
	cerr << "Unable to read file " << worklist[0].first << "\n";
	return 1;
    }

    struct event_base *evbase = event_base_new();

    unsigned int myip;
    unsigned short myport;

    ctrlstate.listener = controller_create(evbase, bindport, &myip, &myport);
    if (ctrlstate.listener && boundcb) {
	struct in_addr myaddr = { myip };
	boundcb(inet_ntoa(myaddr), ntohs(myport));
    }

    // Kick off the first problem to solve
    generate_problem();

    event_base_dispatch(evbase);

    return 0;
}