pgq: Ability to limit the number of events inserted by one TX

As PgQ can split batches only on transaction boundaries,
it does not tolerate huge number of events inserted by one TX.

This batch adds per-queue field queue_per_tx_limit, which can be
set to reasonable number and PgQ will enforce the limit, by throwing
error if event counts gets larger, thus rollbacking the problematic TX.

Such limit also adds possibility to survive UPDATE/DELETE statements,
where the WHERE clause is missing or buggy.

4 files changed, 115 insertions, 5 deletions

diff --git a/sql/pgq/expected/pgq_core.out b/sql/pgq/expected/pgq_core.out
index e90786ca..ca1fcb71 100644
--- a/sql/pgq/expected/pgq_core.out
+++ b/sql/pgq/expected/pgq_core.out
@@ -344,3 +344,55 @@ select pgq.insert_event('myqueue', 'test', 'event');
          2009
 (1 row)
 
+-- test limit
+update pgq.queue set queue_per_tx_limit = 2 where queue_name = 'myqueue';
+begin;
+select pgq.insert_event('myqueue', 'test', 'event1');
+ insert_event 
+--------------
+         2010
+(1 row)
+
+select pgq.insert_event('myqueue', 'test', 'event2');
+ insert_event 
+--------------
+         2011
+(1 row)
+
+select pgq.insert_event('myqueue', 'test', 'event3');
+ERROR:  Queue 'myqueue' allows max 2 events from one TX
+CONTEXT:  PL/pgSQL function "insert_event" line 19 at RETURN
+PL/pgSQL function "insert_event" line 15 at RETURN
+end;
+update pgq.queue set queue_per_tx_limit = 0 where queue_name = 'myqueue';
+begin;
+select pgq.insert_event('myqueue', 'test', 'event1');
+ERROR:  Queue 'myqueue' allows max 0 events from one TX
+CONTEXT:  PL/pgSQL function "insert_event" line 19 at RETURN
+PL/pgSQL function "insert_event" line 15 at RETURN
+select pgq.insert_event('myqueue', 'test', 'event2');
+ERROR:  current transaction is aborted, commands ignored until end of transaction block
+select pgq.insert_event('myqueue', 'test', 'event3');
+ERROR:  current transaction is aborted, commands ignored until end of transaction block
+end;
+update pgq.queue set queue_per_tx_limit = null where queue_name = 'myqueue';
+begin;
+select pgq.insert_event('myqueue', 'test', 'event1');
+ insert_event 
+--------------
+         2014
+(1 row)
+
+select pgq.insert_event('myqueue', 'test', 'event2');
+ insert_event 
+--------------
+         2015
+(1 row)
+
+select pgq.insert_event('myqueue', 'test', 'event3');
+ insert_event 
+--------------
+         2016
+(1 row)
+
+end;
diff --git a/sql/pgq/lowlevel/insert_event.c b/sql/pgq/lowlevel/insert_event.c
index 87f5ad1b..e55a18a6 100644
--- a/sql/pgq/lowlevel/insert_event.c
+++ b/sql/pgq/lowlevel/insert_event.c
@@ -32,6 +32,11 @@
 PG_MODULE_MAGIC;
 #endif
 
+#ifndef TextDatumGetCString
+#define TextDatumGetCString(d) DatumGetCString(DirectFunctionCall1(textout, d))
+#endif
+
+
 /*
  * Function tag
  */
@@ -47,13 +52,15 @@ PG_FUNCTION_INFO_V1(pgq_insert_event_raw);
 #define QUEUE_SQL \
 	"select queue_id::int4, queue_data_pfx::text," \
 	" queue_cur_table::int4, nextval(queue_event_seq)::int8," \
-	" queue_disable_insert::bool" \
+	" queue_disable_insert::bool," \
+	" queue_per_tx_limit::int4" \
 	" from pgq.queue where queue_name = $1"
 #define COL_QUEUE_ID	1
 #define COL_PREFIX	2
 #define COL_TBLNO	3
 #define COL_EVENT_ID	4
 #define COL_DISABLED	5
+#define COL_LIMIT	6
 
 /*
  * Plan cache entry in HTAB.
@@ -61,6 +68,10 @@ PG_FUNCTION_INFO_V1(pgq_insert_event_raw);
 struct InsertCacheEntry {
 	Oid queue_id;		/* actually int32, but we want to use oid_hash */
 	int cur_table;
+
+	TransactionId last_xid;
+	int last_count;
+
 	void *plan;
 };
 
@@ -73,6 +84,7 @@ struct QueueState {
 	char *table_prefix;
 	Datum next_event_id;
 	bool disabled;
+	int per_tx_limit;
 };
 
 /*
@@ -145,7 +157,7 @@ static void *make_plan(struct QueueState *state)
 /*
  * fetch insert plan from cache.
  */
