feat: ♻️ refactored snowflake generation logic

Updates to imports. Closes #17
AkashRajpurohit · Oct 31, 2023 · 73bb6cf · 73bb6cf
1 parent e65cdaf
commit 73bb6cf
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 157 deletions.
diff --git a/src/helpers.ts b/src/helpers.ts
@@ -6,26 +6,13 @@ export const waitUntilNextTimestamp = (currentTimestamp: number) => {
   return nextTimestamp;
 };
 
-export const getValidNodeId = (newNodeId: number, nodeIdBits: number) => {
-  const maxNodeId = 1 << nodeIdBits;
-  let nodeId;
-
-  if (typeof newNodeId !== 'number' || Number.isNaN(newNodeId)) {
-    console.warn(`Invalid node ID provided: ${newNodeId}, using default ID: 0`);
-    nodeId = 0;
-  } else {
-    nodeId = Math.floor(newNodeId) % maxNodeId;
-    if (nodeId < 0) {
-      nodeId = maxNodeId - Math.abs(nodeId);
-    }
-  }
-
-  return nodeId;
-};
-
 export const DEFAULTS = {
   WORKER_ID: 0,
-  NODE_ID_BITS: 12,
-  SEQUENCE_BITS: 10,
   EPOCH: 1597017600000, // August 10, 2020 at 00:00:00 UTC
 };
+
+export const CONFIG = {
+  TIMESTAMP_BITS: 42,
+  WORKER_ID_BITS: 10,
+  SEQUENCE_BITS: 12,
+};
diff --git a/src/snowflake.ts b/src/snowflake.ts
@@ -1,35 +1,15 @@
-import { DEFAULTS, getValidNodeId, waitUntilNextTimestamp } from './helpers';
+import { CONFIG, DEFAULTS, waitUntilNextTimestamp } from './helpers';
 
