proxysql/include/PgSQL_Monitor.hpp

#ifndef __PGSQL_MONITOR_H
#define __PGSQL_MONITOR_H

#include "libpq-fe.h"

#include "DNS_Cache.hpp"
#include "sqlite3db.h"
#include "proxysql_structs.h"

#include <atomic>
#include <cassert>
#include <memory>
#include <mutex>
#include <vector>

#define MONITOR_SQLITE_TABLE_PGSQL_SERVER_CONNECT_LOG "CREATE TABLE pgsql_server_connect_log (hostname VARCHAR NOT NULL , port INT NOT NULL DEFAULT 3306 , time_start_us INT NOT NULL DEFAULT 0 , connect_success_time_us INT DEFAULT 0 , connect_error VARCHAR , PRIMARY KEY (hostname, port, time_start_us))"

#define MONITOR_SQLITE_TABLE_PGSQL_SERVER_PING_LOG "CREATE TABLE pgsql_server_ping_log ( hostname VARCHAR NOT NULL , port INT NOT NULL DEFAULT 3306 , time_start_us INT NOT NULL DEFAULT 0 , ping_success_time_us INT DEFAULT 0 , ping_error VARCHAR , PRIMARY KEY (hostname, port, time_start_us))"

#define MONITOR_SQLITE_TABLE_PGSQL_SERVER_READ_ONLY_LOG "CREATE TABLE pgsql_server_read_only_log ( hostname VARCHAR NOT NULL , port INT NOT NULL DEFAULT 3306 , time_start_us INT NOT NULL DEFAULT 0 , success_time_us INT DEFAULT 0 , read_only INT DEFAULT 1 , error VARCHAR , PRIMARY KEY (hostname, port, time_start_us))"

#define MONITOR_SQLITE_TABLE_PGSQL_SERVER_REPLICATION_LAG_LOG "CREATE TABLE pgsql_server_replication_lag_log ( hostname VARCHAR NOT NULL , port INT NOT NULL DEFAULT 3306 , time_start_us INT NOT NULL DEFAULT 0 , success_time_us INT DEFAULT 0 , repl_lag INT DEFAULT 0 , error VARCHAR , PRIMARY KEY (hostname, port, time_start_us))"

#define MONITOR_SQLITE_TABLE_PGSQL_SERVERS "CREATE TABLE pgsql_servers (hostname VARCHAR NOT NULL , port INT NOT NULL DEFAULT 3306 , status INT CHECK (status IN (0, 1, 2, 3, 4)) NOT NULL DEFAULT 0 , use_ssl INT CHECK (use_ssl IN(0,1)) NOT NULL DEFAULT 0 , PRIMARY KEY (hostname, port) )"

#define MONITOR_SQLITE_TABLE_PROXYSQL_SERVERS "CREATE TABLE proxysql_servers (hostname VARCHAR NOT NULL , port INT NOT NULL DEFAULT 6032 , weight INT CHECK (weight >= 0) NOT NULL DEFAULT 0 , comment VARCHAR NOT NULL DEFAULT '' , PRIMARY KEY (hostname, port) )"

struct PgSQL_Monitor {
	// @brief Flags if monitoring threads should be shutdown.  Atomic because
	// the resolver / monitor loops read it from worker threads while
	// lifecycle code writes it from the main thread.
	std::atomic<bool> shutdown { false };
	// @brief Mutex to hold to update `monitor_internal.pgsql_servers`
	std::mutex pgsql_srvs_mutex {};
	// @brief Mutex to hold to update/read `pgsql_servers` to monitor
	std::mutex pgsql_srvs_to_monitor_mutex {};
	// @brief Used to access monitor database
	SQLite3DB monitordb {};
	// @brief Used to access internal monitor database
	SQLite3DB monitor_internal_db {};
	// Internal counters for metrics
	///////////////////////////////////////////////////////////////////////////
	uint64_t connect_check_ERR { 0 };
	uint64_t connect_check_OK { 0 };
	uint64_t ping_check_ERR { 0 };
	uint64_t ping_check_OK { 0 };
	uint64_t readonly_check_ERR { 0 };
	uint64_t readonly_check_OK { 0 };
	uint64_t repl_lag_check_ERR { 0 };
	uint64_t repl_lag_check_OK { 0 };
	uint64_t ssl_connections_OK { 0 };
	uint64_t non_ssl_connections_OK { 0 };
	// DNS cache counters (independent of MySQL_Monitor's counters).  Atomic
	// because the resolver workers increment them while readers (stats,
	// p_update_metrics) load them from other threads.
	std::atomic<unsigned long long> dns_cache_queried { 0 };
	std::atomic<unsigned long long> dns_cache_lookup_success { 0 };
	std::atomic<unsigned long long> dns_cache_record_updated { 0 };
	///////////////////////////////////////////////////////////////////////////

