sqlxmq/migrations/20210316025847_setup.up.sql

287 lines
8.5 KiB
MySQL
Raw Normal View History

2021-03-28 01:57:17 +00:00
-- The UDT for creating messages
CREATE 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 messages
CREATE 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 uuid_nil() REFERENCES mq_msgs(id) ON DELETE SET DEFAULT
);
-- Insert dummy message so that the 'nil' UUID can be referenced
INSERT INTO mq_msgs (id, channel_name, channel_args, after_message_id) VALUES (uuid_nil(), '', '', NULL);
-- Internal helper function to check that a UUID is neither NULL nor NIL
CREATE FUNCTION mq_uuid_exists(
id UUID
) RETURNS BOOLEAN AS $$
SELECT id IS NOT NULL AND id != uuid_nil()
$$ LANGUAGE SQL IMMUTABLE;
-- Index for polling
CREATE INDEX ON mq_msgs(channel_name, channel_args, attempt_at) WHERE id != uuid_nil() AND NOT mq_uuid_exists(after_message_id);
-- Index for adding messages
CREATE INDEX ON mq_msgs(channel_name, channel_args, created_at, id) WHERE id != uuid_nil() AND after_message_id IS NOT NULL;
-- Index for ensuring strict message order
CREATE UNIQUE INDEX ON mq_msgs(channel_name, channel_args, after_message_id);
-- Large, less frequently updated table of message payloads
CREATE 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 != uuid_nil()
ORDER BY created_at DESC, id DESC
LIMIT 1
),
uuid_nil()
)
$$ 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 != uuid_nil()
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 task 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 != uuid_nil()
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 != uuid_nil()
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 messages
CREATE 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 = uuid_nil()
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;