pgLab/pglablib/ASyncDBConnection.cpp

#include "ASyncDBConnection.h"
#include "ScopeGuard.h"
#include "util.h"
#include "Pgsql_PgException.h"
#include <atomic>
#include <chrono>
#include <queue>
#include <QDebug>
#include <QEventLoop>
#include <QSocketNotifier>
#include <QTimer>

namespace {

    class registerMetaTypes {
    public:
        registerMetaTypes()
        {
            qRegisterMetaType<ASyncDBConnection::State>();
            qRegisterMetaType<Pgsql::ErrorDetails>();
			qRegisterMetaType<std::shared_ptr<Pgsql::Result>>();
		}
    } registerMetaTypes_instance;
    
	class Command {
	public:
		std::string command;
		Pgsql::Params params;
		ASyncDBConnection::on_result_callback on_result;

		Command() = default;
		Command(const std::string &cmd, ASyncDBConnection::on_result_callback cb)
			: command(cmd), on_result(cb)
		{}
		Command(const std::string &cmd, Pgsql::Params &&p, ASyncDBConnection::on_result_callback cb)
			: command(cmd), params(p), on_result(cb)
		{}
	};

}

class ASyncDBConnectionThread {
public:
	using t_CommandQueue = std::queue<Command>;

	struct t_Command {
		std::mutex m_mutex;
		t_CommandQueue m_queue;
	} m_commandQueue;

	ConnectionConfig m_config;
	ASyncDBConnection::State m_state = ASyncDBConnection::State::NotConnected;

	ASyncDBConnectionThread(ASyncDBConnection *asco);

	/// Is started as a seperate thread by ASyncDBConnection
	void run();

	/// Sends a cancel request to the DB server
	bool cancel();

	void stop();

	/// Wakes the worker thread when it is waiting in its event loop
	void wake();

private:
	enum class WaitResult {
		Socket,  ///< the socket became ready
		Wakeup,  ///< woken up by stop() or by a new command
		Timeout  ///< the timeout expired
	};

	ASyncDBConnection *asyncConnObject = nullptr;

	std::atomic<bool> terminateRequested{false}; ///< is set when the thread should stop
	std::mutex m_wakeMutex;
	QEventLoop *m_loop = nullptr; ///< loop the worker blocks in, guarded by m_wakeMutex
	Pgsql::Connection m_connection;
	bool m_terminated = true;
	Pgsql::Canceller m_canceller;
	QElapsedTimer m_timer;


	bool hasQueuedCommand();
	WaitResult waitForSocket(QSocketNotifier::Type type, int timeout_ms, bool wakeOnNewCommand);
	bool makeConnection();
	void communicate();

    void doStateCallback(ASyncDBConnection::StateData state);
	/// Wait's for a command to come in and send's it to the server
    bool waitForAndSendCommand();
    void doNewCommand();
	void waitForResult();


	void processNotice(const PGresult *result);
	/** Function to call when after sending a command the socket is ready for reading.
	 *
	 * It might take several consumes before all data is read.
	 */
	bool consumeResultInput();

};

ASyncDBConnectionThread::ASyncDBConnectionThread(ASyncDBConnection *asco)
	: asyncConnObject(asco)
{}

void ASyncDBConnectionThread::run()
{
	m_terminated = false;
	SCOPE_EXIT {
		m_state = ASyncDBConnection::State::NotConnected;
		m_terminated = true;
        m_connection.close();
//        doStateCallback(ASyncDBConnection::State::NotConnected);
	};

    try {
        while (!terminateRequested) {

            // make or recover connection
            if (makeConnection()) {
                m_connection.setNoticeReceiver(
                    [this](const PGresult *result) { processNotice(result); });
                m_canceller = m_connection.getCancel();


                // send commands and receive results
                communicate();
            }
            else {
                // It is not possible to determine the source of the problem.
                // Accept for PQconnectionNeedsPassword

                // Pass problem to main thread and stop this thread

                // Main thread needs to know it has to restart connecting if it want's to.
                // TODO: add status functions to help main thread so it doesn't have to remember
                // everything reported through callbacks.

                break;
            }
        }
        doStateCallback({ ASyncDBConnection::State::NotConnected, QString("terminating") });
    }
    catch (const std::exception &ex)
    {
        doStateCallback({ ASyncDBConnection::State::NotConnected, QString::fromUtf8(ex.what()) });
    }
    catch (...)
    {
        doStateCallback({ ASyncDBConnection::State::NotConnected, QString::fromUtf8("Unknown error") });
    }

}

bool ASyncDBConnectionThread::cancel()
{
	return m_canceller.cancel(nullptr);
}

bool ASyncDBConnectionThread::hasQueuedCommand()
{
	std::lock_guard<std::mutex> lg(m_commandQueue.m_mutex);
	return !m_commandQueue.m_queue.empty();
}