-static void *load_insert_plan(struct QueueState *state)
+static void *load_insert_plan(Datum qname, struct QueueState *state)
 {
 	struct InsertCacheEntry *entry;
 	Oid queue_id = state->queue_id;
@@ -154,11 +166,27 @@ static void *load_insert_plan(struct QueueState *state)
 	entry = hash_search(insert_cache, &queue_id, HASH_ENTER, &did_exist);
 	if (did_exist) {
 		if (state->cur_table == entry->cur_table)
-			return entry->plan;
+			goto valid_table;
 		SPI_freeplan(entry->plan);
+	} else {
+		entry->last_xid = 0;
 	}
 	entry->cur_table = state->cur_table;
 	entry->plan = make_plan(state);
+valid_table:
+
+	if (state->per_tx_limit >= 0) {
+		TransactionId xid = GetTopTransactionId();
+		if (entry->last_xid != xid) {
+			entry->last_xid = xid;
+			entry->last_count = 0;
+		}
+		entry->last_count++;
+		if (entry->last_count > state->per_tx_limit)
+			elog(ERROR, "Queue '%s' allows max %d events from one TX",
+			     TextDatumGetCString(qname), state->per_tx_limit);
+	}
+
 	return entry->plan;
 }
 
@@ -197,6 +225,9 @@ static void load_queue_info(Datum queue_name, struct QueueState *state)
 	state->disabled = SPI_getbinval(row, desc, COL_DISABLED, &isnull);
 	if (isnull)
 		elog(ERROR, "insert_disabled NULL");
+	state->per_tx_limit = SPI_getbinval(row, desc, COL_LIMIT, &isnull);
+	if (isnull)
+		state->per_tx_limit = -1;
 }
 
 /*
@@ -222,18 +253,20 @@ Datum pgq_insert_event_raw(PG_FUNCTION_ARGS)
 	void *ins_plan;
 	Datum ev_id, ev_time;
 	int i, res;
+	Datum qname;
 
 	if (PG_NARGS() < 6)
 		elog(ERROR, "Need at least 6 arguments");
 	if (PG_ARGISNULL(0))
 		elog(ERROR, "Queue name must not be NULL");
+	qname = PG_GETARG_DATUM(0);
 
 	if (SPI_connect() < 0)
 		elog(ERROR, "SPI_connect() failed");
 
 	init_cache();
 
-	load_queue_info(PG_GETARG_DATUM(0), &state);
+	load_queue_info(qname, &state);
 
 	if (state.disabled)
 		elog(ERROR, "Insert into queue disallowed");
@@ -269,7 +302,7 @@ Datum pgq_insert_event_raw(PG_FUNCTION_ARGS)
 	/*
 	 * Perform INSERT into queue table.
 	 */
-	ins_plan = load_insert_plan(&state);
+	ins_plan = load_insert_plan(qname, &state);
 	res = SPI_execute_plan(ins_plan, values, nulls, false, 0);
 	if (res != SPI_OK_INSERT)
 		elog(ERROR, "Queue insert failed");
diff --git a/sql/pgq/sql/pgq_core.sql b/sql/pgq/sql/pgq_core.sql
index 1d655fb6..c76d7cab 100644
--- a/sql/pgq/sql/pgq_core.sql
+++ b/sql/pgq/sql/pgq_core.sql
@@ -93,3 +93,26 @@ select pgq.insert_event('myqueue', 'test', 'event');
 update pgq.queue set queue_disable_insert = false where queue_name = 'myqueue';
 select pgq.insert_event('myqueue', 'test', 'event');
 
+-- test limit
+update pgq.queue set queue_per_tx_limit = 2 where queue_name = 'myqueue';
+begin;
+select pgq.insert_event('myqueue', 'test', 'event1');
+select pgq.insert_event('myqueue', 'test', 'event2');
+select pgq.insert_event('myqueue', 'test', 'event3');
+end;
+
+update pgq.queue set queue_per_tx_limit = 0 where queue_name = 'myqueue';
+begin;
+select pgq.insert_event('myqueue', 'test', 'event1');
+select pgq.insert_event('myqueue', 'test', 'event2');
+select pgq.insert_event('myqueue', 'test', 'event3');
+end;
+
+update pgq.queue set queue_per_tx_limit = null where queue_name = 'myqueue';
+begin;
+select pgq.insert_event('myqueue', 'test', 'event1');
+select pgq.insert_event('myqueue', 'test', 'event2');
+select pgq.insert_event('myqueue', 'test', 'event3');
+end;
+
+
diff --git a/sql/pgq/structure/tables.sql b/sql/pgq/structure/tables.sql
index 02d2f0a4..6e88c7c5 100644
--- a/sql/pgq/structure/tables.sql
+++ b/sql/pgq/structure/tables.sql
@@ -63,6 +63,7 @@ create table pgq.consumer (
 --      queue_ticker_max_count      - batch should not contain more events
 --      queue_ticker_max_lag        - events should not age more
 --      queue_ticker_idle_period    - how often to tick when no events happen
+--      queue_per_tx_limit          - Max number of events single TX can insert
 --      queue_data_pfx              - prefix for data table names
 --      queue_event_seq             - sequence for event id's
 --      queue_tick_seq              - sequence for tick id's
@@ -84,6 +85,7 @@ create table pgq.queue (
         queue_ticker_max_count      integer     not null default 500,
         queue_ticker_max_lag        interval    not null default '3 seconds',
         queue_ticker_idle_period    interval    not null default '1 minute',
+        queue_per_tx_limit          integer,
 
         queue_data_pfx              text        not null,
         queue_event_seq             text        not null,