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

#include <zmq.hpp>
#include <unistd.h>
#include <string>
#include <deque>
#include <tuple>
#include <vector>
#include <fstream>
#include <stdexcept>
#include <iostream>

//=============================================================================
//
extern std::ofstream	trace;

//=============================================================================
//
class Pool;
class Bundle;
class Stream;
typedef std::deque<std::string> strings_t;

//=============================================================================
//
std::ofstream	trace("/tmp/out.log", std::ios::app);

//=============================================================================
//
class Stream
{
public:
	struct handle_t {
		handle_t(zmq::socket_t* zsock);
		handle_t(int fd);

		zmq::socket_t* zsock;
		int fd;
	};

	typedef uint16_t mode_t;
	static const mode_t eNone = 0;
	static const mode_t eInput = 1;
	static const mode_t eOutput = 2;

public:
	Stream(zmq::socket_t* zsock, const std::string& name, mode_t d, mode_t m);
	Stream(int fd, const std::string& name, mode_t d, mode_t m);
	~Stream();

	const std::string& name() const;

	mode_t mode() const;
	void mode(mode_t m);

	int getHandle() const;
	zmq::socket_t* getSocket() const;

	Bundle* getBundle() const;
	void setParent(Bundle* parent);

private:
	std::string m_name;		// application name for stream
	mode_t m_direction;		// allowed directions
	mode_t m_mode;			// enabled directions
	handle_t m_handle;		// object handle
	Bundle* m_parent;		// owning bundle
};

//=============================================================================
//
Stream::handle_t::handle_t(zmq::socket_t* zsock) :
	zsock(zsock),
	fd(-1)
{
}

Stream::handle_t::handle_t(int fd) :
	zsock(nullptr),
	fd(fd)
{
}

Stream::Stream(zmq::socket_t* zsock, const std::string& name, Stream::mode_t d, Stream::mode_t m) :
	m_name(name),
	m_direction(d),
	m_mode(m),
	m_handle(zsock),
	m_parent(nullptr)
{
}

Stream::Stream(int fd, const std::string& name, Stream::mode_t d, Stream::mode_t m) :
	m_name(name),
	m_direction(d),
	m_mode(m),
	m_handle(fd),
	m_parent(nullptr)
{
}

Stream::~Stream()
{
}

const std::string& Stream::name() const
{
	return m_name;
}

Stream::mode_t Stream::mode() const
{
	return m_mode;
}

void Stream::mode(Stream::mode_t m)
{
	if ((m_direction & Stream::eInput) && (m_mode & Stream::eInput)) {
		m_mode = m;
		return;
	}

	if ((m_direction & eOutput) && (m_mode & eOutput)) {
		m_mode = m;
		return;
	}

	throw std::runtime_error("bad mode in Stream::mode(mode_t)");
}

int Stream::getHandle() const
{
	return m_handle.fd;
}

zmq::socket_t* Stream::getSocket() const
{
	return m_handle.zsock;
}

Bundle* Stream::getBundle() const
{
	return m_parent;
}

void Stream::setParent(Bundle* parent)
{
	m_parent = parent;
}

//=============================================================================
//
class Bundle
{
public:
	typedef Stream*					stream_t;
	typedef std::vector<stream_t>	streams_t;

public:
	Bundle(const std::string& name);
	~Bundle();

	const std::string& name() const;
	void add(Stream* e);
	void remove(Stream* e);
	streams_t getStreams(Stream::mode_t m) const;

private:
	const std::string& m_name;
	streams_t m_streams;	// streams within the bundle
};

//-----------------------------------------------------------------------------
//
Bundle::Bundle(const std::string& name) :
	m_name(name)
{
}

Bundle::~Bundle()
{
	for (auto p = m_streams.rbegin(); p != m_streams.rend(); ++p)
		delete *p;
}

const std::string& Bundle::name() const
{
	return m_name;
}

void Bundle::add(Stream* e)
{
	m_streams.push_back(e);
	e->setParent(this);
}

void Bundle::remove(Stream* e)
{
	for (auto p = m_streams.begin(); p != m_streams.end(); ++p)
		if (*p == e) {
			m_streams.erase(std::remove(m_streams.begin(), m_streams.end(), e), m_streams.end());
			delete e;
			break;
		}
}

Bundle::streams_t Bundle::getStreams(Stream::mode_t m) const
{
	streams_t streams;

	for (const stream_t& s : m_streams)
		if (s) {
			if (static_cast<int>(s->mode()) & static_cast<int>(m))
				streams.push_back(s);
		}

	return streams;
}

//=============================================================================
//
class Pool
{
public:
	typedef std::vector<zmq_pollitem_t>	zmq_pollitems_t;
	typedef ::zmq_pollitem_t			zmq_pollitem_t;
	typedef std::vector<Bundle*>		bundles_t;
	typedef std::vector<Stream*>		streams_t;
	typedef Stream*						stream_t;

public:
	Pool();
	~Pool();

