Something went wrong. Try again.
Highly ambitious ATProtocol AppView service and sdks forked from slices.network/slices
Something went wrong. Try again.
9.2 kB · 289 lines
SQL
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
-- The UDT for creating messagesCREATE TYPE mq_new_t AS ( -- Unique message ID id UUID, -- Delay before message is processed delay INTERVAL, -- Number of retries if initial processing fails retries INT, -- Initial backoff between retries retry_backoff INTERVAL, -- Name of channel channel_name TEXT, -- Arguments to channel channel_args TEXT, -- Interval for two-phase commit (or NULL to disable two-phase commit) commit_interval INTERVAL, -- Whether this message should be processed in order with respect to other -- ordered messages. ordered BOOLEAN, -- Name of message name TEXT, -- JSON payload payload_json TEXT, -- Binary payload payload_bytes BYTEA);
-- Small, frequently updated table of messagesCREATE TABLE mq_msgs ( id UUID PRIMARY KEY, created_at TIMESTAMPTZ DEFAULT NOW(), attempt_at TIMESTAMPTZ DEFAULT NOW(), attempts INT NOT NULL DEFAULT 5, retry_backoff INTERVAL NOT NULL DEFAULT INTERVAL '1 second', channel_name TEXT NOT NULL, channel_args TEXT NOT NULL, commit_interval INTERVAL, after_message_id UUID DEFAULT '00000000-0000-0000-0000-000000000000'::uuid REFERENCES mq_msgs(id) ON DELETE SET DEFAULT);
-- Insert dummy message so that the 'nil' UUID can be referencedINSERT INTO mq_msgs (id, channel_name, channel_args, after_message_id) VALUES ('00000000-0000-0000-0000-000000000000'::uuid, '', '', NULL);
-- Internal helper function to check that a UUID is neither NULL nor NILCREATE FUNCTION mq_uuid_exists( id UUID) RETURNS BOOLEAN AS $$ SELECT id IS NOT NULL AND id != '00000000-0000-0000-0000-000000000000'::uuid$$ LANGUAGE SQL IMMUTABLE;
-- Index for pollingCREATE INDEX ON mq_msgs(channel_name, channel_args, attempt_at) WHERE id != '00000000-0000-0000-0000-000000000000'::uuid AND NOT mq_uuid_exists(after_message_id);-- Index for adding messagesCREATE INDEX ON mq_msgs(channel_name, channel_args, created_at, id) WHERE id != '00000000-0000-0000-0000-000000000000'::uuid AND after_message_id IS NOT NULL;
-- Index for ensuring strict message orderCREATE UNIQUE INDEX mq_msgs_channel_name_channel_args_after_message_id_idx ON mq_msgs(channel_name, channel_args, after_message_id);
-- Large, less frequently updated table of message payloadsCREATE TABLE mq_payloads( id UUID PRIMARY KEY, name TEXT NOT NULL, payload_json JSONB, payload_bytes BYTEA);
-- Internal helper function to return the most recently added message in a queue.CREATE FUNCTION mq_latest_message(from_channel_name TEXT, from_channel_args TEXT)RETURNS UUID AS $$ SELECT COALESCE( ( SELECT id FROM mq_msgs WHERE channel_name = from_channel_name AND channel_args = from_channel_args AND after_message_id IS NOT NULL AND id != '00000000-0000-0000-0000-000000000000'::uuid ORDER BY created_at DESC, id DESC LIMIT 1 ), '00000000-0000-0000-0000-000000000000'::uuid )$$ LANGUAGE SQL STABLE;
-- Internal helper function to randomly select a set of channels with "ready" messages.CREATE FUNCTION mq_active_channels(channel_names TEXT[], batch_size INT)RETURNS TABLE(name TEXT, args TEXT) AS $$ SELECT channel_name, channel_args FROM mq_msgs WHERE id != '00000000-0000-0000-0000-000000000000'::uuid AND attempt_at <= NOW() AND (channel_names IS NULL OR channel_name = ANY(channel_names)) AND NOT mq_uuid_exists(after_message_id) GROUP BY channel_name, channel_args ORDER BY RANDOM() LIMIT batch_size$$ LANGUAGE SQL STABLE;
-- Main entry-point for job runner: pulls a batch of messages from the queue.CREATE FUNCTION mq_poll(channel_names TEXT[], batch_size INT DEFAULT 1)RETURNS TABLE( id UUID, is_committed BOOLEAN, name TEXT, payload_json TEXT, payload_bytes BYTEA, retry_backoff INTERVAL, wait_time INTERVAL) AS $$BEGIN RETURN QUERY UPDATE mq_msgs SET attempt_at = CASE WHEN mq_msgs.attempts = 1 THEN NULL ELSE NOW() + mq_msgs.retry_backoff END, attempts = mq_msgs.attempts - 1, retry_backoff = mq_msgs.retry_backoff * 2 FROM ( SELECT msgs.id FROM mq_active_channels(channel_names, batch_size) AS active_channels INNER JOIN LATERAL ( SELECT * FROM mq_msgs WHERE mq_msgs.id != '00000000-0000-0000-0000-000000000000'::uuid AND mq_msgs.attempt_at <= NOW() AND mq_msgs.channel_name = active_channels.name AND mq_msgs.channel_args = active_channels.args AND NOT mq_uuid_exists(mq_msgs.after_message_id) ORDER BY mq_msgs.attempt_at ASC LIMIT batch_size ) AS msgs ON TRUE LIMIT batch_size ) AS messages_to_update LEFT JOIN mq_payloads ON mq_payloads.id = messages_to_update.id WHERE mq_msgs.id = messages_to_update.id RETURNING mq_msgs.id, mq_msgs.commit_interval IS NULL, mq_payloads.name, mq_payloads.payload_json::TEXT, mq_payloads.payload_bytes, mq_msgs.retry_backoff / 2, interval '0' AS wait_time;
IF NOT FOUND THEN RETURN QUERY SELECT NULL::UUID, NULL::BOOLEAN, NULL::TEXT, NULL::TEXT, NULL::BYTEA, NULL::INTERVAL, MIN(mq_msgs.attempt_at) - NOW() FROM mq_msgs WHERE mq_msgs.id != '00000000-0000-0000-0000-000000000000'::uuid AND NOT mq_uuid_exists(mq_msgs.after_message_id) AND (channel_names IS NULL OR mq_msgs.channel_name = ANY(channel_names)); END IF;END;$$ LANGUAGE plpgsql;
-- Creates new messagesCREATE FUNCTION mq_insert(new_messages mq_new_t[])RETURNS VOID AS $$BEGIN PERFORM pg_notify(CONCAT('mq_', channel_name), '') FROM unnest(new_messages) AS new_msgs GROUP BY channel_name;
IF FOUND THEN PERFORM pg_notify('mq', ''); END IF;
INSERT INTO mq_payloads ( id, name, payload_json, payload_bytes ) SELECT id, name, payload_json::JSONB, payload_bytes FROM UNNEST(new_messages);
INSERT INTO mq_msgs ( id, attempt_at, attempts, retry_backoff, channel_name, channel_args, commit_interval, after_message_id ) SELECT id, NOW() + delay + COALESCE(commit_interval, INTERVAL '0'), retries + 1, retry_backoff, channel_name, channel_args, commit_interval, CASE WHEN ordered THEN LAG(id, 1, mq_latest_message(channel_name, channel_args)) OVER (PARTITION BY channel_name, channel_args, ordered ORDER BY id) ELSE NULL END FROM UNNEST(new_messages);END;$$ LANGUAGE plpgsql;
-- Commits messages previously created with a non-NULL commit interval.CREATE FUNCTION mq_commit(msg_ids UUID[])RETURNS VOID AS $$BEGIN UPDATE mq_msgs SET attempt_at = attempt_at - commit_interval, commit_interval = NULL WHERE id = ANY(msg_ids) AND commit_interval IS NOT NULL;END;$$ LANGUAGE plpgsql;
-- Deletes messages from the queue. This occurs when a message has been-- processed, or when it expires without being processed.CREATE FUNCTION mq_delete(msg_ids UUID[])RETURNS VOID AS $$BEGIN PERFORM pg_notify(CONCAT('mq_', channel_name), '') FROM mq_msgs WHERE id = ANY(msg_ids) AND after_message_id = '00000000-0000-0000-0000-000000000000'::uuid GROUP BY channel_name;
IF FOUND THEN PERFORM pg_notify('mq', ''); END IF;
DELETE FROM mq_msgs WHERE id = ANY(msg_ids); DELETE FROM mq_payloads WHERE id = ANY(msg_ids);END;$$ LANGUAGE plpgsql;
-- Can be called during the initial commit interval, or when processing-- a message. Indicates that the caller is still active and will prevent either-- the commit interval elapsing or the message being retried for the specified-- interval.CREATE FUNCTION mq_keep_alive(msg_ids UUID[], duration INTERVAL)RETURNS VOID AS $$ UPDATE mq_msgs SET attempt_at = NOW() + duration, commit_interval = commit_interval + ((NOW() + duration) - attempt_at) WHERE id = ANY(msg_ids) AND attempt_at < NOW() + duration;$$ LANGUAGE SQL;
-- Called during lengthy processing of a message to checkpoint the progress.-- As well as behaving like `mq_keep_alive`, the message payload can be-- updated.CREATE FUNCTION mq_checkpoint( msg_id UUID, duration INTERVAL, new_payload_json TEXT, new_payload_bytes BYTEA, extra_retries INT)RETURNS VOID AS $$ UPDATE mq_msgs SET attempt_at = GREATEST(attempt_at, NOW() + duration), attempts = attempts + COALESCE(extra_retries, 0) WHERE id = msg_id;
UPDATE mq_payloads SET payload_json = COALESCE(new_payload_json::JSONB, payload_json), payload_bytes = COALESCE(new_payload_bytes, payload_bytes) WHERE id = msg_id;$$ LANGUAGE SQL;