如何在PostgreSQL中使用C语言UDF生成Snowflake ID

我完全懂你的需求——有时候确实需要直接在数据库层面生成Snowflake ID，不用绕到应用层去处理。你已经在SQL Server用C# CLR、MySQL用C++搞定了Snowflake ID的生成逻辑（也就是你说的SnowSeed），现在卡在PostgreSQL用C写UDF这一步，我来给你详细捋清楚怎么做。

PostgreSQL支持用C语言编写自定义函数（UDF），而且性能比PL/pgSQL之类的脚本语言高很多，非常适合实现Snowflake这种需要高效生成的逻辑。下面是完整的实现步骤：

一、编写C语言的Snowflake生成代码 #

首先要写适配PostgreSQL环境的C代码，核心是实现Snowflake算法，同时要注意PostgreSQL的多线程/多进程环境下的线程安全，避免生成重复ID。

这里给你一个示例代码，包含了基础的Snowflake逻辑，还支持通过配置参数设置worker id（避免不同节点重复）：

#include "postgres.h"
#include "fmgr.h"
#include "utils/timestamp.h"
#include "utils/guc.h"
#include <stdint.h>
#include <sys/time.h>

// Snowflake基础配置
#define EPOCH 1609459200000LL // 自定义起始时间（2021-01-01 00:00:00 UTC），可按需修改
#define MAX_WORKER_ID 1023    // Worker ID最大取值（10位）
#define MAX_SEQUENCE 4095     // 序列号最大取值（12位）

// 全局变量：Worker ID（通过GUC参数配置）
static int worker_id = 1;
// 线程安全的序列号，用PostgreSQL原子操作保证并发安全
static pg_atomic_uint32 sequence = PG_ATOMIC_INIT(0);

// 模块标识，PostgreSQL要求必须有
PG_MODULE_MAGIC;

// 初始化函数：注册GUC参数，让Worker ID可以动态配置
void _PG_init(void) {
    DefineCustomIntVariable(
        "snowflake.worker_id",
        "Worker ID for Snowflake ID generator (range: 0-1023)",
        "Each PostgreSQL node should have a unique worker_id to avoid duplicate IDs",
        &worker_id,
        1,
        0,
        MAX_WORKER_ID,
        PGC_SIGHUP,
        0,
        NULL,
        NULL,
        NULL
    );
}

// 获取当前毫秒级时间戳
static int64_t get_current_timestamp_ms(void) {
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return (int64_t)tv.tv_sec * 1000 + tv.tv_usec / 1000;
}

// 核心生成函数：PostgreSQL UDF入口
PG_FUNCTION_INFO_V1(snowflake_generate);

Datum snowflake_generate(PG_FUNCTION_ARGS) {
    // 计算相对于起始epoch的时间差（毫秒）
    int64_t timestamp = get_current_timestamp_ms() - EPOCH;
    if (timestamp < 0) {
        ereport(ERROR, (errcode(ERRCODE_INVALID_PARAMETER_VALUE), 
                        errmsg("Snowflake epoch is set to a future time")));
    }

    // 原子递增序列号，超过最大值则等待下一毫秒
    uint32_t seq = pg_atomic_fetch_add_u32(&sequence, 1);
    if (seq > MAX_SEQUENCE) {
        // 等待时间推进到下一毫秒
        while (get_current_timestamp_ms() - EPOCH == timestamp) {
            pg_usleep(100); // 休眠100微秒，避免空转占用CPU
        }
        // 重置序列号为0，重新获取时间戳
        pg_atomic_set_u32(&sequence, 0);
        seq = 0;
        timestamp = get_current_timestamp_ms() - EPOCH;
    }

    // 组装Snowflake ID：41位时间戳 + 10位Worker ID + 12位序列号
    int64_t snowflake_id = (timestamp << 22) | ((int64_t)worker_id << 12) | seq;

    PG_RETURN_INT64(snowflake_id);
}

二、编写扩展配置文件 #

要把这个C代码打包成PostgreSQL的扩展，需要两个额外文件：

1. Makefile #

用来编译代码成PostgreSQL可识别的模块：

MODULES = snowflake
EXTENSION = snowflake
DATA = snowflake--1.0.sql
PG_CONFIG = pg_config
PGXS := $(shell $(PG_CONFIG) --pgxs)
include $(PGXS)

2. snowflake--1.0.sql #

用来定义UDF的SQL接口：

CREATE OR REPLACE FUNCTION snowflake_generate()
RETURNS bigint
LANGUAGE C
STABLE
PARALLEL SAFE;

三、编译安装扩展 #

把上面三个文件（snowflake.c、Makefile、snowflake--1.0.sql）放在同一个目录下，然后执行以下命令：

# 编译
make
# 安装到PostgreSQL的扩展目录
sudo make install

四、在PostgreSQL中启用并使用 #

登录PostgreSQL数据库，执行以下命令启用扩展：

-- 创建扩展
CREATE EXTENSION snowflake;

-- （可选）配置当前节点的Worker ID，确保不同节点唯一
ALTER SYSTEM SET snowflake.worker_id = 2;
-- 重载配置使生效
SELECT pg_reload_conf();

现在就可以调用函数生成Snowflake ID了：

SELECT snowflake_generate();

一些注意事项 #

• Worker ID唯一性：每个PostgreSQL节点的worker_id必须唯一，否则不同节点可能生成重复的Snowflake ID。
• 时间回拨问题：如果服务器时间发生回拨（比如同步时间导致），代码里会报错或者等待时间追上，避免生成重复ID，你可以根据实际需求调整这部分逻辑。
• 性能优化：这个实现用了原子操作保证并发安全，在高并发场景下也能稳定生成ID，性能比PL/pgSQL实现高很多。

备注：内容来源于stack exchange，提问作者研小艾