cudadl/worker.cc

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

#include <vector>
#include <stdio.h>

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

typedef enum {
    WRKCCSTATE_AWAITCMD,
    WRKCCSTATE_RDPROBLEM,
    WRKCCSTATE_RDDPNODES,
    WRKCCSTATE_END
} WrkCCState;

struct WrkControllerConnInfo {
    WrkCCState state;
    unsigned char subproblem[SUBPROBLEM_DESC_LEN];
    unsigned short num_dpnodes;

    WrkControllerConnInfo() : state(WRKCCSTATE_AWAITCMD) {}
};

static struct WrkControllerState {
    Subproblem *current_problem;
    unsigned short num_expected_dpnodes;
    vector<struct bufferevent *> dpnodes;
    unsigned short num_connected_dpnodes;

    WrkControllerState(): current_problem(NULL) {}
} wrkctrlstate;

static void dpconn_event_cb(struct bufferevent *bev, short events,
    void *ctx)
{
    if (events & BEV_EVENT_CONNECTED) {
	// We have successfully connected to the dpnode
	++wrkctrlstate.num_connected_dpnodes;
	if (wrkctrlstate.num_connected_dpnodes ==
		wrkctrlstate.num_expected_dpnodes) {
	    cerr << "Starting work\n";
	    // start_working();
	}
    } else if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR)) {
	fprintf(stderr, "Closing connection to controller and restarting\n");
	bufferevent_free(bev);
	cerr << "Stopping work\n";
	// stop_working();
    }
}

static void controllerconn_reader(struct bufferevent *bev, void *ctx)
{
    struct evbuffer *input = bufferevent_get_input(bev);
    WrkControllerConnInfo *info = (WrkControllerConnInfo *)ctx;
    unsigned char cmd[1];

    while(1) {
	size_t len = evbuffer_get_length(input);
	switch(info->state) {
	    case WRKCCSTATE_AWAITCMD:
		if (len < 1) return;
		bufferevent_read(bev, cmd, 1);
		switch(cmd[0]) {
		    case 'P':
			info->state = WRKCCSTATE_RDPROBLEM;
			break;
		    case 'S':
			// stop_working();
			break;
		    default:
			/* Unknown command received */
			fprintf(stderr, "Unknown command in "
				"controllerconn_reader: %c\n", cmd[0]);
			info->state = WRKCCSTATE_END;
			break;
		}
		break;

	    case WRKCCSTATE_RDPROBLEM:
		if (len < SUBPROBLEM_DESC_LEN+2) return;
		// stop_working();
		bufferevent_read(bev, info->subproblem, SUBPROBLEM_DESC_LEN);
		bufferevent_read(bev, &(info->num_dpnodes), 2);
		info->state = WRKCCSTATE_RDDPNODES;
		/* FALLTHROUGH */

	    case WRKCCSTATE_RDDPNODES:
		if (len < 6*(info->num_dpnodes)) return;
		wrkctrlstate.num_expected_dpnodes = info->num_dpnodes;
		{
		    unsigned short i;
		    for(i=0;i<info->num_dpnodes;++i) {
			unsigned char ipport[6];
			bufferevent_read(bev, ipport, 6);
			// XXX: Start a connection to this DPnode
			cerr << "Connecting to DPnode " <<
				    int(ipport[0]) << "." <<
				    int(ipport[1]) << "." <<
			            int(ipport[2]) << "." <<
				    int(ipport[3]) << ":" <<
				    ((ipport[4] << 8) + ipport[5]) <<
				    "\n";
			struct bufferevent *dpbev = client_create(
				bufferevent_get_base(bev), ipport,
				dpconn_event_cb);
			if (dpbev) {
			    wrkctrlstate.dpnodes.push_back(dpbev);
			} else {
			    // stop_working();
			}
		    }
		}
		info->state = WRKCCSTATE_AWAITCMD;
		break;

	    case WRKCCSTATE_END:
		// Shut down
		delete info;
		event_base_loopbreak(bufferevent_get_base(bev));
		bufferevent_free(bev);
		return;
	}
    }
}

static void controllerconn_event_cb(struct bufferevent *bev, short events,
    void *ctx)
{
    if (events & BEV_EVENT_CONNECTED) {
	// We have successfully connected to the controller
	char id[1] = { 'W' };

        bufferevent_enable(bev, EV_READ|EV_WRITE);
	bufferevent_write(bev, id, 1);
	bufferevent_setcb(bev, controllerconn_reader, NULL,
		controllerconn_event_cb, new WrkControllerConnInfo());
    } else if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR)) {
	fprintf(stderr, "Closing connection to controller and exiting\n");
	event_base_loopbreak(bufferevent_get_base(bev));
	bufferevent_free(bev);
    }
}

int main(int argc, char **argv)
{
    if (argc != 3) {
	fprintf(stderr, "Usage: %s controller_host controller_port\n", argv[0]);
	return 1;
    }

    unsigned short controller_port = strtoul(argv[2], NULL, 10);

    return controller_client(argv[1], controller_port,
				controllerconn_event_cb);
}