ASyncDBConnectionThread::WaitResult ASyncDBConnectionThread::waitForSocket(QSocketNotifier::Type type, int timeout_ms, bool wakeOnNewCommand)
{
	QEventLoop loop;
	WaitResult result = WaitResult::Wakeup;

	QSocketNotifier notifier(m_connection.socket(), type);
	QObject::connect(&notifier, &QSocketNotifier::activated, &loop, [&result, &loop] {
		result = WaitResult::Socket;
		loop.quit();
	});

	QTimer timer;
	timer.setSingleShot(true);
	if (timeout_ms >= 0) {
		QObject::connect(&timer, &QTimer::timeout, &loop, [&result, &loop] {
			result = WaitResult::Timeout;
			loop.quit();
		});
		timer.start(timeout_ms);
	}

	{
		std::lock_guard<std::mutex> lg(m_wakeMutex);
		m_loop = &loop;
	}

	// Checked after publishing the loop, so a wake-up that happened before is not lost
	bool pending = terminateRequested || (wakeOnNewCommand && hasQueuedCommand());
	if (!pending) {
		loop.exec();
	}

	{
		std::lock_guard<std::mutex> lg(m_wakeMutex);
		m_loop = nullptr;
	}
	return result;
}

void ASyncDBConnectionThread::wake()
{
	std::lock_guard<std::mutex> lg(m_wakeMutex);
	if (m_loop) {
		QMetaObject::invokeMethod(m_loop, [loop = m_loop] { loop->quit(); }, Qt::QueuedConnection);
	}
}

bool ASyncDBConnectionThread::makeConnection()
{
	using namespace std::literals::chrono_literals;

	// start connecting
	QString conn_string = m_config.connectionString();
	while (!terminateRequested) {

		bool ok = m_connection.connectStart(conn_string); //keywords, values);
		auto start = std::chrono::steady_clock::now();
		if (ok && m_connection.status() != CONNECTION_BAD) {
			QSocketNotifier::Type type = QSocketNotifier::Write;
			while (true) {
				// poll till complete or failed (we can get an abort command)
				auto elapsed = std::chrono::steady_clock::now() - start;
				auto remaining = std::chrono::duration_cast<std::chrono::milliseconds>(30s - elapsed).count();
				WaitResult res = waitForSocket(type, remaining > 0 ? static_cast<int>(remaining) : 0, false);

				if (res == WaitResult::Socket) {
					auto poll_state = m_connection.connectPoll();
					if (poll_state == PGRES_POLLING_OK) {
                        qDebug() << "ASyncDBConnection Connected";
						// if connected return true
                        doStateCallback({ ASyncDBConnection::State::Connected, "Success" });
						return true;
					}
					else if (poll_state == PGRES_POLLING_FAILED) {
                        doStateCallback({ ASyncDBConnection::State::NotConnected, "Failed to connect" });
						return false;
					}
					else if (poll_state == PGRES_POLLING_READING) {
                        doStateCallback({ ASyncDBConnection::State::Connecting, "Negotiating" });
						type = QSocketNotifier::Read;
					}
					else if (poll_state == PGRES_POLLING_WRITING) {
                        doStateCallback({ ASyncDBConnection::State::Connecting, "Negotiating" });
						type = QSocketNotifier::Write;
					}
				}
				else if (terminateRequested) {
					m_connection.close();
					return false;
				}
				else if (res == WaitResult::Timeout) {
					doStateCallback({ ASyncDBConnection::State::NotConnected, "Connection timed out" });
					m_connection.close();
					return false;
				}
			} // end while (true)
		}
	}
	return false;
}

void ASyncDBConnectionThread::communicate()
{
	while (!terminateRequested) {
		// wait for something to do:
		// - command to send to server
		//   - wait for results and (notifies can also come in)
		//   - pass each result on to the completion routine
		// - notify comming in from the server
		//   - pass to notify callback
		// - connection raises an error
		//   - return
		// - stop signal
		//   - return


		if (m_state == ASyncDBConnection::State::Connected) {
            if (!waitForAndSendCommand()) {
                return;
            }
		}
		else if (m_state == ASyncDBConnection::State::QuerySend || m_state == ASyncDBConnection::State::CancelSend) {
			// Wait for result, even after a cancel we should wait, for all results
			// New command's are not excepted when one has been send
			waitForResult();
		}

	}
}

void ASyncDBConnectionThread::stop()
{
	terminateRequested = true;
	wake();
}

void ASyncDBConnectionThread::doStateCallback(ASyncDBConnection::StateData state)
{
    qDebug() << "State change " + state.Message;
    m_state = state.State;
    Q_EMIT asyncConnObject->onStateChanged();
}

