PageRenderTime 67ms CodeModel.GetById 37ms RepoModel.GetById 0ms app.codeStats 0ms

/network/zmq/znc/2/main.cc

https://bitbucket.org/kbw/wb-samples
C++ | 522 lines | 388 code | 87 blank | 47 comment | 53 complexity | 78356a7ae3a827c81595c9d301b7ff51 MD5 | raw file
Possible License(s): BSD-3-Clause, MPL-2.0-no-copyleft-exception
  1. #include <zmq.hpp>
  2. #include <unistd.h>
  3. #include <string>
  4. #include <deque>
  5. #include <tuple>
  6. #include <vector>
  7. #include <fstream>
  8. #include <stdexcept>
  9. #include <iostream>
  10. //=============================================================================
  11. //
  12. extern std::ofstream trace;
  13. //=============================================================================
  14. //
  15. class Pool;
  16. class Bundle;
  17. class Stream;
  18. typedef std::deque<std::string> strings_t;
  19. //=============================================================================
  20. //
  21. std::ofstream trace("/tmp/out.log", std::ios::app);
  22. //=============================================================================
  23. //
  24. class Stream
  25. {
  26. public:
  27. struct handle_t {
  28. handle_t(zmq::socket_t* zsock);
  29. handle_t(int fd);
  30. zmq::socket_t* zsock;
  31. int fd;
  32. };
  33. typedef uint16_t mode_t;
  34. static const mode_t eNone = 0;
  35. static const mode_t eInput = 1;
  36. static const mode_t eOutput = 2;
  37. public:
  38. Stream(zmq::socket_t* zsock, const std::string& name, mode_t d, mode_t m);
  39. Stream(int fd, const std::string& name, mode_t d, mode_t m);
  40. ~Stream();
  41. const std::string& name() const;
  42. mode_t mode() const;
  43. void mode(mode_t m);
  44. int getHandle() const;
  45. zmq::socket_t* getSocket() const;
  46. Bundle* getBundle() const;
  47. void setParent(Bundle* parent);
  48. private:
  49. std::string m_name; // application name for stream
  50. mode_t m_direction; // allowed directions
  51. mode_t m_mode; // enabled directions
  52. handle_t m_handle; // object handle
  53. Bundle* m_parent; // owning bundle
  54. };
  55. //=============================================================================
  56. //
  57. Stream::handle_t::handle_t(zmq::socket_t* zsock) :
  58. zsock(zsock),
  59. fd(-1)
  60. {
  61. }
  62. Stream::handle_t::handle_t(int fd) :
  63. zsock(nullptr),
  64. fd(fd)
  65. {
  66. }
  67. Stream::Stream(zmq::socket_t* zsock, const std::string& name, Stream::mode_t d, Stream::mode_t m) :
  68. m_name(name),
  69. m_direction(d),
  70. m_mode(m),
  71. m_handle(zsock),
  72. m_parent(nullptr)
  73. {
  74. }
  75. Stream::Stream(int fd, const std::string& name, Stream::mode_t d, Stream::mode_t m) :
  76. m_name(name),
  77. m_direction(d),
  78. m_mode(m),
  79. m_handle(fd),
  80. m_parent(nullptr)
  81. {
  82. }
  83. Stream::~Stream()
  84. {
  85. }
  86. const std::string& Stream::name() const
  87. {
  88. return m_name;
  89. }
  90. Stream::mode_t Stream::mode() const
  91. {
  92. return m_mode;
  93. }
  94. void Stream::mode(Stream::mode_t m)
  95. {
  96. if ((m_direction & Stream::eInput) && (m_mode & Stream::eInput)) {
  97. m_mode = m;
  98. return;
  99. }
  100. if ((m_direction & eOutput) && (m_mode & eOutput)) {
  101. m_mode = m;
  102. return;
  103. }
  104. throw std::runtime_error("bad mode in Stream::mode(mode_t)");
  105. }
  106. int Stream::getHandle() const
  107. {
  108. return m_handle.fd;
  109. }
  110. zmq::socket_t* Stream::getSocket() const
  111. {
  112. return m_handle.zsock;
  113. }
  114. Bundle* Stream::getBundle() const
  115. {
  116. return m_parent;
  117. }
  118. void Stream::setParent(Bundle* parent)
  119. {
  120. m_parent = parent;
  121. }
  122. //=============================================================================
  123. //
  124. class Bundle
  125. {
  126. public:
  127. typedef Stream* stream_t;
  128. typedef std::vector<stream_t> streams_t;
  129. public:
  130. Bundle(const std::string& name);
  131. ~Bundle();
  132. const std::string& name() const;
  133. void add(Stream* e);
  134. void remove(Stream* e);
  135. streams_t getStreams(Stream::mode_t m) const;
  136. private:
  137. const std::string& m_name;
  138. streams_t m_streams; // streams within the bundle
  139. };
  140. //-----------------------------------------------------------------------------
  141. //
  142. Bundle::Bundle(const std::string& name) :
  143. m_name(name)
  144. {
  145. }
  146. Bundle::~Bundle()
  147. {
  148. for (auto p = m_streams.rbegin(); p != m_streams.rend(); ++p)
  149. delete *p;
  150. }
  151. const std::string& Bundle::name() const
  152. {
  153. return m_name;
  154. }
  155. void Bundle::add(Stream* e)
  156. {
  157. m_streams.push_back(e);
  158. e->setParent(this);
  159. }
  160. void Bundle::remove(Stream* e)
  161. {
  162. for (auto p = m_streams.begin(); p != m_streams.end(); ++p)
  163. if (*p == e) {
  164. m_streams.erase(std::remove(m_streams.begin(), m_streams.end(), e), m_streams.end());
  165. delete e;
  166. break;
  167. }
  168. }
  169. Bundle::streams_t Bundle::getStreams(Stream::mode_t m) const
  170. {
  171. streams_t streams;
  172. for (const stream_t& s : m_streams)
  173. if (s) {
  174. if (static_cast<int>(s->mode()) & static_cast<int>(m))
  175. streams.push_back(s);
  176. }
  177. return streams;
  178. }
  179. //=============================================================================
  180. //
  181. class Pool
  182. {
  183. public:
  184. typedef std::vector<zmq_pollitem_t> zmq_pollitems_t;
  185. typedef ::zmq_pollitem_t zmq_pollitem_t;
  186. typedef std::vector<Bundle*> bundles_t;
  187. typedef std::vector<Stream*> streams_t;
  188. typedef Stream* stream_t;
  189. public:
  190. Pool();
  191. ~Pool();
  192. void add(Bundle* bundle);
  193. std::tuple<zmq_pollitems_t, streams_t> getItems(Stream::mode_t mode) const;
  194. private:
  195. bundles_t m_bundles; // bundles
  196. };
  197. //-----------------------------------------------------------------------------
  198. //
  199. Pool::Pool()
  200. {
  201. }
  202. Pool::~Pool()
  203. {
  204. for (auto p = m_bundles.rbegin(); p != m_bundles.rend(); ++p)
  205. delete *p;
  206. }
  207. void Pool::add(Bundle* bundle)
  208. {
  209. m_bundles.push_back(bundle);
  210. }
  211. std::tuple<Pool::zmq_pollitems_t, Pool::streams_t> Pool::getItems(Stream::mode_t mode) const
  212. {
  213. zmq_pollitems_t items;
  214. streams_t streams;
  215. for (Bundle* pbundle : m_bundles)
  216. for (Stream* pstream : pbundle->getStreams(mode)) {
  217. streams.push_back(pstream);
  218. short zmq_mode = 0;
  219. if (pstream->mode() & Stream::eInput)
  220. zmq_mode |= ZMQ_POLLIN;
  221. if (pstream->mode() & Stream::eOutput)
  222. zmq_mode |= ZMQ_POLLOUT;
  223. if (zmq_mode == 0) {
  224. trace << " ignoring streeam: " << pbundle->name() << ":" << pstream->name() << std::endl;
  225. continue;
  226. }
  227. if (zmq::socket_t* zsock = pstream->getSocket()) {
  228. zmq_pollitem_t rec = { static_cast<void*>(zsock), -1, zmq_mode, 0 };
  229. items.push_back(rec);
  230. trace
  231. << " add stream: " << pbundle->name() << ":" << pstream->name() << " for "
  232. << ((zmq_mode & ZMQ_POLLIN) ? "input" :
  233. (zmq_mode & ZMQ_POLLOUT) ? "output" :
  234. (zmq_mode & (ZMQ_POLLIN|ZMQ_POLLOUT)) ? "io" : "unknown")
  235. << std::endl;
  236. }
  237. else {
  238. zmq_pollitem_t rec = { nullptr, pstream->getHandle(), zmq_mode, 0 };
  239. items.push_back(rec);
  240. trace
  241. << " add stream: " << pbundle->name() << ":" << pstream->name() << " for "
  242. << ((zmq_mode & ZMQ_POLLIN) ? "input" :
  243. (zmq_mode & ZMQ_POLLOUT) ? "output" :
  244. (zmq_mode & (ZMQ_POLLIN|ZMQ_POLLOUT)) ? "io" : "unknown")
  245. << std::endl;
  246. }
  247. }
  248. return std::make_tuple(items, streams);
  249. }
  250. namespace
  251. {
  252. bool reset(const strings_t& msgs, Stream::mode_t& mode)
  253. {
  254. // we always want read
  255. Stream::mode_t working_mode = Stream::eInput;
  256. if (msgs.empty())
  257. working_mode &= ~Stream::eOutput;
  258. else
  259. working_mode |= Stream::eOutput;
  260. // calculate mode change
  261. bool changed = working_mode != mode;
  262. mode = working_mode;
  263. trace << " mode: eInput=" << ((mode & Stream::eInput) ? "on" : "off") << " eOutput=" << ((mode & Stream::eOutput) ? "on" : "off") << std::endl;
  264. return changed;
  265. }
  266. void parse_stream_buffer(
  267. strings_t& msgs,
  268. bool more,
  269. std::unique_ptr<char[]>& buf,
  270. std::unique_ptr<char[]>& spare_buf, char*& spare_buf_ptr)
  271. {
  272. char* p = buf.get();
  273. for (char* q; (q = strchr(p, '\n')); p = q + 1) {
  274. if (p != q) {
  275. msgs.emplace_back(p, q);
  276. trace << " msg=" << msgs.back() << std::endl;
  277. }
  278. }
  279. if (more) {
  280. spare_buf.swap(buf);
  281. spare_buf_ptr = p;
  282. }
  283. else {
  284. if (*p) {
  285. msgs.emplace_back(p);
  286. trace << " msg=" << msgs.back() << std::endl;
  287. }
  288. }
  289. trace << " msgs.size=" << msgs.size() << std::endl;
  290. }
  291. bool read_fd(
  292. strings_t& msgs,
  293. Stream::mode_t mode, Pool* pool, Pool::zmq_pollitems_t& items, Pool::streams_t& streams, size_t i)
  294. {
  295. const size_t bufsz = 4*1024; // one page
  296. std::unique_ptr<char[]> buf; // main buffer
  297. std::unique_ptr<char[]> spare_buf; // overflow buffer
  298. char* spare_buf_ptr = nullptr; // overflow buffer offset
  299. bool more;
  300. do {
  301. more = false;
  302. // allocate the buffer
  303. if (!buf.get()) {
  304. buf.reset(new char[bufsz]);
  305. trace << " allocate buf: bufsz=" << bufsz << std::endl;
  306. }
  307. bzero(buf.get(), bufsz);
  308. // read the device
  309. int n = read(items[i].fd, buf.get(), bufsz);
  310. trace << " read=" << n << std::endl;
  311. if (n == 0) {
  312. // end of stream
  313. trace << " closing fd=" << items[i].fd << std::endl;
  314. close(items[i].fd);
  315. streams[i]->getBundle()->remove(streams[i]);
  316. if (more) {
  317. trace << " flushing buffer" << std::endl;
  318. spare_buf.swap(buf);
  319. parse_stream_buffer(msgs, more, buf, spare_buf, spare_buf_ptr);
  320. }
  321. std::tie(items, streams) = pool->getItems(mode | Stream::eInput); // reinitialize pool items
  322. trace << " " << items.size() << " items" << std::endl;
  323. trace << " " << streams.size() << " streams" << std::endl;
  324. return false;
  325. }
  326. trace << " processing buffer" << std::endl;
  327. more = n == bufsz; // do we think there's more to read?
  328. parse_stream_buffer(msgs, more, buf, spare_buf, spare_buf_ptr);
  329. }
  330. while (more);
  331. return true;
  332. }
  333. }
  334. int poll(zmq_pollitem_t* data, size_t size, int ms_sleep)
  335. {
  336. int rc = zmq_poll(data, size, ms_sleep);
  337. trace << " zmq_poll=" << rc << std::endl;
  338. return rc;
  339. }
  340. Pool* create_pool(zmq::context_t& ctx)
  341. {
  342. Pool* pool = new Pool;
  343. if (Bundle* reader_bundle = new Bundle("reader")) {
  344. zmq::socket_t* pub = new zmq::socket_t(ctx, zmq::socket_type::pub);
  345. pub->bind("inproc://news.znc");
  346. reader_bundle->add(new Stream(pub, "repeater", Stream::eOutput, Stream::eOutput));
  347. reader_bundle->add(new Stream(STDIN_FILENO, "stdin", Stream::eInput, Stream::eInput));
  348. pool->add(reader_bundle);
  349. }
  350. if (Bundle* writer_bundle = new Bundle("writer")) {
  351. zmq::socket_t* sub = new zmq::socket_t(ctx, zmq::socket_type::sub);
  352. sub->connect("inproc://news.znc");
  353. sub->setsockopt(ZMQ_SUBSCRIBE, "", 0); // subscribe to everything
  354. writer_bundle->add(new Stream(sub, "receiver", Stream::eInput, Stream::eInput));
  355. writer_bundle->add(new Stream(STDOUT_FILENO, "stdout", Stream::eOutput, Stream::eOutput));
  356. pool->add(writer_bundle);
  357. }
  358. return pool;
  359. }
  360. int main()
  361. try
  362. {
  363. trace << "begin" << std::endl;
  364. zmq::context_t ctx(1);
  365. zmq_ctx_set(static_cast<void*>(ctx), ZMQ_IO_THREADS, 4);
  366. std::unique_ptr<Pool> pool( create_pool(ctx) );
  367. strings_t msgs; // messages read to be written
  368. strings_t msgs2;
  369. Stream::mode_t mode = 0; // stream i/o mode
  370. Pool::zmq_pollitems_t items;// parallel array: poll items
  371. Pool::streams_t streams; // parallel array: streams
  372. bool stop = false;
  373. while (!stop) {
  374. // reset poolitems to wait for
  375. if (reset(msgs, mode)) {
  376. std::tie(items, streams) = pool->getItems(mode); // initialize pool items
  377. trace << " " << items.size() << " items" << std::endl;
  378. trace << " " << streams.size() << " streams" << std::endl;
  379. }
  380. // wait for poolitems
  381. int rc = poll(items.data(), items.size(), 1 * 1000);
  382. // handle what fired
  383. bool check_items = rc > 0;
  384. for (size_t i = 0; check_items && i != items.size(); ++i) {
  385. // handle read
  386. if (items[i].revents & ZMQ_POLLIN) {
  387. trace << " stream(" << i << ") " << streams[i]->name() << " fired for input" << std::endl;
  388. if (streams[i]->getSocket()) {
  389. // handle read socket
  390. zmq::message_t msg;
  391. streams[i]->getSocket()->recv(&msg);
  392. msgs2.push_back(std::string(static_cast<const char*>(msg.data()), msg.size()));
  393. }
  394. else if (items[i].fd != -1) {
  395. // handle read fd
  396. if (!read_fd(msgs, mode, pool.get(), items, streams, i))
  397. break; // refresh items fd is closed
  398. }
  399. }
  400. // handle write
  401. if (items[i].revents & ZMQ_POLLOUT) {
  402. trace << " stream(" << i << ") " << streams[i]->name() << " fired for output" << std::endl;
  403. if (streams[i]->getSocket()) {
  404. // handle write socket
  405. zmq::message_t msg(msgs.front().c_str(), msgs.front().size());
  406. streams[i]->getSocket()->send(msg);
  407. trace << " msg=" << msgs.front() << std::endl;
  408. msgs.pop_front();
  409. }
  410. /* 9 27 19 67 50 26
  411. else if (items[i].fd != -1) {
  412. // handle write fd
  413. if (streams[i]->name() == "stdout") {
  414. for (int n, offset = 0; offset < (int)msgs2.front().size(); offset += n) {
  415. n = write(items[i].fd, msgs2.front().c_str() + offset, msgs2.front().size() - offset);
  416. if (n <= 0)
  417. break;
  418. }
  419. msgs2.pop_front();
  420. if (msgs2.empty())
  421. std::tie(items, streams) = pool->getItems(mode | Stream::eInput); // reinitialize pool items
  422. }
  423. }
  424. */
  425. }
  426. }
  427. stop = rc == 0;
  428. }
  429. trace << "done\n" << std::endl;
  430. }
  431. catch (const std::exception& e)
  432. {
  433. std::clog << "fatal: " << e.what() << "\n" << std::endl;
  434. }