	void add(Bundle* bundle);
	std::tuple<zmq_pollitems_t, streams_t>	getItems(Stream::mode_t mode) const;

private:
	bundles_t m_bundles;	// bundles
};

//-----------------------------------------------------------------------------
//
Pool::Pool()
{
}

Pool::~Pool()
{
	for (auto p = m_bundles.rbegin(); p != m_bundles.rend(); ++p)
		delete *p;
}

void Pool::add(Bundle* bundle)
{
	m_bundles.push_back(bundle);
}

std::tuple<Pool::zmq_pollitems_t, Pool::streams_t> Pool::getItems(Stream::mode_t mode) const
{
	zmq_pollitems_t items;
	streams_t streams;

	for (Bundle* pbundle : m_bundles)
		for (Stream* pstream : pbundle->getStreams(mode)) {
			streams.push_back(pstream);

			short zmq_mode = 0;
			if (pstream->mode() & Stream::eInput)
				zmq_mode |= ZMQ_POLLIN;
			if (pstream->mode() & Stream::eOutput)
				zmq_mode |= ZMQ_POLLOUT;

			if (zmq_mode == 0) {
				trace << "  ignoring streeam: " << pbundle->name() << ":" << pstream->name() << std::endl;
				continue;
			}

			if (zmq::socket_t* zsock = pstream->getSocket()) {
				zmq_pollitem_t rec = { static_cast<void*>(zsock), -1, zmq_mode, 0 };
				items.push_back(rec);
				trace
					<< "  add stream: " << pbundle->name() << ":" << pstream->name() << " for "
					<< ((zmq_mode & ZMQ_POLLIN)               ? "input" :
						(zmq_mode & ZMQ_POLLOUT)              ? "output" :
						(zmq_mode & (ZMQ_POLLIN|ZMQ_POLLOUT)) ? "io" : "unknown")
					<< std::endl;
			}
			else {
				zmq_pollitem_t rec = { nullptr, pstream->getHandle(), zmq_mode, 0 };
				items.push_back(rec);
				trace
					<< "  add stream: " << pbundle->name() << ":" << pstream->name() << " for "
					<< ((zmq_mode & ZMQ_POLLIN)               ? "input" :
						(zmq_mode & ZMQ_POLLOUT)              ? "output" :
						(zmq_mode & (ZMQ_POLLIN|ZMQ_POLLOUT)) ? "io" : "unknown")
					<< std::endl;
			}
		}

	return std::make_tuple(items, streams);
}

namespace
{
	bool reset(const strings_t& msgs, Stream::mode_t& mode)
	{
		// we always want read
		Stream::mode_t working_mode = Stream::eInput;
		if (msgs.empty())
			working_mode &= ~Stream::eOutput;
		else
			working_mode |= Stream::eOutput;

		// calculate mode change
		bool changed = working_mode != mode;
		mode = working_mode;
		trace << " mode: eInput=" << ((mode & Stream::eInput) ? "on" : "off") << " eOutput=" << ((mode & Stream::eOutput) ? "on" : "off") << std::endl;
		return changed;
	}

	void parse_stream_buffer(
		strings_t& msgs,
		bool more,
		std::unique_ptr<char[]>& buf,
		std::unique_ptr<char[]>& spare_buf, char*& spare_buf_ptr)
	{
		char* p = buf.get();
		for (char* q; (q = strchr(p, '\n')); p = q + 1) {
			if (p != q) {
				msgs.emplace_back(p, q);
				trace << "   msg=" << msgs.back() << std::endl;
			}
		}

		if (more) {
			spare_buf.swap(buf);
			spare_buf_ptr = p;
		}
		else {
			if (*p) {
				msgs.emplace_back(p);
				trace << "   msg=" << msgs.back() << std::endl;
			}
		}

		trace << "   msgs.size=" << msgs.size() << std::endl;
	}

	bool read_fd(
			strings_t& msgs,
			Stream::mode_t mode, Pool* pool, Pool::zmq_pollitems_t& items, Pool::streams_t& streams, size_t i)
	{
		const size_t bufsz = 4*1024;		// one page
		std::unique_ptr<char[]> buf;		// main buffer
		std::unique_ptr<char[]> spare_buf;	// overflow buffer
		char* spare_buf_ptr = nullptr;		// overflow buffer offset

		bool more;
		do {
			more = false;

			// allocate the buffer
			if (!buf.get()) {
				buf.reset(new char[bufsz]);
				trace << "  allocate buf: bufsz=" << bufsz << std::endl;
			}
			bzero(buf.get(), bufsz);

			// read the device
			int n = read(items[i].fd, buf.get(), bufsz);
			trace << "  read=" << n << std::endl;

			if (n == 0) {
				// end of stream
				trace << "  closing fd=" << items[i].fd << std::endl;
				close(items[i].fd);
				streams[i]->getBundle()->remove(streams[i]);

				if (more) {
					trace << "  flushing buffer" << std::endl;
					spare_buf.swap(buf);
					parse_stream_buffer(msgs, more, buf, spare_buf, spare_buf_ptr);
				}

				std::tie(items, streams) = pool->getItems(mode | Stream::eInput);	// reinitialize pool items
				trace << " " << items.size() << " items" << std::endl;
				trace << " " << streams.size() << " streams" << std::endl;
				return false;
			}

			trace << "  processing buffer" << std::endl;
			more = n == bufsz;	// do we think there's more to read?
			parse_stream_buffer(msgs, more, buf, spare_buf, spare_buf_ptr);
		}
		while (more);

		return true;
	}
}