-const SnowflakeId = ({
-  workerId = DEFAULTS.WORKER_ID,
-  nodeIdBits = DEFAULTS.NODE_ID_BITS,
-  sequenceBits = DEFAULTS.SEQUENCE_BITS,
-  epoch = DEFAULTS.EPOCH,
-} = {}) => {
-  if (typeof workerId !== 'number' || Number.isNaN(workerId)) {
-    console.warn(`Invalid worker ID provided: ${workerId}, using default ID: ${DEFAULTS.WORKER_ID}`);
-    workerId = DEFAULTS.WORKER_ID;
-  }
-
-  if (typeof nodeIdBits !== 'number' || Number.isNaN(nodeIdBits) || nodeIdBits < 1 || nodeIdBits > 31) {
-    console.warn(`Invalid node ID bits provided: ${nodeIdBits}, using default value: ${DEFAULTS.NODE_ID_BITS}`);
-    nodeIdBits = DEFAULTS.NODE_ID_BITS;
-  }
-
-  if (typeof sequenceBits !== 'number' || Number.isNaN(sequenceBits) || sequenceBits < 1 || sequenceBits > 22) {
-    console.warn(`Invalid sequence bits provided: ${sequenceBits}, using default value: ${DEFAULTS.SEQUENCE_BITS}`);
-    sequenceBits = DEFAULTS.SEQUENCE_BITS;
-  }
+const SnowflakeId = ({ workerId = DEFAULTS.WORKER_ID, epoch = DEFAULTS.EPOCH } = {}) => {
+  const currentTimestamp = Date.now();
 
-  if (typeof epoch !== 'number' || Number.isNaN(epoch)) {
-    console.warn(`Invalid epoch provided: ${epoch}, using default value: ${DEFAULTS.EPOCH}`);
-    epoch = DEFAULTS.EPOCH;
+  if (epoch > currentTimestamp) {
+    throw new Error(`Invalid epoch: ${epoch}, it can't be greater than the current timestamp!`);
   }
 
   let lastTimestamp = -1;
   let sequence = 0;
-  let nodeId = getValidNodeId(workerId, nodeIdBits);
-  const maxSequence = (1 << sequenceBits) - 1;
+  const maxSequence = (1 << CONFIG.SEQUENCE_BITS) - 1;
 
   function generate() {
     let timestamp = Date.now();
@@ -49,13 +29,16 @@ const SnowflakeId = ({
 
     lastTimestamp = timestamp;
 
-    const high =
-      (BigInt(timestamp - epoch) << BigInt(nodeIdBits + sequenceBits)) |
-      (BigInt(nodeId) << BigInt(sequenceBits)) |
-      BigInt(sequence);
-    const low = BigInt(0);
-    const snowflakeId = (high << BigInt(32)) | low;
-    return snowflakeId.toString();
+    const timestampOffset = timestamp - epoch;
+
+    const timestampBits = timestampOffset.toString(2).padStart(CONFIG.TIMESTAMP_BITS, '0');
+    const workerIdBits = workerId.toString(2).padStart(CONFIG.WORKER_ID_BITS, '0');
+    const sequenceBits = sequence.toString(2).padStart(CONFIG.SEQUENCE_BITS, '0');
+
+    const idBinary = `${timestampBits}${workerIdBits}${sequenceBits}`;
+    const idDecimal = BigInt('0b' + idBinary).toString();
+
+    return idDecimal.toString();
   }
 
   return {

diff --git a/tests/helpers.test.ts b/tests/helpers.test.ts
@@ -1,36 +1,8 @@
 import { describe, it, expect } from 'vitest';
 
-import { getValidNodeId, waitUntilNextTimestamp } from '~/helpers';
+import { waitUntilNextTimestamp } from '~/helpers';
 
 describe('helpers', () => {
-  describe('getValidNodeId', () => {
-    it('should return 0 if an invalid nodeId is provided', () => {
-      const nodeId = getValidNodeId(NaN, 10);
-
-      expect(nodeId).toBe(0);
-    });
-
-    it('should return the same nodeId if a valid nodeId is provided', () => {
-      const nodeId = getValidNodeId(512, 10);
-
-      expect(nodeId).toBe(512);
-    });
-
-    it('should return a valid nodeId within the specified range', () => {
-      const nodeId = getValidNodeId(2048, 10);
-
-      expect(nodeId).toBeGreaterThanOrEqual(0);
-      expect(nodeId).toBeLessThanOrEqual(2 ** 10);
-    });
-
-    it('should return a valid nodeId within the specified range for negative numbers', () => {
-      const nodeId = getValidNodeId(-100, 10);
-
-      expect(nodeId).toBeGreaterThanOrEqual(0);
-      expect(nodeId).toBeLessThanOrEqual(2 ** 10);
-    });
-  });
-
   describe('waitUntilNextTimestamp', () => {
     it('should return a timestamp greater than the current timestamp', () => {
       const currentTimestamp = Date.now();

diff --git a/tests/snowflake.test.ts b/tests/snowflake.test.ts
@@ -1,9 +1,8 @@
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
-import { DEFAULTS } from '~/helpers';
 
 import SnowflakeId from '~/snowflake';
 
-describe('snowflake', () => {
+describe.only('snowflake', () => {
   describe('Invalid input parameters check', () => {
     beforeEach(() => {
       // tell vitest we use mocked time
@@ -15,83 +14,23 @@ describe('snowflake', () => {
       vi.useRealTimers();
     });
 
-    it('should use default workerId if invalid workedId is passed', () => {
-      const logSpy = vi.spyOn(global.console, 'warn');
-
-      SnowflakeId({
-        workerId: NaN,
-      });
-
-      expect(logSpy).toHaveBeenCalled();
-      expect(logSpy).toBeCalledWith(`Invalid worker ID provided: NaN, using default ID: ${DEFAULTS.WORKER_ID}`);
-    });
-
-    it('should use default NODE_ID_BITS if invalid nodeIdBits is passed', () => {
-      const logSpy = vi.spyOn(global.console, 'warn');
-
-      SnowflakeId({
-        workerId: 121,
-        nodeIdBits: NaN,
-      });
-
-      expect(logSpy).toHaveBeenCalled();
-      expect(logSpy).toBeCalledWith(
-        `Invalid node ID bits provided: NaN, using default value: ${DEFAULTS.NODE_ID_BITS}`
-      );
-    });
-
-    it('should use default NODE_ID_BITS if nodeIdBits provided is not in the range', () => {
-      const logSpy = vi.spyOn(global.console, 'warn');
-
-      SnowflakeId({
-        workerId: 121,
-        nodeIdBits: 35,
-      });
-
-      expect(logSpy).toHaveBeenCalled();
-      expect(logSpy).toBeCalledWith(`Invalid node ID bits provided: 35, using default value: ${DEFAULTS.NODE_ID_BITS}`);
-    });
-
-    it('should use default SEQUENCE_BITS if invalid sequenceBits is passed', () => {
-      const logSpy = vi.spyOn(global.console, 'warn');
-
-      SnowflakeId({
-        workerId: 121,
-        sequenceBits: NaN,
-      });
-
-      expect(logSpy).toHaveBeenCalled();
-      expect(logSpy).toBeCalledWith(
-        `Invalid sequence bits provided: NaN, using default value: ${DEFAULTS.SEQUENCE_BITS}`
-      );
-    });
-
-    it('should use default SEQUENCE_BITS if sequenceBits provided is not in the range', () => {
-      const logSpy = vi.spyOn(global.console, 'warn');
-
-      SnowflakeId({
-        workerId: 121,
-        sequenceBits: 35,
-      });
+    it('should throw error if epoch is greater than current timestamp', () => {
+      let error: Error | null = null;
+      try {
+        SnowflakeId({
+          workerId: 121,
+          epoch: Date.now() + 10000,
+        });
+      } catch (err) {
+        error = err as Error;
+      }
 
-      expect(logSpy).toHaveBeenCalled();
-      expect(logSpy).toBeCalledWith(
-        `Invalid sequence bits provided: 35, using default value: ${DEFAULTS.SEQUENCE_BITS}`
+      expect(error).toBeInstanceOf(Error);
+      expect(error?.message).toEqual(
+        `Invalid epoch: ${Date.now() + 10000}, it can't be greater than the current timestamp!`
       );
     });
 
-    it('should use default EPOCH if invalid epoch is passed', () => {
-      const logSpy = vi.spyOn(global.console, 'warn');
-
-      SnowflakeId({
-        workerId: 121,
-        epoch: NaN,
-      });
-
-      expect(logSpy).toHaveBeenCalled();
-      expect(logSpy).toBeCalledWith(`Invalid epoch provided: NaN, using default value: ${DEFAULTS.EPOCH}`);
-    });
-
     it('should throw error if currentTimeStamp is less than lastTimeStamp', () => {
       let error: Error | null = null;
       try {
@@ -120,10 +59,15 @@ describe('snowflake', () => {
       workerId: 121,
     });
 
-    it('should generate a new id', () => {
+    it('should generate a new valid snowflake id', () => {
       const id = snowflakeId.generate();
 
       expect(id).toBeTypeOf('string');
+
+      // Generated id should be at max 64 bits
+      // We can padStart with 0 for the binary representation
+      const idBinary = BigInt(id).toString(2);
+      expect(idBinary.length).toBeLessThanOrEqual(64);
     });
 
     it('should not generate same id ever', () => {