From 7e87051962459cf7d006f8266de443943c23f045 Mon Sep 17 00:00:00 2001 From: Hayato Kuroda Date: Wed, 22 Nov 2023 08:42:34 +0000 Subject: [PATCH] initial commit for snowflake_sequence --- contrib/Makefile | 1 + contrib/meson.build | 1 + contrib/snowflake_sequence/.gitignore | 4 + contrib/snowflake_sequence/Makefile | 23 +++++ contrib/snowflake_sequence/meson.build | 25 +++++ .../snowflake_sequence--1.0.sql | 46 +++++++++ .../snowflake_sequence/snowflake_sequence.c | 94 +++++++++++++++++++ .../snowflake_sequence.control | 6 ++ 8 files changed, 200 insertions(+) create mode 100644 contrib/snowflake_sequence/.gitignore create mode 100644 contrib/snowflake_sequence/Makefile create mode 100644 contrib/snowflake_sequence/meson.build create mode 100644 contrib/snowflake_sequence/snowflake_sequence--1.0.sql create mode 100644 contrib/snowflake_sequence/snowflake_sequence.c create mode 100644 contrib/snowflake_sequence/snowflake_sequence.control diff --git a/contrib/Makefile b/contrib/Makefile index da4e2316a3..9540245b04 100644 --- a/contrib/Makefile +++ b/contrib/Makefile @@ -43,6 +43,7 @@ SUBDIRS = \ pg_walinspect \ postgres_fdw \ seg \ + snowflake_sequence \ spi \ tablefunc \ tcn \ diff --git a/contrib/meson.build b/contrib/meson.build index c0b267c632..3f4ae350ee 100644 --- a/contrib/meson.build +++ b/contrib/meson.build @@ -58,6 +58,7 @@ subdir('pg_walinspect') subdir('postgres_fdw') subdir('seg') subdir('sepgsql') +subdir('snowflake_sequence') subdir('spi') subdir('sslinfo') # start-scripts doesn't contain build products diff --git a/contrib/snowflake_sequence/.gitignore b/contrib/snowflake_sequence/.gitignore new file mode 100644 index 0000000000..5dcb3ff972 --- /dev/null +++ b/contrib/snowflake_sequence/.gitignore @@ -0,0 +1,4 @@ +# Generated subdirectories +/log/ +/results/ +/tmp_check/ diff --git a/contrib/snowflake_sequence/Makefile b/contrib/snowflake_sequence/Makefile new file mode 100644 index 0000000000..110f4e7ddb --- /dev/null +++ b/contrib/snowflake_sequence/Makefile @@ -0,0 +1,23 @@ +# contrib/snowflake_sequence/Makefile + +MODULE_big = snowflake_sequence +OBJS = \ + $(WIN32RES) \ + snowflake_sequence.o + +EXTENSION = snowflake_sequence +DATA = snowflake_sequence--1.0.sql +PGFILEDESC = "System-wide unique value generator" + +# REGRESS = snowflake_sequence + +ifdef USE_PGXS +PG_CONFIG = pg_config +PGXS := $(shell $(PG_CONFIG) --pgxs) +include $(PGXS) +else +subdir = contrib/snowflake_sequence +top_builddir = ../.. +include $(top_builddir)/src/Makefile.global +include $(top_srcdir)/contrib/contrib-global.mk +endif diff --git a/contrib/snowflake_sequence/meson.build b/contrib/snowflake_sequence/meson.build new file mode 100644 index 0000000000..7ae35d4a25 --- /dev/null +++ b/contrib/snowflake_sequence/meson.build @@ -0,0 +1,25 @@ +# Copyright (c) 2022-2023, PostgreSQL Global Development Group + +snowflake_sequence_sources = files( + 'snowflake_sequence.c', +) + +if host_system == 'windows' + snowflake_sequence_sources += rc_lib_gen.process(win32ver_rc, extra_args: [ + '--NAME', 'snowflake_sequence', + '--FILEDESC', 'System-wide unique value generator',]) +endif + +snowflake_sequence = shared_module('snowflake_sequence', + snowflake_sequence_sources, + kwargs: contrib_mod_args + { + 'dependencies': contrib_mod_args['dependencies'], + }, +) +contrib_targets += snowflake_sequence + +install_data( + 'snowflake_sequence.control', + 'snowflake_sequence--1.0.sql', + kwargs: contrib_data_args, +) diff --git a/contrib/snowflake_sequence/snowflake_sequence--1.0.sql b/contrib/snowflake_sequence/snowflake_sequence--1.0.sql new file mode 100644 index 0000000000..1712081a96 --- /dev/null +++ b/contrib/snowflake_sequence/snowflake_sequence--1.0.sql @@ -0,0 +1,46 @@ +/* contrib/snowflake_sequence/snowflake_sequence--1.0.sql */ + +-- complain if script is sourced in psql, rather than via CREATE EXTENSION +\echo Use "CREATE EXTENSION snowflake_sequence" to load this file. \quit + +-- Create a small table to record a Machine ID. +CREATE TABLE snowflake_sequence.machine_id ( + identifier smallint CHECK (identifier > 0), + CHECK (identifier < 512) +) WITH (user_catalog_table=true); + +-- Set a default value of the ID. Here, a random value is used. +INSERT INTO snowflake_sequence.machine_id + SELECT round((random() * (0 - 511))::numeric, 0) + 511; + +-- Create a snowflake sequence. Actually, this function just creates a normal +-- sequence in the snowflake_sequence schema. Other parts in snowflake id is +-- dynamically calculated so that we do not have to do anything. +CREATE FUNCTION snowflake_sequence.create_sequence(sequence_name name) +RETURNS void AS $$ +BEGIN + EXECUTE 'CREATE SEQUENCE snowflake_sequence.' || sequence_name || ' AS int MINVALUE 1 MAXVALUE 4095 CYCLE'; +END; +$$ LANGUAGE plpgsql; + + +-- Returns a nextval counted by a snowflake sequence. +CREATE FUNCTION snowflake_sequence.nextval(sequence_name name) +RETURNS bigint AS $$ +DECLARE + machine_id smallint; + ret bigint; +BEGIN + SELECT identifier FROM snowflake_sequence.machine_id INTO machine_id; + SELECT snowflake_sequence.snowflake_nextval_internal(sequence_name::text, machine_id) INTO ret; + + return ret; +END; +$$ LANGUAGE plpgsql; + +-- Internal function for nextval. Should not be called from users. +CREATE FUNCTION snowflake_sequence.snowflake_nextval_internal(sequence_name text, machine_id int) +RETURNS bigint +AS 'MODULE_PATHNAME' +LANGUAGE C; +REVOKE EXECUTE ON FUNCTION snowflake_sequence.snowflake_nextval_internal(text, int) FROM PUBLIC; diff --git a/contrib/snowflake_sequence/snowflake_sequence.c b/contrib/snowflake_sequence/snowflake_sequence.c new file mode 100644 index 0000000000..c2be75aa41 --- /dev/null +++ b/contrib/snowflake_sequence/snowflake_sequence.c @@ -0,0 +1,94 @@ +/*------------------------------------------------------------------------- + * + * snowflake_sequence.c + * Set of functions for generating system-wide unique values + * + * Copyright (c) 2023, PostgreSQL Global Development Group + * + * IDENTIFICATION + * snowflake_sequence/snowflake_sequence.c + * + * Snowflake ID is a globally unique identifier based on the time, machine ID, + * and local sequence number. This extension adds a new type of sequence for + * counting snowflake IDs. Currently, Snowflake ID is represented by a 64-bit + * integer, and its format is as follows: + * + * [1bit - unused] + * + [41bit - millisecond timestamp] + * + [10bit - machie ID] + * + [12bit - local sequence number] + *------------------------------------------------------------------------- + */ + +#include "postgres.h" +#include "fmgr.h" + + +#include "catalog/namespace.h" +#include "commands/sequence.h" +#include "nodes/makefuncs.h" +#include "utils/builtins.h" +#include "utils/timestamp.h" +#include "utils/varlena.h" + +PG_MODULE_MAGIC; + +#define LOCAL_SEQUENCE_BIT_LENGTH 12 +#define MACHINE_ID_BIT_LENGTH 10 + +#define TIMESTAMP_SHIFT_LENGTH (LOCAL_SEQUENCE_BIT_LENGTH + MACHINE_ID_BIT_LENGTH) + +static int snowflake_get_local_nextval(text *name); + +PG_FUNCTION_INFO_V1(snowflake_nextval_internal); + +/* + * Find a specified sequence from our schema, and get a next value. + * + * This function is basically ported from native nextval(). + */ +static int +snowflake_get_local_nextval(text *name) +{ + RangeVar *sequence; + Oid relid; + + sequence = makeRangeVar("snowflake_sequence", text_to_cstring(name), -1); + + /* + * XXX: This is not safe in the presence of concurrent DDL, but acquiring + * a lock here is more expensive than letting nextval_internal do it, + * since the latter maintains a cache that keeps us from hitting the lock + * manager more than once per transaction. It's not clear whether the + * performance penalty is material in practice, but for now, we do it this + * way. + */ + relid = RangeVarGetRelid(sequence, NoLock, false); + + return (int) nextval_internal(relid, true); +} + +/* + * Construct a snowflake ID and return it. + */ +Datum +snowflake_nextval_internal(PG_FUNCTION_ARGS) +{ + text *name = PG_GETARG_TEXT_PP(0); + int machine_id; + int64 millisecond_time; + int local_nextval; + int64 ret; + + /* Gather information used by snowflake ID */ + millisecond_time = GetCurrentTimestamp() / 1000; + machine_id = PG_GETARG_INT32(1); + local_nextval = snowflake_get_local_nextval(name); + + /* And construct them */ + ret = millisecond_time << (TIMESTAMP_SHIFT_LENGTH) | + machine_id << (LOCAL_SEQUENCE_BIT_LENGTH) | + local_nextval; + + PG_RETURN_INT64(ret); +} diff --git a/contrib/snowflake_sequence/snowflake_sequence.control b/contrib/snowflake_sequence/snowflake_sequence.control new file mode 100644 index 0000000000..42f118124e --- /dev/null +++ b/contrib/snowflake_sequence/snowflake_sequence.control @@ -0,0 +1,6 @@ +# snowflake_sequence extension +comment = 'System-wide unique value generator' +default_version = '1.0' +module_pathname = '$libdir/snowflake_sequence' +relocatable = false +schema = snowflake_sequence -- 2.27.0