Refactored and Optimized Query Cache with PgSQL Support

deps/postgresql/handle_row_data.patch

@@ -6,10 +6,10 @@ index 2265ab5..56883ec 100644
         return conn->result;
  }
 
-+int PShandleRowData(PGconn *conn, PSresult* result) {
++int PShandleRowData(PGconn *conn, bool is_first_packet, PSresult* result) {
 +    if (!conn || !result)
 +		return 1;
-+    return psHandleRowData(conn, result);
++    return psHandleRowData(conn, is_first_packet, result);
 +}
 +
 diff --git src/interfaces/libpq/fe-misc.c src/interfaces/libpq/fe-misc.c
@@ -60,7 +60,7 @@ diff --git src/interfaces/libpq/fe-protocol3.c src/interfaces/libpq/fe-protocol3
 index 9c4aa7e..de0746c 100644
 --- src/interfaces/libpq/fe-protocol3.c
 +++ src/interfaces/libpq/fe-protocol3.c
-@@ -2299,3 +2299,105 @@ build_startup_packet(const PGconn *conn, char *packet,
+@@ -2299,3 +2299,109 @@ build_startup_packet(const PGconn *conn, char *packet,
 
  	return packet_len;
  }
@@ -78,7 +78,7 @@ index 9c4aa7e..de0746c 100644
 + * -1  -> Not enough data to process the message; the next call should be to PQconsumeInput.
 + */
 +int
-+psHandleRowData(PGconn *conn, PSresult* result)
++psHandleRowData(PGconn *conn, bool isFirstPacket, PSresult* result)
 +{
 +	char	id;
 +	int		msgLength;
@@ -122,6 +122,10 @@ index 9c4aa7e..de0746c 100644
 +       return 1;
 +   }
 +
++   /* First data row should be skipped since it is part of PGresult, which contains row description */
++   if (isFirstPacket)
++		return 1;
++
 +	if (conn->result != NULL &&
 +		conn->result->resultStatus == PGRES_TUPLES_OK)
 +	{
@@ -194,7 +198,7 @@ index c5170d1..3e3cc34 100644
  extern const PGresult *PQgetResultFromPGconn(PGconn *conn);
 
 +/* ProxySQL special handler function */
-+extern int PShandleRowData(PGconn *conn, PSresult* result);
++extern int PShandleRowData(PGconn *conn, bool is_first_packet, PSresult* result);
 +
  #ifdef __cplusplus
  }
@@ -210,7 +214,7 @@ index a951f49..e1df8b5 100644
 + /*
 +  * ProxySQL light weight routines
 +  */
-+extern int  psHandleRowData(PGconn *conn, PSresult* result);								 
++extern int  psHandleRowData(PGconn *conn, bool is_first_packet, PSresult* result);								 
 +
  /* === in fe-misc.c === */
 

include/MySQL_Query_Cache.h

@@ -0,0 +1,26 @@
+#ifndef __CLASS_MYSQL_QUERY_CACHE_H
+#define __CLASS_MYSQL_QUERY_CACHE_H
+
+#include "proxysql.h"
+#include "cpp.h"
+#include "query_cache.hpp"
+
+typedef struct _MySQL_QC_entry : public QC_entry_t {
+	uint32_t column_eof_pkt_offset;
+	uint32_t row_eof_pkt_offset;
+	uint32_t ok_pkt_offset;
+} MySQL_QC_entry_t;
+
+class MySQL_Query_Cache : public Query_Cache<MySQL_Query_Cache> {
+public:
+	MySQL_Query_Cache() = default;
+	~MySQL_Query_Cache() = default;
+
+	bool set(uint64_t user_hash, const unsigned char* kp, uint32_t kl, unsigned char* vp, uint32_t vl, 
+		uint64_t create_ms, uint64_t curtime_ms, uint64_t expire_ms, bool deprecate_eof_active);
+	unsigned char* get(uint64_t user_hash, const unsigned char* kp, const uint32_t kl, uint32_t* lv, 
+		uint64_t curtime_ms, uint64_t cache_ttl, bool deprecate_eof_active);
+	//void* purgeHash_thread(void*);
+};
+
+#endif /* __CLASS_MYSQL_QUERY_CACHE_H */

include/PgSQL_Data_Stream.h

@@ -108,8 +108,8 @@ class PgSQL_Data_Stream
 	FixedSizeQueue data_packets_history_IN;
 	FixedSizeQueue data_packets_history_OUT;
 	//PtrSizeArray *PSarrayOUTpending;
-	PtrSizeArray* resultset;
-	unsigned int resultset_length;
+	//PtrSizeArray* resultset;
+	//unsigned int resultset_length;
 
 	ProxySQL_Poll<PgSQL_Data_Stream>* mypolls;
 	//int listener;
@@ -201,8 +201,9 @@ class PgSQL_Data_Stream
 	void check_data_flow();
 	int assign_fd_from_mysql_conn();
 
-	unsigned char* resultset2buffer(bool);
-	void buffer2resultset(unsigned char*, unsigned int);
+	static unsigned char* copy_array_to_buffer(PtrSizeArray* resultset, size_t resultset_length, bool del);
+	static void copy_buffer_to_resultset(PtrSizeArray* resultset, unsigned char* ptr, uint64_t size, 
+		char current_transaction_state);
 
 	// safe way to attach a PgSQL Connection
 	void attach_connection(PgSQL_Connection* mc) {

include/PgSQL_Query_Cache.h

@@ -0,0 +1,22 @@
+#ifndef __CLASS_PGSQL_QUERY_CACHE_H
+#define __CLASS_PGSQL_QUERY_CACHE_H
+
+#include "proxysql.h"
+#include "cpp.h"
+#include "query_cache.hpp"
+
+typedef struct _PgSQL_QC_entry : public QC_entry_t {} PgSQL_QC_entry_t;
+
+class PgSQL_Query_Cache : public Query_Cache<PgSQL_Query_Cache> {
+public:
+	PgSQL_Query_Cache() = default;
+	~PgSQL_Query_Cache() = default;
+
+	bool set(uint64_t user_hash, const unsigned char* kp, uint32_t kl, unsigned char* vp, uint32_t vl, 
+		uint64_t create_ms, uint64_t curtime_ms, uint64_t expire_ms);
+	const std::shared_ptr<PgSQL_QC_entry_t> get(uint64_t user_hash, const unsigned char* kp, const uint32_t kl, 
+		uint64_t curtime_ms, uint64_t cache_ttl);
+	//void* purgeHash_thread(void*);
+};
+
+#endif /* __CLASS_PGSQL_QUERY_CACHE_H */

include/cpp.h

@@ -10,7 +10,7 @@
 #include "PgSQL_Backend.h"
 #include "ProxySQL_Poll.h"
 //#include "MySQL_Data_Stream.h"
-#include "query_cache.hpp"
+//#include "MySQL_Query_Cache.h"
 #include "mysql_connection.h"
 #include "sqlite3db.h"
 //#include "StatCounters.h"

include/gen_utils.h

@@ -325,6 +325,7 @@ char *trim_spaces_in_place(char *str);
 char *trim_spaces_and_quotes_in_place(char *str);
 bool mywildcmp(const char *p, const char *str);
 std::string trim(const std::string& s);
+char* escape_string_single_quotes_and_backslashes(char* input, bool free_it);
 
 /**
  * @brief Helper function that converts a MYSQL_RES into a 'SQLite3_result'.

include/proxysql_structs.h

@@ -706,7 +706,8 @@ class SimpleKV;
 class AdvancedKV;
 template <class T>
 class ProxySQL_Poll;
-class Query_Cache;
+class MySQL_Query_Cache;
+class PgSQL_Query_Cache;
 class MySQL_Authentication;
 class MySQL_Connection;
 class PgSQL_Connection;
@@ -1095,6 +1096,10 @@ __thread int pgsql_thread___monitor_threads;
 __thread char* pgsql_thread___monitor_username;
 __thread char* pgsql_thread___monitor_password;
 
+// PgSQL Query Cache
+__thread int pgsql_thread___query_cache_size_MB;
+__thread int pgsql_thread___query_cache_soft_ttl_pct;
+__thread int pgsql_thread___query_cache_handle_warnings;
 //---------------------------
 
 __thread char *mysql_thread___default_schema;
@@ -1382,6 +1387,10 @@ extern __thread int pgsql_thread___monitor_threads;
 extern __thread char* pgsql_thread___monitor_username;
 extern __thread char* pgsql_thread___monitor_password;
 
+// PgSQL Query Cache
+extern __thread int pgsql_thread___query_cache_size_MB;
+extern __thread int pgsql_thread___query_cache_soft_ttl_pct;
+extern __thread int pgsql_thread___query_cache_handle_warnings;
 //---------------------------
 
 extern __thread char *mysql_thread___default_schema;

include/query_cache.hpp

@@ -1,9 +1,9 @@
 #ifndef __CLASS_QUERY_CACHE_H
 #define __CLASS_QUERY_CACHE_H
-
 #include "proxysql.h"
 #include "cpp.h"
-#include <tuple>
+#include "prometheus/counter.h"
+#include "prometheus/gauge.h"
 
 #define EXPIRE_DROPIT   0
 #define SHARED_QUERY_CACHE_HASH_TABLES  32
@@ -13,30 +13,6 @@
 #define DEFAULT_purge_threshold_pct_min 3
 #define DEFAULT_purge_threshold_pct_max 90
 
-#include "prometheus/counter.h"
-#include "prometheus/gauge.h"
-
-class KV_BtreeArray;
-
-typedef struct __QC_entry_t QC_entry_t;
-
-struct __QC_entry_t {
-	uint64_t key; // primary key
-	char *value;  // pointer to value
-	KV_BtreeArray *kv; // pointer to the KV_BtreeArray where the entry is stored
-	QC_entry_t *self; // pointer to itself
-	uint32_t klen; // length of the key : FIXME: not sure if still relevant
-	uint32_t length; // length of the value
-	unsigned long long create_ms; // when the entry was created, monotonic, millisecond granularity
-	unsigned long long expire_ms; // when the entry will expire, monotonic , millisecond granularity
-	unsigned long long access_ms; // when the entry was read last , monotonic , millisecond granularity
-	bool refreshing; // true when a client will hit the backend to refresh the entry
-	uint32_t column_eof_pkt_offset = 0;
-	uint32_t row_eof_pkt_offset = 0;
-	uint32_t ok_pkt_offset = 0;
-	uint32_t ref_count; // reference counter
-};
-
 struct p_qc_counter {
 	enum metric {
 		query_cache_count_get = 0,
@@ -65,34 +41,67 @@ struct qc_metrics_map_idx {
 };
 
 class KV_BtreeArray;
+class MySQL_Query_Cache;
+class PgSQL_Query_Cache;
+struct _MySQL_QC_entry;
+struct _PgSQL_QC_entry;
+typedef struct _MySQL_QC_entry MySQL_QC_entry_t;
+typedef struct _PgSQL_QC_entry PgSQL_QC_entry_t;
+
+typedef struct _QC_entry {
+	uint64_t key;			// primary key
+	unsigned char *value;	// pointer to value
+	uint32_t length;		// length of the value
+	uint32_t klen;			// length of the key : FIXME: not sure if still relevant
+	uint64_t create_ms;		// when the entry was created, monotonic, millisecond granularity
+	uint64_t expire_ms;		// when the entry will expire, monotonic , millisecond granularity
+	uint64_t access_ms;		// when the entry was read last , monotonic , millisecond granularity
+	bool refreshing;		// true when a client will hit the backend to refresh the entry
+	KV_BtreeArray* kv;		// pointer to the KV_BtreeArray where the entry is stored (used for troubleshooting)
+	//struct _QC_entry* self; // pointer to itself
+} QC_entry_t;
+
+template <typename QC_DERIVED>
 class Query_Cache {
-	private:
-	KV_BtreeArray * KVs[SHARED_QUERY_CACHE_HASH_TABLES];
-	uint64_t get_data_size_total();
-	unsigned int current_used_memory_pct();
-	struct {
-		std::array<prometheus::Counter*, p_qc_counter::__size> p_counter_array {};
-		std::array<prometheus::Gauge*, p_qc_gauge::__size> p_gauge_array {};
-	} metrics;
-	public:
+	static_assert(std::is_same_v<QC_DERIVED,MySQL_Query_Cache> || std::is_same_v<QC_DERIVED,PgSQL_Query_Cache>,
+		"Invalid QC_DERIVED Query Cache type");
+	using TypeQCEntry = typename std::conditional<std::is_same_v<QC_DERIVED, MySQL_Query_Cache>, 
+		MySQL_QC_entry_t, PgSQL_QC_entry_t>::type;
+public:
+	static bool shutting_down;
+	static pthread_t purge_thread_id;
+	constexpr static unsigned int purge_loop_time = DEFAULT_purge_loop_time;
+
+	void print_version();
+	uint64_t flush();
 	void p_update_metrics();
-	void * purgeHash_thread(void *);
-	int size;
-	int shutdown;
-	unsigned long long QCnow_ms;
-	pthread_t purge_thread_id;
-	unsigned int purge_loop_time;
-	unsigned int purge_total_time;
-	unsigned int purge_threshold_pct_min;
-	unsigned int purge_threshold_pct_max;
-	uint64_t max_memory_size;
+	SQLite3_result* SQL3_getStats();
+	void purgeHash(uint64_t max_memory_size);
+
+protected:
 	Query_Cache();
 	~Query_Cache();
-	void print_version();
-	bool set(uint64_t user_hash, const unsigned char *kp, uint32_t kl, unsigned char *vp, uint32_t vl, unsigned long long create_ms, unsigned long long curtime_ms, unsigned long long expire_ms, bool deprecate_eof_active);
-	unsigned char * get(uint64_t , const unsigned char *, const uint32_t, uint32_t *, unsigned long long, unsigned long long, bool deprecate_eof_active);
-	uint64_t flush();
-	SQLite3_result * SQL3_getStats();
+
+	bool set(QC_entry_t* entry, uint64_t user_hash, const unsigned char *kp, uint32_t kl, unsigned char *vp,
+		uint32_t vl, uint64_t create_ms, uint64_t curtime_ms, uint64_t expire_ms);
+	std::shared_ptr<QC_entry_t> get(uint64_t user_hash, const unsigned char* kp, const uint32_t kl, 
+		uint64_t curtime_ms, uint64_t cache_ttl);
+
+	constexpr static unsigned int purge_total_time = DEFAULT_purge_total_time;
+	constexpr static unsigned int purge_threshold_pct_min = DEFAULT_purge_threshold_pct_min;
+	constexpr static unsigned int purge_threshold_pct_max = DEFAULT_purge_threshold_pct_max;
+	//uint64_t max_memory_size;
+
+private:
+	KV_BtreeArray* KVs[SHARED_QUERY_CACHE_HASH_TABLES];
+	uint64_t get_data_size_total();
+	unsigned int current_used_memory_pct(uint64_t max_memory_size);
+	void purgeHash(uint64_t QCnow_ms, unsigned int curr_pct);
+
+	struct {
+		std::array<prometheus::Counter*, p_qc_counter::__size> p_counter_array{};
+		std::array<prometheus::Gauge*, p_qc_gauge::__size> p_gauge_array{};
+	} metrics;
 };
-#endif /* __CLASS_QUERY_CACHE_H */
 
+#endif /* __CLASS_QUERY_CACHE_H */

include/query_processor.h

@@ -177,7 +177,7 @@ class Query_Processor_Output {
 		mirror_flagOUT=-1;
 		next_query_flagIN=-1;
 		cache_ttl=-1;
-		cache_empty_result=1;
+		cache_empty_result=-1;
 		cache_timeout=-1;
 		reconnect=-1;
 		timeout=-1;

lib/Admin_Bootstrap.cpp

@@ -148,7 +148,7 @@ struct cpu_timer
 extern int admin_load_main_;
 extern bool admin_nostart_;
 
-extern Query_Cache *GloQC;
+//extern MySQL_Query_Cache *GloMyQC;
 extern MySQL_Authentication *GloMyAuth;
 extern PgSQL_Authentication *GloPgAuth;
 extern MySQL_LDAP_Authentication *GloMyLdapAuth;

lib/Admin_FlushVariables.cpp

@@ -125,7 +125,7 @@ extern char * proxysql_version;
 
 #include "proxysql_find_charset.h"
 
-extern Query_Cache *GloQC;
+//extern MySQL_Query_Cache *GloMyQC;
 extern MySQL_Authentication *GloMyAuth;
 extern PgSQL_Authentication *GloPgAuth;
 extern MySQL_LDAP_Authentication *GloMyLdapAuth;