int poll(zmq_pollitem_t* data, size_t size, int ms_sleep)
{
	int rc = zmq_poll(data, size, ms_sleep);
	trace << " zmq_poll=" << rc << std::endl;

	return rc;
}

Pool* create_pool(zmq::context_t& ctx)
{
	Pool* pool = new Pool;

	if (Bundle* reader_bundle = new Bundle("reader")) {
		zmq::socket_t* pub = new zmq::socket_t(ctx, zmq::socket_type::pub);
		pub->bind("inproc://news.znc");
		reader_bundle->add(new Stream(pub, "repeater", Stream::eOutput, Stream::eOutput));

		reader_bundle->add(new Stream(STDIN_FILENO, "stdin", Stream::eInput, Stream::eInput));

		pool->add(reader_bundle);
	}

	if (Bundle* writer_bundle = new Bundle("writer")) {
		zmq::socket_t* sub = new zmq::socket_t(ctx, zmq::socket_type::sub);
		sub->connect("inproc://news.znc");
		sub->setsockopt(ZMQ_SUBSCRIBE, "", 0);					// subscribe to everything
		writer_bundle->add(new Stream(sub, "receiver", Stream::eInput, Stream::eInput));

		writer_bundle->add(new Stream(STDOUT_FILENO, "stdout", Stream::eOutput, Stream::eOutput));

		pool->add(writer_bundle);
	}

	return pool;
}

int main()
try
{
	trace << "begin" << std::endl;

	zmq::context_t	ctx(1);
	zmq_ctx_set(static_cast<void*>(ctx), ZMQ_IO_THREADS, 4);

	std::unique_ptr<Pool> pool( create_pool(ctx) );

	strings_t msgs;				// messages read to be written
	strings_t msgs2;
	Stream::mode_t mode = 0;	// stream i/o mode
	Pool::zmq_pollitems_t items;// parallel array: poll items
	Pool::streams_t streams;	// parallel array: streams

	bool stop = false;
	while (!stop) {
		// reset poolitems to wait for
		if (reset(msgs, mode)) {
			std::tie(items, streams) = pool->getItems(mode);	// initialize pool items
			trace << " " << items.size() << " items" << std::endl;
			trace << " " << streams.size() << " streams" << std::endl;
		}

		// wait for poolitems
		int rc = poll(items.data(), items.size(), 1 * 1000);

		// handle what fired
		bool check_items = rc > 0;
		for (size_t i = 0; check_items && i != items.size(); ++i) {
			// handle read
			if (items[i].revents & ZMQ_POLLIN) {
				trace << "  stream(" << i << ") " << streams[i]->name() << " fired for input" << std::endl;

				if (streams[i]->getSocket()) {
					// handle read socket
					zmq::message_t msg;
					streams[i]->getSocket()->recv(&msg);
					msgs2.push_back(std::string(static_cast<const char*>(msg.data()), msg.size()));
				}
				else if (items[i].fd != -1) {
					// handle read fd
					if (!read_fd(msgs, mode, pool.get(), items, streams, i))
						break;	// refresh items fd is closed
				}
			}

			// handle write
			if (items[i].revents & ZMQ_POLLOUT) {
				trace << "  stream(" << i << ") " << streams[i]->name() << " fired for output" << std::endl;

				if (streams[i]->getSocket()) {
					// handle write socket
					zmq::message_t msg(msgs.front().c_str(), msgs.front().size());
					streams[i]->getSocket()->send(msg);
					trace << "  msg=" <<  msgs.front() << std::endl;
					msgs.pop_front();
				}
/* 9 27 19 67 50 26
				else if (items[i].fd != -1) {
					// handle write fd
					if (streams[i]->name() == "stdout") {
						for (int n, offset = 0; offset < (int)msgs2.front().size(); offset += n) {
							n = write(items[i].fd, msgs2.front().c_str() + offset, msgs2.front().size() - offset);
							if (n <= 0)
								break;
						}
						msgs2.pop_front();

						if (msgs2.empty())
							std::tie(items, streams) = pool->getItems(mode | Stream::eInput);	// reinitialize pool items
					}
				}
 */
			}
		}

		stop = rc == 0;
	}

	trace << "done\n" << std::endl;
}
catch (const std::exception& e)
{
	std::clog << "fatal: " << e.what() << "\n" << std::endl;
}