bool ASyncDBConnectionThread::waitForAndSendCommand()
{
	if (hasQueuedCommand()) {
		doNewCommand();
		return true;
	}

	// Idle connection: a readable socket carries notifies or notices, or signals a lost connection
	if (waitForSocket(QSocketNotifier::Read, -1, true) == WaitResult::Socket) {
		if (!m_connection.consumeInput()) {
			doStateCallback({ ASyncDBConnection::State::NotConnected, "Connection lost" });
			return false;
		}
	}
    // Note if it was stop we can just return and function
    // above will stop looping because terminateRequested has been set too by stop
    return true;
}

void ASyncDBConnectionThread::doNewCommand()
{
	// get command from top of queue (but leave it in the queue, we need the callback)
	if (! m_commandQueue.m_queue.empty())	{
		const Command &command = m_commandQueue.m_queue.front();
		if (command.params.empty())
            m_connection.sendQuery(command.command.c_str());
		else
            m_connection.sendQueryParams(command.command.c_str(), command.params);

        m_timer.start();
        doStateCallback(ASyncDBConnection::State::QuerySend);
	}
}

bool ASyncDBConnectionThread::consumeResultInput()
{
	bool finished = false;
	if (m_connection.consumeInput()) {
		while ( ! finished && ! m_connection.isBusy()) {
			auto res(m_connection.getResultNoThrow());
			{
				qint64 ms = m_timer.restart();
				std::lock_guard<std::mutex> lg(m_commandQueue.m_mutex);
				m_commandQueue.m_queue.front().on_result(res, ms);
				if (res == nullptr) {
					m_timer.invalidate();
					m_commandQueue.m_queue.pop();
					doStateCallback(ASyncDBConnection::State::Connected);
					finished = true;
				}
			}

		}
		// else is still waiting for more data

	}
	else {
		// error during consume
		auto error_msg = m_connection.getErrorMessage();
		qDebug() << "error while communicating with server " << QString::fromStdString(error_msg);
        doStateCallback({ASyncDBConnection::State::NotConnected, QString::fromStdString(error_msg)});
		finished = true;
		stop();
	}
	return finished;
}

void ASyncDBConnectionThread::waitForResult()
{
	bool finished = false;
	while ( ! finished) {
		WaitResult res = waitForSocket(QSocketNotifier::Read, -1, false);

		if (res == WaitResult::Socket) {
			if (consumeResultInput()) {
				finished = true;
			}
		}
		else if (terminateRequested) {
			// Send cancel, close connection and terminate thread
			cancel();
			doStateCallback(ASyncDBConnection::State::Terminating);
			finished = true;
		}
	} // end while
}

void ASyncDBConnectionThread::processNotice(const PGresult *result)
{
	Pgsql::ErrorDetails details = Pgsql::ErrorDetails::createErrorDetailsFromPGresult(result);
    Q_EMIT asyncConnObject->onNotice(details);
}




ASyncDBConnection::ASyncDBConnection()
	: m_threadData(std::make_unique<ASyncDBConnectionThread>(this))
{}

ASyncDBConnection::~ASyncDBConnection()
{
	closeConnection();
}

ASyncDBConnection::State ASyncDBConnection::state() const
{
	return m_threadData->m_state;
}

void ASyncDBConnection::setupConnection(const ConnectionConfig &config)
{
	if (m_thread.joinable()) {
		m_threadData->stop();
		m_thread.join();
	}

	m_threadData.reset(new ASyncDBConnectionThread(this));
	m_threadData->m_config = config;
	m_thread = std::thread([this] () { m_threadData->run(); });

}

void ASyncDBConnection::doStateCallback(State state)
{
    m_state = state;
    if (state == State::Connected) {
        m_canceller = m_connection.getCancel();
        m_connection.setNoticeReceiver(
            [this](const PGresult *result) { processNotice(result); });
    }
    Q_EMIT onStateChanged();
}


void ASyncDBConnection::closeConnection()
{
	m_threadData->stop();
	if (m_thread.joinable()) {
		m_thread.join();
	}
}

bool ASyncDBConnection::send(const std::string &command, on_result_callback on_result)
{
    if (command.empty() || state() == State::NotConnected)
		return false;

	{
		std::lock_guard<std::mutex> lg(m_threadData->m_commandQueue.m_mutex);
		m_threadData->m_commandQueue.m_queue.emplace(command, on_result);
	}
	m_threadData->wake();
	return true;
}

bool ASyncDBConnection::send(const std::string &command, Pgsql::Params params, on_result_callback on_result)
{
    if (command.empty() || state() == State::NotConnected)
        return false;

    {
		std::lock_guard<std::mutex> lg(m_threadData->m_commandQueue.m_mutex);
		m_threadData->m_commandQueue.m_queue.emplace(command, std::move(params), on_result);
	}
	m_threadData->wake();
	return true;
}

bool ASyncDBConnection::cancel()
{
	return m_threadData->cancel();
}

void ASyncDBConnection::processNotice(const PGresult *result)
{
    Pgsql::ErrorDetails details = Pgsql::ErrorDetails::createErrorDetailsFromPGresult(result);
    Q_EMIT onNotice(details);
}