controller.cc 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865
  1. extern "C" {
  2. #include <event2/listener.h>
  3. #include <event2/bufferevent.h>
  4. #include <event2/buffer.h>
  5. }
  6. #include <NTL/vec_ZZ.h>
  7. #include <NTL/ZZ.h>
  8. #include <NTL/ZZ_p.h>
  9. #include <sys/socket.h>
  10. #include <netinet/in.h>
  11. #include <arpa/inet.h>
  12. #include <sys/time.h>
  13. #include <fstream>
  14. #include <vector>
  15. #include <set>
  16. #include <map>
  17. #include <string.h>
  18. #include "evutils.h"
  19. #include "subproblem.h"
  20. #include "controller.h"
  21. NTL_CLIENT
  22. // #undef VERBOSE
  23. // Used to prefix output from this process
  24. std::string output_prefix;
  25. struct SubproblemProgress;
  26. typedef std::set<struct bufferevent *> BESet;
  27. typedef std::map<struct bufferevent *, SubproblemProgress *> BEMap;
  28. static void besetdump(const BESet &bes, ostream &os)
  29. {
  30. BESet::const_iterator besit;
  31. os << hex << " ";
  32. for (besit = bes.begin(); besit != bes.end(); ++besit) {
  33. os << *besit << " ";
  34. }
  35. os << dec << "\n";
  36. }
  37. static void besetfree(BESet &bes)
  38. {
  39. BESet::iterator besit;
  40. for (besit = bes.begin(); besit != bes.end(); ++besit) {
  41. bufferevent_free(*besit);
  42. }
  43. bes.clear();
  44. }
  45. static void bemapdump(const BEMap &bem, ostream &os)
  46. {
  47. BEMap::const_iterator bemit;
  48. os << hex << " ";
  49. for (bemit = bem.begin(); bemit != bem.end(); ++bemit) {
  50. os << bemit->first << "->" << bemit->second << " ";
  51. }
  52. os << dec << "\n";
  53. }
  54. static void bemapfree(BEMap &bem)
  55. {
  56. BEMap::iterator bemit;
  57. for (bemit = bem.begin(); bemit != bem.end(); ++bemit) {
  58. bufferevent_free(bemit->first);
  59. }
  60. bem.clear();
  61. }
  62. struct Statuses {
  63. BESet idle;
  64. BEMap working;
  65. // Dump the state for debug purposes
  66. void dump(ostream &os) const {
  67. os << " idle (" << idle.size() << "):\n";
  68. besetdump(idle, os);
  69. os << " working (" << working.size() << "):\n";
  70. bemapdump(working, os);
  71. }
  72. void free(void) {
  73. besetfree(idle);
  74. bemapfree(working);
  75. }
  76. };
  77. struct FactorDecomp {
  78. ZZ factor;
  79. vec_ZZ fvec;
  80. };
  81. static void vsppdump(const vector<SubproblemProgress> &spv, ostream &os);
  82. static struct ControllerState {
  83. ZZ rho;
  84. FactorDecomp p, q;
  85. ZZ base, target;
  86. bool working;
  87. struct timeval started_working;
  88. vector<SubproblemProgress> subproblems_p, subproblems_q;
  89. Statuses dpnodes, workers;
  90. unsigned int num_unsolved_subproblems;
  91. Worklist worklist;
  92. struct evconnlistener *listener;
  93. ControllerState() : working(false), num_unsolved_subproblems(0),
  94. listener(NULL) {}
  95. // Reset the state for a new problem with the same modulus
  96. void reset(void) {
  97. base = 0;
  98. target = 0;
  99. working = false;
  100. subproblems_p.clear();
  101. subproblems_q.clear();
  102. num_unsolved_subproblems = 0;
  103. started_working.tv_sec = 0;
  104. started_working.tv_usec = 0;
  105. }
  106. // Dump the state for debug purposes
  107. void dump(ostream &os) const {
  108. if (!working) {
  109. os << "Not working\n";
  110. return;
  111. }
  112. os << "P:\n";
  113. vsppdump(subproblems_p, os);
  114. os << "Q:\n";
  115. vsppdump(subproblems_q, os);
  116. os << "dpnodes:\n";
  117. dpnodes.dump(os);
  118. os << "workers:\n";
  119. workers.dump(os);
  120. }
  121. } ctrlstate;
  122. struct IPPort {
  123. unsigned char ipport[6];
  124. IPPort(unsigned char *ipp) {
  125. memmove(ipport, ipp, 6);
  126. }
  127. void dump(ostream &os) const {
  128. os << int(ipport[0]) << "." << int(ipport[1]) << "." <<
  129. int(ipport[2]) << "." << int(ipport[3]) << ":" <<
  130. ((ipport[4] << 8) + ipport[5]) << " ";
  131. }
  132. };
  133. typedef vector<IPPort> IPPortSet;
  134. static void ipportsetdump(const IPPortSet &ipps, ostream &os)
  135. {
  136. IPPortSet::const_iterator ippsit;
  137. os << " ";
  138. for (ippsit = ipps.begin(); ippsit != ipps.end(); ++ippsit) {
  139. ippsit->dump(os);
  140. }
  141. os << "\n";
  142. }
  143. void desired_resources(const ZZ &order, unsigned short &desired_dpnodes,
  144. unsigned int &max_workers, unsigned int &dpfreq);
  145. struct SubproblemProgress : Subproblem {
  146. // The sets of dpnodes and workers currently working on this subproblem
  147. BESet dpnodes, workers;
  148. // The dpnode IPPorts registered for this subproblem
  149. IPPortSet ipports;
  150. // The desired number of DPnodes for this subproblem
  151. unsigned short desired_dpnodes;
  152. // The maximum number of workers useful for this subproblem
  153. unsigned int max_workers;
  154. // Have we found a solution?
  155. bool solved;
  156. // The solution, if found.
  157. ZZ solution;
  158. SubproblemProgress(unsigned short id, const ZZ &b, const ZZ &t,
  159. const ZZ &m, const ZZ &o) :
  160. // By default, 1 in 1000 points are distinguihed points. The
  161. // number in the next line is 2^32/1000
  162. Subproblem(id, b, t, m, o, 4294967), solved(false) {
  163. desired_resources(order, desired_dpnodes, max_workers, dpfreq);
  164. // We need to update the binary description of the subproblem,
  165. // as we may have just changed dpfreq.
  166. updatedesc();
  167. }
  168. // Stop all dpnodes and workers
  169. void stop(void) {
  170. BESet::iterator iter;
  171. unsigned char stopcmd[1] = { 'S' };
  172. for (BESet::iterator iter = dpnodes.begin(); iter != dpnodes.end();
  173. ++iter) {
  174. bufferevent_write(*iter, stopcmd, 1);
  175. ctrlstate.dpnodes.working.erase(*iter);
  176. ctrlstate.dpnodes.idle.insert(*iter);
  177. }
  178. for (BESet::iterator iter = workers.begin(); iter != workers.end();
  179. ++iter) {
  180. bufferevent_write(*iter, stopcmd, 1);
  181. ctrlstate.workers.working.erase(*iter);
  182. ctrlstate.workers.idle.insert(*iter);
  183. }
  184. dpnodes.clear();
  185. workers.clear();
  186. ipports.clear();
  187. }
  188. // Dump for debugging purposes
  189. void dump(ostream &os) const {
  190. os << " Subproblem " << problemid << "\n";
  191. os << " dpnodes (" << dpnodes.size() << "):\n";
  192. besetdump(dpnodes, os);
  193. os << " workers (" << workers.size() << "):\n";
  194. besetdump(workers, os);
  195. os << " ipports (" << ipports.size() << "):\n";
  196. ipportsetdump(ipports, os);
  197. if (solved) {
  198. os << " solution: " << solution << "\n\n";
  199. }
  200. }
  201. void worker_write(struct bufferevent *bev) {
  202. bev_write(bev);
  203. unsigned short num_ipports = ipports.size();
  204. bufferevent_write(bev, &num_ipports, 2);
  205. for (unsigned short i = 0; i < num_ipports; ++i) {
  206. bufferevent_write(bev, ipports[i].ipport, 6);
  207. }
  208. #ifdef VERBOSE
  209. cerr << "Added worker " << bev << " to subproblem "
  210. << problemid << "\n";
  211. #endif
  212. }
  213. };
  214. // Dump the state for debug purposes
  215. static void vsppdump(const vector<SubproblemProgress> &spv, ostream &os)
  216. {
  217. vector<SubproblemProgress>::const_iterator spiter;
  218. int count = 0;
  219. for (spiter = spv.begin(); spiter != spv.end(); ++spiter) {
  220. ++count;
  221. os << " " << count << ":\n";
  222. spiter->dump(os);
  223. }
  224. os << "\n";
  225. }
  226. // Find a subproblem in the given vector that could use another DPnode,
  227. // and give it one of the idle ones. Only allocate it to a subproblem
  228. // with no current DPnodes if consider_empty is true.
  229. static void find_subproblem_for_dpnode(vector<SubproblemProgress> &spv,
  230. bool consider_empty)
  231. {
  232. vector<SubproblemProgress>::iterator spiter;
  233. for (spiter = spv.begin(); spiter != spv.end(); ++spiter) {
  234. if (spiter->solved) continue;
  235. if (spiter->dpnodes.size() == 0 && consider_empty == false) continue;
  236. // How many DPnodes would we like to have for this subproblem?
  237. while (spiter->dpnodes.size() < spiter->desired_dpnodes &&
  238. ctrlstate.dpnodes.idle.size() > 0) {
  239. // Get the first idle DPnode
  240. BESet::iterator beviter = ctrlstate.dpnodes.idle.begin();
  241. struct bufferevent *firstbev = *beviter;
  242. // Allocate it to the subproblem
  243. spiter->dpnodes.insert(firstbev);
  244. ctrlstate.dpnodes.working[firstbev] = &(*spiter);
  245. ctrlstate.dpnodes.idle.erase(firstbev);
  246. // Tell it to start listening for DPs
  247. spiter->bev_write(firstbev);
  248. }
  249. }
  250. }
  251. // Find a subproblem in the given vector that has all of its DPnodes and
  252. // could use another worker, and give it one of the idle ones.
  253. static void find_subproblem_for_worker(vector<SubproblemProgress> &spv)
  254. {
  255. vector<SubproblemProgress>::iterator spiter;
  256. for (spiter = spv.begin(); spiter != spv.end(); ++spiter) {
  257. if (spiter->solved) continue;
  258. while (spiter->ipports.size() == spiter->desired_dpnodes &&
  259. spiter->workers.size() < spiter->max_workers &&
  260. ctrlstate.workers.idle.size() > 0) {
  261. // Get the first idle worker
  262. BESet::iterator beviter = ctrlstate.workers.idle.begin();
  263. struct bufferevent *firstbev = *beviter;
  264. // Allocate it to the subproblem
  265. spiter->workers.insert(firstbev);
  266. ctrlstate.workers.working[firstbev] = &(*spiter);
  267. ctrlstate.workers.idle.erase(firstbev);
  268. // Tell it to start working on the subproblem
  269. spiter->worker_write(firstbev);
  270. }
  271. }
  272. }
  273. // See if there are any idle DPnodes or workers we can put to use
  274. static void schedule(void)
  275. {
  276. // Check the DPnodes
  277. // Iterate through the subproblems, looking for one that can use
  278. // another DPnode. First look for subproblems that already have
  279. // some, but not all, of their DPnodes
  280. if (ctrlstate.dpnodes.idle.size() > 0) {
  281. find_subproblem_for_dpnode(ctrlstate.subproblems_p, false);
  282. }
  283. if (ctrlstate.dpnodes.idle.size() > 0) {
  284. find_subproblem_for_dpnode(ctrlstate.subproblems_q, false);
  285. }
  286. // If there are still more dpnodes to place, start assigning them to
  287. // subproblems with no current dpnodes
  288. if (ctrlstate.dpnodes.idle.size() > 0) {
  289. find_subproblem_for_dpnode(ctrlstate.subproblems_p, true);
  290. }
  291. if (ctrlstate.dpnodes.idle.size() > 0) {
  292. find_subproblem_for_dpnode(ctrlstate.subproblems_q, true);
  293. }
  294. // Check the workers
  295. // Iterate through the subproblems, looking for one that can use
  296. // another worker.
  297. if (ctrlstate.workers.idle.size() > 0) {
  298. find_subproblem_for_worker(ctrlstate.subproblems_p);
  299. }
  300. if (ctrlstate.workers.idle.size() > 0) {
  301. find_subproblem_for_worker(ctrlstate.subproblems_q);
  302. }
  303. // cerr << "After schedule:\n"; ctrlstate.dump(cerr);
  304. }
  305. static ZZ computation_complete_p(const vector<SubproblemProgress> &v)
  306. {
  307. ZZ curmodulus, curexp;
  308. curmodulus = 2;
  309. curexp = 0;
  310. vector<SubproblemProgress>::const_iterator vit;
  311. for(vit = v.begin(); vit != v.end(); ++vit) {
  312. CRT(curexp, curmodulus, vit->solution, vit->order);
  313. }
  314. if (curexp < 0) {
  315. curexp += curmodulus;
  316. }
  317. return curexp;
  318. }
  319. // All subproblems are solved. Combine the results.
  320. static void computation_complete(void)
  321. {
  322. ZZ exp_p = computation_complete_p(ctrlstate.subproblems_p);
  323. ZZ exp_q = computation_complete_p(ctrlstate.subproblems_q);
  324. ZZ pm1 = (ctrlstate.p.factor - 1)/2;
  325. ZZ qm1 = (ctrlstate.q.factor - 1)/2;
  326. CRT(exp_p, pm1, exp_q, qm1);
  327. if (exp_p < 0) {
  328. exp_p += pm1;
  329. }
  330. ZZ& expon = exp_p;
  331. struct timeval ended_working;
  332. gettimeofday(&ended_working, NULL);
  333. unsigned long long computation_length_ms =
  334. (ended_working.tv_sec - ctrlstate.started_working.tv_sec) * 1000 +
  335. (ended_working.tv_usec - ctrlstate.started_working.tv_usec) / 1000;
  336. char length_buf[50];
  337. sprintf(length_buf, "%lld.%03lld s", computation_length_ms / 1000,
  338. computation_length_ms % 1000);
  339. char timestamp[20];
  340. sprintf(timestamp, "%d.%06d",
  341. (int)ended_working.tv_sec, (int)ended_working.tv_usec);
  342. cout << timestamp << ":" << output_prefix << ": " <<
  343. "expon = " << expon << "\n";
  344. cout << timestamp << ":" << output_prefix << ": " <<
  345. ctrlstate.worklist[0].first << " ";
  346. ZZ base_exp = PowerMod(ctrlstate.base, expon, ctrlstate.rho);
  347. if (base_exp == ctrlstate.target) {
  348. cout << "CORRECT in " << length_buf << "\n";
  349. } else {
  350. cout << "INCORRECT in " << length_buf << ":\n";
  351. cout << "base^exp = " << base_exp << "\n";
  352. cout << "target = " << ctrlstate.target << "\n";
  353. }
  354. cout.flush();
  355. ctrlstate.reset();
  356. }
  357. static unsigned short curproblemid = 0;
  358. // Take base and target mod f.factor, then decompose that into small
  359. // subproblems given our knowledge of the factors of phi(f.factor)
  360. static vector<SubproblemProgress> decomp(const ZZ_p &base, const ZZ_p &target,
  361. const FactorDecomp &f)
  362. {
  363. vector<SubproblemProgress> ret;
  364. // Compute phi(factor)
  365. const int fveclen = f.fvec.length();
  366. ZZ phi = to_ZZ(2);
  367. for (int i = 0; i < fveclen; ++i) {
  368. phi *= f.fvec[i];
  369. }
  370. ZZ_p::init(f.factor);
  371. for (int i = 0; i < fveclen; ++i) {
  372. const ZZ& order = f.fvec[i];
  373. ZZ quotient = phi / order;
  374. ZZ_p subgroup_base = to_ZZ_p(rep(base));
  375. subgroup_base = power(subgroup_base, quotient);
  376. ZZ_p subgroup_target = to_ZZ_p(rep(target));
  377. subgroup_target = power(subgroup_target, quotient);
  378. if (subgroup_base == 1) {
  379. // The original base wasn't a generator of the whole group
  380. if (subgroup_target == 1) {
  381. // But the target is in the subgroup. Lucky us.
  382. continue;
  383. } else {
  384. ret.clear();
  385. return ret;
  386. }
  387. }
  388. ++curproblemid;
  389. ret.push_back(SubproblemProgress(curproblemid, rep(subgroup_base),
  390. rep(subgroup_target), f.factor, order));
  391. }
  392. return ret;
  393. }
  394. // Read the modulus (and the factorization of the modulus and its
  395. // totient) from the given file. "-" means cin. Returns true if
  396. // successful.
  397. static bool read_modulus(const char *filename)
  398. {
  399. if (strcmp(filename, "-")) {
  400. ifstream ins(filename);
  401. if (!ins.good()) return false;
  402. ins >> ctrlstate.rho >> ctrlstate.p.factor >> ctrlstate.p.fvec >>
  403. ctrlstate.q.factor >> ctrlstate.q.fvec;
  404. ins.close();
  405. } else {
  406. cin >> ctrlstate.rho >> ctrlstate.p.factor >> ctrlstate.p.fvec >>
  407. ctrlstate.q.factor >> ctrlstate.q.fvec;
  408. }
  409. return true;
  410. }
  411. static int generate_problem(struct event_base *evbase)
  412. {
  413. while (ctrlstate.worklist[0].second == 0) {
  414. ctrlstate.worklist.erase(ctrlstate.worklist.begin());
  415. if (ctrlstate.worklist.size() > 0) {
  416. read_modulus(ctrlstate.worklist[0].first);
  417. } else {
  418. break;
  419. }
  420. }
  421. // If there are no more problems to generate, close the listener
  422. if (ctrlstate.worklist.size() == 0) {
  423. evconnlistener_free(ctrlstate.listener);
  424. ctrlstate.dpnodes.free();
  425. ctrlstate.workers.free();
  426. return -1;
  427. }
  428. // If there's already a problem on the go, don't generate another one
  429. if (ctrlstate.working == true) {
  430. return -1;
  431. }
  432. ctrlstate.working = true;
  433. int num_subproblems_p = 0;
  434. int num_subproblems_q = 0;
  435. // Generate a DLP mod rho (in the large odd-order subgroup)
  436. ZZ_p::init(ctrlstate.rho);
  437. do {
  438. ZZ_p base = power(random_ZZ_p(), 2);
  439. ZZ_p target = power(random_ZZ_p(), 2);
  440. gettimeofday(&ctrlstate.started_working, NULL);
  441. ctrlstate.base = rep(base);
  442. ctrlstate.target = rep(target);
  443. // Decompose it mod p and mod q
  444. ctrlstate.subproblems_p = decomp(base, target, ctrlstate.p);
  445. ctrlstate.subproblems_q = decomp(base, target, ctrlstate.q);
  446. num_subproblems_p = ctrlstate.subproblems_p.size();
  447. num_subproblems_q = ctrlstate.subproblems_q.size();
  448. } while (num_subproblems_p == 0 || num_subproblems_q == 0);
  449. ctrlstate.num_unsolved_subproblems =
  450. num_subproblems_p + num_subproblems_q;
  451. schedule();
  452. --(ctrlstate.worklist[0].second);
  453. return 0;
  454. }
  455. typedef enum {
  456. CCSTATE_START,
  457. CCSTATE_DPWAITRESP,
  458. CCSTATE_DPLISTENING,
  459. CCSTATE_DPEXPON,
  460. CCSTATE_END
  461. } CCState;
  462. struct ControllerConnInfo {
  463. CCState state;
  464. ControllerConnInfo() : state(CCSTATE_DPWAITRESP) {}
  465. };
  466. static void controller_dpnode_reader(struct bufferevent *bev, void *ctx)
  467. {
  468. struct evbuffer *input = bufferevent_get_input(bev);
  469. ControllerConnInfo *info = (ControllerConnInfo *)ctx;
  470. unsigned char cmd[1];
  471. ZZ expon;
  472. while(1) {
  473. size_t len = evbuffer_get_length(input);
  474. switch (info->state) {
  475. case CCSTATE_START:
  476. case CCSTATE_DPWAITRESP:
  477. if (len < 1) return;
  478. bufferevent_read(bev, cmd, 1);
  479. switch (cmd[0]) {
  480. case 'L':
  481. info->state = CCSTATE_DPLISTENING;
  482. break;
  483. case 'E':
  484. info->state = CCSTATE_DPEXPON;
  485. break;
  486. default:
  487. /* Unknown DPnode command received */
  488. fprintf(stderr, "Unknown command in "
  489. "controller_dpnode_reader: "
  490. "%c\n", cmd[0]);
  491. info->state = CCSTATE_END;
  492. break;
  493. }
  494. break;
  495. case CCSTATE_DPLISTENING:
  496. // Read 6 bytes
  497. if (len < 6) return;
  498. unsigned char ipport[6];
  499. unsigned int DPip;
  500. unsigned short DPport;
  501. bufferevent_read(bev, ipport, 6);
  502. memmove(&DPip, ipport, 4);
  503. memmove(&DPport, ipport+4, 2);
  504. {
  505. #ifdef VERBOSE
  506. struct in_addr DPaddr = { DPip };
  507. fprintf(stderr, "DP node at %s:%d\n", inet_ntoa(DPaddr), ntohs(DPport));
  508. #endif
  509. if (ctrlstate.dpnodes.working.count(bev) > 0) {
  510. ctrlstate.dpnodes.working[bev]->ipports.push_back(
  511. IPPort(ipport));
  512. schedule();
  513. }
  514. }
  515. info->state = CCSTATE_DPWAITRESP;
  516. break;
  517. case CCSTATE_DPEXPON:
  518. // Read the subproblemid and the answer to the subproblem
  519. if (len < 2 + 3*sizeof(unsigned int)) return;
  520. unsigned char exponbytes[2 + 3*sizeof(unsigned int)];
  521. unsigned short problemid;
  522. bufferevent_read(bev, exponbytes, 2 + 3*sizeof(unsigned int));
  523. memmove(&problemid, exponbytes, 2);
  524. ZZFromBytes(expon, exponbytes+2, 3*sizeof(unsigned int));
  525. // Find the subproblem and check the answer
  526. if (ctrlstate.dpnodes.working.count(bev) > 0) {
  527. SubproblemProgress *spp = ctrlstate.dpnodes.working[bev];
  528. if (spp->problemid == problemid &&
  529. spp->solved == false &&
  530. spp->target ==
  531. PowerMod(spp->base, expon, spp->modulus)) {
  532. // Subproblem solved!
  533. spp->solution = expon;
  534. spp->solved = true;
  535. spp->stop();
  536. ctrlstate.num_unsolved_subproblems--;
  537. if (ctrlstate.num_unsolved_subproblems == 0) {
  538. computation_complete();
  539. generate_problem(bufferevent_get_base(bev));
  540. }
  541. schedule();
  542. }
  543. }
  544. info->state = CCSTATE_DPWAITRESP;
  545. break;
  546. case CCSTATE_END:
  547. // Shut down the connection
  548. delete info;
  549. bufferevent_free(bev);
  550. return;
  551. }
  552. }
  553. }
  554. static void controller_dpnode_event_cb(struct bufferevent *bev, short events,
  555. void *ctx)
  556. {
  557. if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR)) {
  558. ControllerConnInfo *info = (ControllerConnInfo*)ctx;
  559. fprintf(stderr, "Closing dpnode connection\n");
  560. if (ctrlstate.dpnodes.working.count(bev)) {
  561. // If we lose a dpnode from an active computation, the
  562. // computation is useless.
  563. SubproblemProgress *spp = ctrlstate.dpnodes.working[bev];
  564. ctrlstate.dpnodes.working.erase(bev);
  565. spp->dpnodes.erase(bev);
  566. spp->stop();
  567. } else {
  568. ctrlstate.dpnodes.idle.erase(bev);
  569. }
  570. delete info;
  571. bufferevent_free(bev);
  572. schedule();
  573. }
  574. }
  575. static void controller_worker_event_cb(struct bufferevent *bev, short events,
  576. void *ctx)
  577. {
  578. if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR)) {
  579. ControllerConnInfo *info = (ControllerConnInfo*)ctx;
  580. fprintf(stderr, "Closing worker connection\n");
  581. if (ctrlstate.workers.working.count(bev)) {
  582. SubproblemProgress *spp = ctrlstate.workers.working[bev];
  583. ctrlstate.workers.working.erase(bev);
  584. spp->workers.erase(bev);
  585. } else {
  586. ctrlstate.workers.idle.erase(bev);
  587. }
  588. delete info;
  589. bufferevent_free(bev);
  590. schedule();
  591. }
  592. }
  593. static void controller_event_cb(struct bufferevent *bev, short events,
  594. void *ctx)
  595. {
  596. if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR)) {
  597. fprintf(stderr, "Closing connection\n");
  598. ControllerConnInfo *info = (ControllerConnInfo*)ctx;
  599. delete info;
  600. bufferevent_free(bev);
  601. }
  602. }
  603. // We're just going to read a single byte that will tell us whether the
  604. // peer is a DPnode or a Worker
  605. static void controller_master_reader(struct bufferevent *bev, void *ctx)
  606. {
  607. struct evbuffer *input = bufferevent_get_input(bev);
  608. size_t len = evbuffer_get_length(input);
  609. if (len < 1) return;
  610. char indata[1];
  611. bufferevent_read(bev, indata, 1);
  612. switch(indata[0]) {
  613. case 'D':
  614. /* Add this DPnode to the list of available ones */
  615. ctrlstate.dpnodes.idle.insert(bev);
  616. bufferevent_setcb(bev, controller_dpnode_reader, NULL,
  617. controller_dpnode_event_cb, ctx);
  618. controller_dpnode_reader(bev, ctx);
  619. schedule();
  620. return;
  621. case 'W':
  622. ctrlstate.workers.idle.insert(bev);
  623. // We don't actually read anything from workers
  624. bufferevent_enable(bev, EV_WRITE);
  625. bufferevent_setcb(bev, NULL, NULL,
  626. controller_worker_event_cb, ctx);
  627. schedule();
  628. return;
  629. default:
  630. fprintf(stderr, "Unknown command in controller_master_reader: "
  631. "%c\n", indata[0]);
  632. ControllerConnInfo *info = (ControllerConnInfo*)ctx;
  633. delete info;
  634. bufferevent_free(bev);
  635. return;
  636. }
  637. }
  638. static void controller_accept_cb(struct evconnlistener *listener,
  639. evutil_socket_t fd, struct sockaddr *address, int socklen,
  640. void *ctx)
  641. {
  642. // Create the state of the new connection
  643. ControllerConnInfo *info = new ControllerConnInfo();
  644. // Create a bufferevent for the new connection
  645. struct event_base *base = evconnlistener_get_base(listener);
  646. struct bufferevent *bev = bufferevent_socket_new(
  647. base, fd, BEV_OPT_CLOSE_ON_FREE);
  648. bufferevent_setcb(bev, controller_master_reader, NULL,
  649. controller_event_cb, info);
  650. bufferevent_enable(bev, EV_READ|EV_WRITE);
  651. }
  652. // Create a new controller socket. bindport is the port to bind to (in
  653. // host byte order), or 0 if any port will do. ip and boundport are set
  654. // to the IP and port of the socket, in network byte order.
  655. static struct evconnlistener *controller_create(struct event_base *evbase,
  656. unsigned short bindport, unsigned int *ip, unsigned short *boundport)
  657. {
  658. return listener_create(evbase, bindport, controller_accept_cb, NULL,
  659. ip, boundport, false);
  660. }
  661. int controller_parse_args(int argc, char **argv, unsigned short &bindport,
  662. Worklist &worklist, unsigned short &total_nodes,
  663. unsigned short &GB_mem_per_node)
  664. {
  665. bindport = 0;
  666. total_nodes = 1;
  667. GB_mem_per_node = 1;
  668. int reps = 1;
  669. while (argc > 2 && argv[1][0] == '-') {
  670. if (!strncmp(argv[1], "-p", 2)) {
  671. // A port number was specified
  672. bindport = strtoul(argv[2], NULL, 10);
  673. argc -= 2;
  674. argv += 2;
  675. } else if (!strncmp(argv[1], "-n", 2)) {
  676. total_nodes = strtoul(argv[2], NULL, 10);
  677. argc -= 2;
  678. argv += 2;
  679. } else if (!strncmp(argv[1], "-m", 2)) {
  680. GB_mem_per_node = strtoul(argv[2], NULL, 10);
  681. argc -= 2;
  682. argv += 2;
  683. } else if (!strncmp(argv[1], "-r", 2)) {
  684. reps = strtoul(argv[2], NULL, 10);
  685. argc -= 2;
  686. argv += 2;
  687. }
  688. }
  689. if (argc < 3 || (argc % 2 == 0)) {
  690. return 1;
  691. }
  692. for (int r=0;r<reps;++r) {
  693. for (int i=1;i<argc;i+=2) {
  694. worklist.push_back(
  695. Workentry(argv[i], strtoul(argv[i+1], NULL, 10)));
  696. }
  697. }
  698. return 0;
  699. }
  700. int controller_main(const Worklist &worklist, unsigned short bindport,
  701. void (*boundcb)(const char *boundaddr, unsigned short boundport))
  702. {
  703. // Initialize the prng with some randomness from the kernel
  704. unsigned char randbuf[1024];
  705. ifstream urand("/dev/urandom");
  706. urand.read((char *)randbuf, sizeof(randbuf));
  707. urand.close();
  708. ZZ randzz = ZZFromBytes(randbuf, sizeof(randbuf));
  709. SetSeed(randzz);
  710. if (worklist.size() == 0) {
  711. return 0;
  712. }
  713. // Initialize the output prefix
  714. char prefix[300];
  715. gethostname(prefix, 256);
  716. strcat(prefix, "-C");
  717. output_prefix = std::string(prefix);
  718. ctrlstate.worklist = worklist;
  719. // Read the modulus and the factorization of its totient from the
  720. // specified file
  721. if (!read_modulus(worklist[0].first)) {
  722. cerr << "Unable to read file " << worklist[0].first << "\n";
  723. return 1;
  724. }
  725. struct event_base *evbase = event_base_new();
  726. unsigned int myip;
  727. unsigned short myport;
  728. ctrlstate.listener = controller_create(evbase, bindport, &myip, &myport);
  729. if (ctrlstate.listener && boundcb) {
  730. struct in_addr myaddr = { myip };
  731. boundcb(inet_ntoa(myaddr), ntohs(myport));
  732. }
  733. // Kick off the first problem to solve
  734. generate_problem(evbase);
  735. event_base_dispatch(evbase);
  736. return 0;
  737. }