	// PgSQL-side DNS cache, independent of MySQL_Monitor's cache.  Lifetime is
	// the same as the PgSQL_Monitor singleton; the resolver loop and the
	// PgSQL_Connection lookup path both go through this shared_ptr.
	std::shared_ptr<DNS_Cache> dns_cache;
	// Set to true by trigger_dns_cache_update() to short-circuit the
	// refresh_interval sleep on the resolver loop.
	std::atomic_bool force_dns_cache_update { false };

	std::vector<table_def_t> tables_defs_monitor {
		{
			const_cast<char*>("pgsql_server_connect_log"),
			const_cast<char*>(MONITOR_SQLITE_TABLE_PGSQL_SERVER_CONNECT_LOG)
		},
		{
			const_cast<char*>("pgsql_server_ping_log"),
			const_cast<char*>(MONITOR_SQLITE_TABLE_PGSQL_SERVER_PING_LOG)
		},
		{
			const_cast<char*>("pgsql_server_read_only_log"),
			const_cast<char*>(MONITOR_SQLITE_TABLE_PGSQL_SERVER_READ_ONLY_LOG)
		},
		{
			const_cast<char*>("pgsql_server_replication_lag_log"),
			const_cast<char*>(MONITOR_SQLITE_TABLE_PGSQL_SERVER_REPLICATION_LAG_LOG)
		},
	};

	std::vector<table_def_t> tables_defs_monitor_internal {
		{
			const_cast<char*>("pgsql_servers"),
			const_cast<char*>(MONITOR_SQLITE_TABLE_PGSQL_SERVERS)
		}
	};

	PgSQL_Monitor();

	// DNS cache facade.  Mirrors MySQL_Monitor's surface so PgSQL_Connection
	// can call PgSQL_Monitor::dns_lookup(hostname) and get a cached IP without
	// blocking on getaddrinfo.
	static std::string dns_lookup(const std::string& hostname,
		bool return_hostname_if_lookup_fails = true, size_t* ip_count = nullptr);
	static std::string dns_lookup(const char* hostname,
		bool return_hostname_if_lookup_fails = true, size_t* ip_count = nullptr);
	// Seed the cache from an already-connected PGconn.  Called on the
	// SUCCESSFUL branch of PgSQL_Connection's handler so the next
	// connect avoids getaddrinfo even before the resolver loop has run.
	static bool update_dns_cache_from_pgsql_conn(PGconn* pgsql_conn);
	// Wake the resolver loop early (e.g. after a pgsql_servers update).
	static void trigger_dns_cache_update();

	// Background loop that drives the resolver pool.  Mirrors
	// MySQL_Monitor::monitor_dns_cache but walks PgHGM->pgsql_servers_to_monitor
	// and reads pgsql-monitor_local_dns_* settings.
	void* monitor_dns_cache();

private:
	static bool _dns_cache_update(const std::string& hostname,
		std::vector<std::string>&& ip_address);
};

struct pgsql_conn_t {
	PGconn* conn { nullptr };
	int fd { 0 };
	uint64_t last_used { 0 };
	ASYNC_ST state { ASYNC_ST::ASYNC_CONNECT_FAILED };
	mf_unique_ptr<char> err {};
};

void* PgSQL_monitor_scheduler_thread();

// pthread entry-point that drives PgSQL_Monitor::monitor_dns_cache() until
// shutdown.  Launched independently of the monitor scheduler so the DNS
// cache stays warm even when pgsql-monitor_enabled=false (same pattern as
// MySQL_Monitor's resolver thread, which is unconditional).
void* PgSQL_monitor_dns_cache_pthread(void* arg);

#endif