Skip to content

Commit

Permalink
Use advisory lock when assigning sequence IDs (#401)
Browse files Browse the repository at this point in the history
Hoping for this to be an invisible change - relying on existing unit tests for verification.

Whenever we add a table that relies on a sequence for ordering, we should consider adding a Postgres stored function that performs insertions by taking an advisory lock first. This ensures that inserts on the table are [processed serially](https://espadrine.github.io/blog/posts/two-postgresql-sequence-misconceptions.html#Order_violation).

For these stored functions, I've matched the params and param ordering to the existing queries. I also allowed unique constraint errors to bubble up rather than specifying `ON CONFLICT`, to match the existing queries.

The `dev/up` script also generates the protos, and I don't see a convenient way to exclude the protos for the replication project. However, I assume that at some point we will need them on this server, for forwards-compatibility purposes.
  • Loading branch information
richardhuaaa authored Aug 29, 2024
1 parent f7a54a0 commit 987ef53
Show file tree
Hide file tree
Showing 45 changed files with 3,775 additions and 732 deletions.
6 changes: 6 additions & 0 deletions pkg/migrations/mls/20240829001344_serial-ids.down.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
DROP FUNCTION insert_group_message;

DROP FUNCTION insert_welcome_message;

DROP FUNCTION insert_inbox_log;

44 changes: 44 additions & 0 deletions pkg/migrations/mls/20240829001344_serial-ids.up.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
CREATE FUNCTION insert_group_message(group_id BYTEA, data BYTEA, group_id_data_hash BYTEA)
RETURNS SETOF group_messages
AS $$
BEGIN
-- Ensures that the generated sequence ID matches the insertion order
-- Only released at the end of the enclosing transaction - beware if called within a long transaction
PERFORM
pg_advisory_xact_lock(hashtext('group_messages_sequence'), hashtext(encode(group_id, 'hex')));
RETURN QUERY INSERT INTO group_messages(group_id, data, group_id_data_hash)
VALUES(group_id, data, group_id_data_hash)
RETURNING
*;
END;
$$
LANGUAGE plpgsql;

CREATE FUNCTION insert_welcome_message(installation_key BYTEA, data BYTEA, installation_key_data_hash BYTEA, hpke_public_key BYTEA)
RETURNS SETOF welcome_messages
AS $$
BEGIN
PERFORM
pg_advisory_xact_lock(hashtext('welcome_messages_sequence'), hashtext(encode(installation_key, 'hex')));
RETURN QUERY INSERT INTO welcome_messages(installation_key, data, installation_key_data_hash, hpke_public_key)
VALUES(installation_key, data, installation_key_data_hash, hpke_public_key)
RETURNING
*;
END;
$$
LANGUAGE plpgsql;

CREATE FUNCTION insert_inbox_log(inbox_id BYTEA, server_timestamp_ns BIGINT, identity_update_proto BYTEA)
RETURNS SETOF inbox_log
AS $$
BEGIN
PERFORM
pg_advisory_xact_lock(hashtext('inbox_log_sequence'), hashtext(encode(inbox_id, 'hex')));
RETURN QUERY INSERT INTO inbox_log(inbox_id, server_timestamp_ns, identity_update_proto)
VALUES(inbox_id, server_timestamp_ns, identity_update_proto)
RETURNING
*;
END;
$$
LANGUAGE plpgsql;

24 changes: 12 additions & 12 deletions pkg/mls/store/queries.sql
Original file line number Diff line number Diff line change
Expand Up @@ -59,10 +59,10 @@ RETURNING
*;

-- name: InsertInboxLog :one
INSERT INTO inbox_log(inbox_id, server_timestamp_ns, identity_update_proto)
VALUES (decode(@inbox_id, 'hex'), @server_timestamp_ns, @identity_update_proto)
RETURNING
sequence_id;
SELECT
sequence_id
FROM
insert_inbox_log(decode(@inbox_id, 'hex'), @server_timestamp_ns, @identity_update_proto);

-- name: RevokeAddressFromLog :exec
UPDATE
Expand Down Expand Up @@ -111,16 +111,16 @@ WHERE
id = ANY (@installation_ids::BYTEA[]);

-- name: InsertGroupMessage :one
INSERT INTO group_messages(group_id, data, group_id_data_hash)
VALUES ($1, $2, $3)
RETURNING
*;
SELECT
*
FROM
insert_group_message(@group_id, @data, @group_id_data_hash);

-- name: InsertWelcomeMessage :one
INSERT INTO welcome_messages(installation_key, data, installation_key_data_hash, hpke_public_key)
VALUES ($1, $2, $3, $4)
RETURNING
*;
SELECT
*
FROM
insert_welcome_message(@installation_key, @data, @installation_key_data_hash, @hpke_public_key);

-- name: GetAllGroupMessages :many
SELECT
Expand Down
2 changes: 1 addition & 1 deletion pkg/mls/store/queries/db.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion pkg/mls/store/queries/models.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

20 changes: 10 additions & 10 deletions pkg/mls/store/queries/queries.sql.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

28 changes: 14 additions & 14 deletions pkg/proto/identity/api/v1/identity.pb.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading

0 comments on commit 987ef53

Please sign in to comment.