controller.cc 25 KB

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