pguid/pguid.c at main · ats1999/pguid · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
#include "postgres.h"
#include "fmgr.h"
#include "utils/guc.h"
#include "storage/lwlock.h"
#include "storage/shmem.h"
#include "miscadmin.h"
#include <sys/time.h>

PG_MODULE_MAGIC;

/* ---- External hook declaration ---- */
extern void (*shmem_startup_hook)(void);

/* ---- Config ---- */
static int node_id = 0;

/* ---- Bit allocation ---- */
#define SEQ_BITS 12
#define NODE_BITS 10

#define MAX_SEQ ((1 << SEQ_BITS) - 1)
#define MAX_NODE ((1 << NODE_BITS) - 1)

/* ---- Epoch (custom) ---- */
#define EPOCH 1704067200000LL  /* Jan 1 2024 */

/* ---- Shared state ---- */
typedef struct {
    int64 last_ts;
    uint16 seq;
    LWLock lock;
} pguid_state;

static pguid_state *state = NULL;

/* ---- Time in ms ---- */
static int64
current_time_millis(void)
{
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return (int64) tv.tv_sec * 1000 + tv.tv_usec / 1000;
}

/* ---- Shmem init ---- */
static void
pguid_shmem_startup(void)
{
    bool found;

    state = ShmemInitStruct("pguid_state",
                            sizeof(pguid_state),
                            &found);

    if (!found)
    {
        state->last_ts = 0;
        state->seq = 0;
        LWLockInitialize(&state->lock, LWLockNewTrancheId());
    }
}

/* ---- Module init ---- */

void _PG_init(void);


void
_PG_init(void)
{
    DefineCustomIntVariable(
        "pguid.node_id",
        "Node ID for distributed ID generator",
        NULL,
        &node_id,
        0,
        0,
        MAX_NODE,
        PGC_POSTMASTER,
        0,
        NULL,
        NULL,
        NULL
    );

    RequestAddinShmemSpace(sizeof(pguid_state));

    shmem_startup_hook = pguid_shmem_startup;
}

/* ---- ID generation ---- */
PG_FUNCTION_INFO_V1(pguid);

Datum
pguid(PG_FUNCTION_ARGS)
{
    int64 now, ts;
    uint16 seq;
    int64 id;

    if (state == NULL)
        ereport(ERROR, (errmsg("pguid shared memory not initialized")));

    if (node_id < 0 || node_id > MAX_NODE)
        ereport(ERROR, (errmsg("invalid node_id")));

    now = current_time_millis() - EPOCH;

    LWLockAcquire(&state->lock, LW_EXCLUSIVE);

    if (now > state->last_ts)
    {
        state->last_ts = now;
        state->seq = 0;
    }
    else
    {
        state->seq++;

        if (state->seq > MAX_SEQ)
        {
            state->last_ts++;
            state->seq = 0;
        }
    }

    ts = state->last_ts;
    seq = state->seq;

    LWLockRelease(&state->lock);

    id = (ts << (NODE_BITS + SEQ_BITS))
       | ((int64) node_id << SEQ_BITS)
       | seq;

    PG_RETURN_INT64(id);
}