perf: rewrite using lua scripting when is possible (#15)

Author: Kikobeats

src/stats.js

@@ -1,18 +1,44 @@
 'use strict'
 
+const { promisify } = require('util')
+const stream = require('stream')
+
+const { Transform } = stream
+
+const pipeline = promisify(stream.pipeline)
+
 const formatYYYMMDDDate = (now = new Date()) => {
   const year = now.getFullYear()
   const month = String(now.getMonth() + 1).padStart(2, '0')
   const day = String(now.getDate()).padStart(2, '0')
   return `${year}-${month}-${day}`
 }
 
+/**
+ * 90 days in milliseconds
+ */
+const TTL = 90 * 24 * 60 * 60 * 1000
+
+/**
+ * Lua script to increment a key and set an expiration time if it doesn't have one.
+ * This script is necessary to perform the operation atomically.
+ */
+const LUA_INCREMENT_AND_EXPIRE = `
+local key = KEYS[1]
+local ttl = ARGV[1]
+local quantity = ARGV[2]
+local current = redis.call('incrby', key, quantity)
+if tonumber(redis.call('ttl', key)) == -1 then
+  redis.call('expire', key, ttl)
+end
+return current
+`
+
 module.exports = ({ redis, prefix }) => {
   const prefixKey = key => `${prefix}stats:${key}`
 
-  const increment = async keyValue => {
-    redis.incr(`${prefixKey(keyValue)}:${formatYYYMMDDDate()}`)
-  }
+  const increment = (keyValue, quantity = 1) =>
+    redis.eval(LUA_INCREMENT_AND_EXPIRE, 1, `${prefixKey(keyValue)}:${formatYYYMMDDDate()}`, TTL, quantity)
 
   /**
    * Get stats for a given key.
@@ -30,15 +56,24 @@ module.exports = ({ redis, prefix }) => {
    * // ]
    */
   const stats = async keyValue => {
-    const keys = await redis.keys(prefixKey(`${keyValue}*`))
     const stats = []
+    const stream = redis.scanStream({ match: `${prefixKey(keyValue)}*` })
+    const dataHandler = new Transform({
+      objectMode: true,
+      transform: async (keys, _, next) => {
+        if (keys.length) {
+          const values = await redis.mget.apply(redis, keys)
+          const statsPart = keys.map((key, i) => ({
+            date: key.replace(`${prefixKey(keyValue)}:`, ''),
+            count: Number(values[i])
+          }))
+          stats.push.apply(stats, statsPart)
+        }
+        next()
+      }
+    })
 
-    for (const key of keys) {
-      const date = key.replace(`${prefixKey(keyValue)}:`, '')
-      const count = Number(await redis.get(key))
-      stats.push({ date, count })
-    }
-
+    await pipeline(stream, dataHandler)
     return stats.sort((a, b) => a.date.localeCompare(b.date))
   }
 

src/usage.js

@@ -29,7 +29,6 @@ module.exports = ({ plans, keys, redis, stats, prefix, serialize, deserialize })
 
     let usage = deserialize(await redis.get(prefixKey(keyValue)))
 
-    // TODO: move into lua script
     if (usage === null) {
       usage = {
         count: quantity,
@@ -45,7 +44,9 @@ module.exports = ({ plans, keys, redis, stats, prefix, serialize, deserialize })
     }
 
     const pending =
-      quantity > 0 && Promise.all([redis.set(prefixKey(keyValue), serialize(usage)), stats.increment(keyValue)])
+      quantity > 0
+        ? Promise.all([redis.set(prefixKey(keyValue), serialize(usage)), stats.increment(keyValue, quantity)])
+        : Promise.resolve([])
 
     return {
       limit: plan.limit,

test/stats.js

@@ -24,7 +24,7 @@ testCleanup({
     ])
 })
 
-test('.stats', async t => {
+test.serial('.increment # by one', async t => {
   const plan = await openkey.plans.create({
     id: randomUUID(),
     limit: 3,
@@ -59,3 +59,31 @@ test('.stats', async t => {
     { date: openkey.stats.formatYYYMMDDDate(addDays(Date.now(), 2)), count: 5 }
   ])
 })
+
+test.serial('.increment # by more than one', async t => {
+  const plan = await openkey.plans.create({
+    id: randomUUID(),
+    limit: 3,
+    period: '100ms'
+  })
+
+  const key = await openkey.keys.create({ plan: plan.id })
+  const data = await openkey.usage.increment(key.value)
+  await data.pending
+
+  Date.setNow(addDays(Date.now(), 1))
+  await openkey.usage.increment(key.value, 10)
+  await data.pending
+
+  Date.setNow(addDays(Date.now(), 1))
+  await openkey.usage.increment(key.value, 5)
+  await data.pending
+
+  Date.setNow()
+
+  t.deepEqual(await openkey.stats(key.value), [
+    { date: openkey.stats.formatYYYMMDDDate(), count: 1 },
+    { date: openkey.stats.formatYYYMMDDDate(addDays(Date.now(), 1)), count: 10 },
+    { date: openkey.stats.formatYYYMMDDDate(addDays(Date.now(), 2)), count: 5 }
+  ])
+})

test/usage.js

@@ -31,7 +31,6 @@ test('.increment', async t => {
 
   const key = await openkey.keys.create({ plan: plan.id })
   let data = await openkey.usage(key.value)
-  t.is(data.pending, false)
   t.is(data.remaining, 3)
   data = await openkey.usage.increment(key.value)
   await data.pending