feat: support configurable number of bits (#2)

.gitignore

@@ -1,2 +1,3 @@
 build
-.idea
+.idea
+.vscode

README.md

@@ -23,15 +23,19 @@ Usage
 
     tttttttttttttttttttttttttttttttttttttttttdddddnnnnnsssssssssssss
 
-where
+where default
 
 * `s` is a 12-bit integer that increments when next_id() is called multiple times for the same millisecond
 * `n` is a 5-bit integer representing the node within a given data center
 * `d` is a 5-bit integer representing a unique data center or group of servers
 * `t` is a 42-bit integer representing the current timestamp in milliseconds
-    * the number of milliseconds to have elapsed since 1413817200000 or 2014-10-20T15:00:00.000Z
+    * the number of milliseconds to have elapsed since 1609459200000 or 2021-01-01T00:00:00.000Z
 
-`sf.init(datacenter_id, node_id)` is used to initialize snowflake and set values for ddddd and nnnnn as follows:
+`sf.init(datacenter_id, node_id, snowflake_epoc, node_id_bits, datacenter_id_bits, sequence_bits)` is used to initialize snowflake and set values for ddddd and nnnnn as follows:
 
-* `datacenter_id` is an integer n, where 0 ≤ n ≤ 0x1f and specifies the `ddddd` portion of the id
-* `node_id` is an integer n, where 0 ≤ n ≤ 0x1f and specifies the `nnnnn` portion of the id
+* `datacenter_id` is an integer n, where 0 ≤ n < (1 << datacenter_id_bits) and specifies the `ddddd` portion of the id
+* `node_id` is an integer n, where 0 ≤ n < (1 << node_id_bits) and specifies the `nnnnn` portion of the id
+* `snowflake_epoc` is an integer n, the starting timestamp is expressed in milliseconds
+* `node_id_bits` is an integer n, representing the node within a given data center
+* `datacenter_id_bits` is an integer n, representing a unique data center or group of servers
+* `sequence_bits` is an integer n, called multiple times for the same millisecond

rockspec/snowflake-1.0-1.rockspec → rockspec/api7-snowflake-2.0-1.rockspec

@@ -1,16 +1,16 @@
-package = "snowflake"
-version = "1.0-1"
+package = "api7-snowflake"
+version = "2.0-1"
 
 source = {
-	url = "git://github.com/stuartcarnie/lua-snowflake.git",
-	tag = "v1.0",
+	url = "git://github.com/api7/lua-snowflake.git",
+	tag = "v2.0",
 }
 
 description = {
 	summary = "An implementation of a distributed ID generator, similar to Snowflake by Twitter",
-	homepage = "http://github.com/stuartcarnie/lua-snowflake",
+	homepage = "http://github.com/api7/lua-snowflake",
 	license = "MIT",
-	maintainer = "Stuart Carnie",
+	maintainer = "dickens7",
 }
 
 dependencies = {
@@ -19,9 +19,8 @@ dependencies = {
 
 build = {
 	type = "builtin",
-
     modules = {
-        snowflake = {
+        ["snowflake"] = {
             sources = { "src/main.c" }
         }
     },

src/main.c

@@ -3,9 +3,8 @@
 #include <stdbool.h>
 #include <sys/time.h>
 
-
 #if LUA_VERSION_NUM < 502
-#define luaL_newlib(L, l) ( lua_newtable( L ), luaL_register( L, NULL, l ) )
+#define luaL_newlib(L, l) (lua_newtable(L), luaL_register(L, NULL, l))
 #endif
 
 static bool initialized = false;
@@ -15,17 +14,23 @@ static int g_node_id = 0;
 static long g_last_timestamp = -1;
 static int g_sequence = 0;
 
-// 2014-10-20T15:00:00Z
-#define SNOWFLAKE_EPOC 1413817200000L
+// 2021-01-01T00:00:00Z
+#define SNOWFLAKE_EPOC 1609459200000L
 
 #define NODE_ID_BITS 5
 #define DATACENTER_ID_BITS 5
 #define SEQUENCE_BITS 12
-#define NODE_ID_SHIFT SEQUENCE_BITS
-#define DATACENTER_ID_SHIFT (NODE_ID_SHIFT + NODE_ID_BITS)
-#define TIMESTAMP_SHIFT (DATACENTER_ID_SHIFT + DATACENTER_ID_BITS)
 
-#define SEQUENCE_MASK (0xffffffff ^ (0xffffffff << SEQUENCE_BITS))
+struct conf {
+    long snowflake_epoc;
+    int node_id_bits;
+    int datacenter_id_bits;
+    int sequence_bits;
+    int node_id_shift;
+    int datacenter_id_shift;
+    int timestamp_shift;
+    int sequence_mask;
+} conf;
 
 static long get_timestamp() {
     struct timeval tv;
@@ -44,16 +49,52 @@ static long get_til_next_millis(long last_timestamp) {
 }
 
 static int luasnowflake_init(lua_State *L) {
+    conf.snowflake_epoc = luaL_optlong(L, 3, SNOWFLAKE_EPOC);
+    struct timeval tv;
+    gettimeofday(&tv, NULL);
+    if (conf.snowflake_epoc < 0 || conf.snowflake_epoc > tv.tv_sec * 1000) {
+        char buffer[32];
+        snprintf(buffer, 32, "%ld", tv.tv_sec * 1000);
+        return luaL_error(L, "snowflake_epoc must be an long n where 1 ≤ n ≤ %s", buffer);
+    }
+
+    conf.node_id_bits = luaL_optint(L, 4, NODE_ID_BITS);
+    if (conf.node_id_bits < 1) {
+        return luaL_error(L, "node_id_bits must be an integer n where ≥ 1");
+    }
+
+    conf.datacenter_id_bits = luaL_optint(L, 5, DATACENTER_ID_BITS);
+    if (conf.datacenter_id_bits < 1) {
+        return luaL_error(L, "datacenter_id_bits must be an integer n where ≥ 1");
+    }
+
+    conf.sequence_bits = luaL_optint(L, 6, SEQUENCE_BITS);
+    if (conf.sequence_bits < 1) {
+        return luaL_error(L, "sequence_bits must be an integer n where ≥ 1");
+    }
+
+    // The timestamp >= 32 is because it is in milliseconds and only lasts 49 days when it equals 32 bits
+    if ((conf.node_id_bits + conf.datacenter_id_bits + conf.sequence_bits) > 32) {
+        return luaL_error(L, "(node_id_bits + datacenter_id_bits + sequence_bits) cannot be > 32");
+    }
+
     g_datacenter_id = luaL_checkint(L, 1);
-    if (g_datacenter_id < 0x00 || g_datacenter_id > 0x1f) {
-        return luaL_error(L, "datacenter_id must be an integer n, where 0 ≤ n ≤ 0x1f");
+    int max_datacenter_id_bits = (1 << conf.datacenter_id_bits) - 1;
+    if (g_datacenter_id < 0x00 || g_datacenter_id > max_datacenter_id_bits) {
+        return luaL_error(L, "datacenter_id must be an integer n, where 0 ≤ n ≤ %d", max_datacenter_id_bits);
     }
 
     g_node_id = luaL_checkint(L, 2);
-    if (g_node_id < 0x00 || g_node_id > 0x1f) {
-        return luaL_error(L, "node_id must be an integer n where 0 ≤ n ≤ 0x1f");
+    int max_node_id_bits = (1 << conf.node_id_bits) -1;
+    if (g_node_id < 0x00 || g_node_id > max_node_id_bits) {
+        return luaL_error(L, "node_id must be an integer n where 0 ≤ n ≤ %d", max_node_id_bits) ;
     }
 
+    conf.node_id_shift = conf.sequence_bits;
+    conf.datacenter_id_shift = (conf.node_id_shift + conf.node_id_bits);
+    conf.timestamp_shift = (conf.datacenter_id_shift + conf.datacenter_id_bits);
+    conf.sequence_mask = (0xffffffff ^ (0xffffffff << conf.sequence_bits));
+
     initialized = true;
 
     return 1;
@@ -67,7 +108,7 @@ static int luasnowflake_next_id(lua_State *L) {
     long ts = get_timestamp();
 
     if (g_last_timestamp == ts) {
-        g_sequence = (g_sequence + 1) & SEQUENCE_MASK;
+        g_sequence = (g_sequence + 1) & conf.sequence_mask;
         if (g_sequence == 0) {
             ts = get_til_next_millis(g_last_timestamp);
         }
@@ -76,9 +117,9 @@ static int luasnowflake_next_id(lua_State *L) {
     }
 
     g_last_timestamp = ts;
-    ts = ((ts - SNOWFLAKE_EPOC) << TIMESTAMP_SHIFT)
-            | (g_datacenter_id << DATACENTER_ID_SHIFT)
-            | (g_node_id << NODE_ID_SHIFT)
+    ts = ((ts - conf.snowflake_epoc) << conf.timestamp_shift)
+            | (g_datacenter_id << conf.datacenter_id_shift)
+            | (g_node_id << conf.node_id_shift)
             | g_sequence;
 
     char buffer[32];
@@ -89,13 +130,13 @@ static int luasnowflake_next_id(lua_State *L) {
 }
 
 static const struct luaL_Reg luasnowflake_lib[] = {
-        {"init", luasnowflake_init},
-        {"next_id", luasnowflake_next_id},
-        {NULL, NULL},
+    {"init", luasnowflake_init},
+    {"next_id", luasnowflake_next_id},
+    {NULL, NULL},
 };
 
 LUALIB_API int luaopen_snowflake(lua_State *const L) {
     luaL_newlib(L, luasnowflake_lib);
 
     return 1;
-}
+}

test.lua

@@ -1,4 +1,4 @@
 local sf = require "snowflake"
 
-sf.init(0x1f, 0x1f)
+sf.init(0x1f, 0x1f, nil, 5, 5, 12)
 print(sf.next_id())