feat(locking): Locking module (#9524)

**What**
- Locking Module to manage concurrency
- Default `in-memory` provider

packages/modules/locking/.gitignore

@@ -0,0 +1,6 @@
+/dist
+node_modules
+.DS_store
+.env*
+.env
+*.sql

packages/modules/locking/integration-tests/__tests__/index.spec.ts

@@ -0,0 +1,122 @@
+import { ILockingModule } from "@medusajs/framework/types"
+import { Modules } from "@medusajs/framework/utils"
+import { moduleIntegrationTestRunner } from "medusa-test-utils"
+import { setTimeout } from "node:timers/promises"
+
+jest.setTimeout(10000)
+
+moduleIntegrationTestRunner<ILockingModule>({
+  moduleName: Modules.LOCKING,
+  testSuite: ({ service }) => {
+    describe("Locking Module Service", () => {
+      let stock = 5
+      function replenishStock() {
+        stock = 5
+      }
+      function hasStock() {
+        return stock > 0
+      }
+      async function reduceStock() {
+        await setTimeout(10)
+        stock--
+      }
+      async function buy() {
+        if (hasStock()) {
+          await reduceStock()
+          return true
+        }
+        return false
+      }
+
+      it("should execute functions respecting the key locked", async () => {
+        // 10 parallel calls to buy should oversell the stock
+        const prom: any[] = []
+        for (let i = 0; i < 10; i++) {
+          prom.push(buy())
+        }
+        await Promise.all(prom)
+        expect(stock).toBe(-5)
+
+        replenishStock()
+
+        // 10 parallel calls to buy with lock should not oversell the stock
+        const promWLock: any[] = []
+        for (let i = 0; i < 10; i++) {
+          promWLock.push(service.execute("item_1", buy))
+        }
+        await Promise.all(promWLock)
+
+        expect(stock).toBe(0)
+      })
+
+      it("should acquire lock and release it", async () => {
+        await service.acquire("key_name", {
+          ownerId: "user_id_123",
+        })
+
+        const userReleased = await service.release("key_name", {
+          ownerId: "user_id_456",
+        })
+        const anotherUserLock = service.acquire("key_name", {
+          ownerId: "user_id_456",
+        })
+
+        expect(userReleased).toBe(false)
+        await expect(anotherUserLock).rejects.toThrowError(
+          `"key_name" is already locked.`
+        )
+
+        const releasing = await service.release("key_name", {
+          ownerId: "user_id_123",
+        })
+
+        expect(releasing).toBe(true)
+      })
+
+      it("should acquire lock and release it during parallel calls", async () => {
+        const keyToLock = "mySpecialKey"
+        const user_1 = {
+          ownerId: "user_id_456",
+        }
+        const user_2 = {
+          ownerId: "user_id_000",
+        }
+
+        expect(service.acquire(keyToLock, user_1)).resolves.toBeUndefined()
+
+        expect(service.acquire(keyToLock, user_1)).resolves.toBeUndefined()
+
+        expect(service.acquire(keyToLock, user_2)).rejects.toThrowError(
+          `"${keyToLock}" is already locked.`
+        )
+
+        expect(service.acquire(keyToLock, user_2)).rejects.toThrowError(
+          `"${keyToLock}" is already locked.`
+        )
+
+        await service.acquire(keyToLock, user_1)
+
+        const releaseNotLocked = await service.release(keyToLock, {
+          ownerId: "user_id_000",
+        })
+        expect(releaseNotLocked).toBe(false)
+
+        const release = await service.release(keyToLock, user_1)
+        expect(release).toBe(true)
+      })
+    })
+
+    it("should release lock in case of failure", async () => {
+      const fn_1 = jest.fn(async () => {
+        throw new Error("Error")
+      })
+      const fn_2 = jest.fn(async () => {})
+
+      await service.execute("lock_key", fn_1).catch(() => {})
+      await service.execute("lock_key", fn_2).catch(() => {})
+
+      expect(fn_1).toBeCalledTimes(1)
+      expect(fn_2).toBeCalledTimes(1)
+    })
+  },
+})

packages/modules/locking/jest.config.js

@@ -0,0 +1,10 @@
+const defineJestConfig = require("../../../define_jest_config")
+module.exports = defineJestConfig({
+  moduleNameMapper: {
+    "^@models": "<rootDir>/src/models",
+    "^@services": "<rootDir>/src/services",
+    "^@repositories": "<rootDir>/src/repositories",
+    "^@types": "<rootDir>/src/types",
+    "^@utils": "<rootDir>/src/utils",
+  },
+})

packages/modules/locking/mikro-orm.config.dev.ts

@@ -0,0 +1,6 @@
+import { defineMikroOrmCliConfig, Modules } from "@medusajs/framework/utils"
+import * as entities from "./src/models"
+
+export default defineMikroOrmCliConfig(Modules.LOCKING, {
+  entities: Object.values(entities),
+})

packages/modules/locking/package.json

@@ -0,0 +1,50 @@
+{
+  "name": "@medusajs/locking",
+  "version": "0.0.1",
+  "description": "Locking Module for Medusa",
+  "main": "dist/index.js",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/medusajs/medusa",
+    "directory": "packages/locking"
+  },
+  "publishConfig": {
+    "access": "public"
+  },
+  "author": "Medusa",
+  "license": "MIT",
+  "scripts": {
+    "watch": "tsc --build --watch",
+    "watch:test": "tsc --build tsconfig.spec.json --watch",
+    "resolve:aliases": "tsc --showConfig -p tsconfig.json > tsconfig.resolved.json && tsc-alias -p tsconfig.resolved.json && rimraf tsconfig.resolved.json",
+    "build": "rimraf dist && tsc --build && npm run resolve:aliases",
+    "test": "jest --passWithNoTests --runInBand --bail --forceExit -- src/",
+    "test:integration": "jest --runInBand --forceExit -- integration-tests/**/__tests__/**/*.spec.ts",
+    "migration:generate": " MIKRO_ORM_CLI=./mikro-orm.config.dev.ts medusa-mikro-orm migration:generate",
+    "migration:initial": " MIKRO_ORM_CLI=./mikro-orm.config.dev.ts medusa-mikro-orm migration:create --initial -n InitialSetupMigration",
+    "migration:create": " MIKRO_ORM_CLI=./mikro-orm.config.dev.ts medusa-mikro-orm migration:create",
+    "migration:up": " MIKRO_ORM_CLI=./mikro-orm.config.dev.ts medusa-mikro-orm migration:up",
+    "orm:cache:clear": " MIKRO_ORM_CLI=./mikro-orm.config.dev.ts medusa-mikro-orm cache:clear"
+  },
+  "devDependencies": {
+    "@medusajs/framework": "^0.0.1",
+    "@mikro-orm/cli": "5.9.7",
+    "@mikro-orm/core": "5.9.7",
+    "@mikro-orm/migrations": "5.9.7",
+    "@mikro-orm/postgresql": "5.9.7",
+    "@swc/core": "^1.7.28",
+    "@swc/jest": "^0.2.36",
+    "jest": "^29.7.0",
+    "medusa-test-utils": "^1.1.44",
+    "rimraf": "^3.0.2",
+    "tsc-alias": "^1.8.6",
+    "typescript": "^5.6.2"
+  },
+  "peerDependencies": {
+    "@medusajs/framework": "^0.0.1",
+    "@mikro-orm/core": "5.9.7",
+    "@mikro-orm/migrations": "5.9.7",
+    "@mikro-orm/postgresql": "5.9.7",
+    "awilix": "^8.0.1"
+  }
+}

packages/modules/locking/src/index.ts

@@ -0,0 +1,11 @@
+import { Module, Modules } from "@medusajs/framework/utils"
+import { LockingModuleService } from "@services"
+import loadProviders from "./loaders/providers"
+
+export default Module(Modules.LOCKING, {
+  service: LockingModuleService,
+  loaders: [loadProviders],
+})
+
+// Module options types
+export { LockingModuleOptions } from "./types"

packages/modules/locking/src/loaders/providers.ts

@@ -0,0 +1,85 @@
+import { moduleProviderLoader } from "@medusajs/framework/modules-sdk"
+import {
+  LoaderOptions,
+  ModuleProvider,
+  ModulesSdkTypes,
+} from "@medusajs/framework/types"
+import { ContainerRegistrationKeys } from "@medusajs/framework/utils"
+import { LockingProviderService } from "@services"
+import {
+  LockingDefaultProvider,
+  LockingIdentifiersRegistrationName,
+  LockingProviderRegistrationPrefix,
+} from "@types"
+import { Lifetime, asFunction, asValue } from "awilix"
+import { InMemoryLockingProvider } from "../providers/in-memory"
+
+const registrationFn = async (klass, container, pluginOptions) => {
+  const key = LockingProviderService.getRegistrationIdentifier(klass)
+
+  container.register({
+    [LockingProviderRegistrationPrefix + key]: asFunction(
+      (cradle) => new klass(cradle, pluginOptions.options),
+      {
+        lifetime: klass.LIFE_TIME || Lifetime.SINGLETON,
+      }
+    ),
+  })
+
+  container.registerAdd(LockingIdentifiersRegistrationName, asValue(key))
+}
+
+export default async ({
+  container,
+  options,
+}: LoaderOptions<
+  (
+    | ModulesSdkTypes.ModuleServiceInitializeOptions
+    | ModulesSdkTypes.ModuleServiceInitializeCustomDataLayerOptions
+  ) & { providers: ModuleProvider[] }
+>): Promise<void> => {
+  const logger = container.resolve(ContainerRegistrationKeys.LOGGER)
+  container.registerAdd(LockingIdentifiersRegistrationName, asValue(undefined))
+
+  // InMemoryLockingProvider - default provider
+  container.register({
+    [LockingProviderRegistrationPrefix + InMemoryLockingProvider.identifier]:
+      asFunction((cradle) => new InMemoryLockingProvider(), {
+        lifetime: Lifetime.SINGLETON,
+      }),
+  })
+  container.registerAdd(
+    LockingIdentifiersRegistrationName,
+    asValue(InMemoryLockingProvider.identifier)
+  )
+  container.register(
+    LockingDefaultProvider,
+    asValue(InMemoryLockingProvider.identifier)
+  )
+
+  // Load other providers
+  await moduleProviderLoader({
+    container,
+    providers: options?.providers || [],
+    registerServiceFn: registrationFn,
+  })
+
+  const isSingleProvider = options?.providers?.length === 1
+  let hasDefaultProvider = false
+  for (const provider of options?.providers || []) {
+    if (provider.is_default || isSingleProvider) {
+      if (provider.is_default) {
+        hasDefaultProvider = true
+      }
+      container.register(LockingDefaultProvider, asValue(provider.id))
+    }
+  }
+
+  if (!hasDefaultProvider) {
+    logger.warn(
+      `No default locking provider explicit defined. Using "${container.resolve(
+        LockingDefaultProvider
+      )}" as default.`
+    )
+  }
+}

packages/modules/locking/src/providers/in-memory.ts

@@ -0,0 +1,176 @@
+import { ILockingProvider } from "@medusajs/framework/types"
+import { isDefined } from "@medusajs/framework/utils"
+
+type LockInfo = {
+  ownerId: string | null
+  expiration: number | null
+  currentPromise?: ResolvablePromise
+}
+
+type ResolvablePromise = {
+  promise: Promise<any>
+  resolve: () => void
+}
+
+export class InMemoryLockingProvider implements ILockingProvider {
+  static identifier = "in-memory"
+
+  private locks: Map<string, LockInfo> = new Map()
+
+  constructor() {}
+
+  private getPromise(): ResolvablePromise {
+    let resolve: any
+    const pro = new Promise((ok) => {
+      resolve = ok
+    })
+
+    return {
+      promise: pro,
+      resolve,
+    }
+  }
+
+  async execute<T>(
+    keys: string | string[],
+    job: () => Promise<T>,
+    args?: {
+      timeout?: number
+    }
+  ): Promise<T> {
+    keys = Array.isArray(keys) ? keys : [keys]
+
+    const timeoutSeconds = args?.timeout ?? 5
+
+    const promises: Promise<any>[] = []
+    if (timeoutSeconds > 0) {
+      promises.push(this.getTimeout(timeoutSeconds))
+    }
+
+    promises.push(
+      this.acquire(keys, {
+        awaitQueue: true,
+      })
+    )
+
+    await Promise.race(promises).catch(async (err) => {
+      await this.release(keys)
+    })
+
+    try {
+      return await job()
+    } finally {
+      await this.release(keys)
+    }
+  }
+
+  async acquire(
+    keys: string | string[],
+    args?: {
+      ownerId?: string | null
+      expire?: number
+      awaitQueue?: boolean
+    }
+  ): Promise<void> {
+    keys = Array.isArray(keys) ? keys : [keys]
+    const { ownerId, expire } = args ?? {}
+
+    for (const key of keys) {
+      const lock = this.locks.get(key)
+      const now = Date.now()
+
+      if (!lock) {
+        this.locks.set(key, {
+          ownerId: ownerId ?? null,
+          expiration: expire ? now + expire * 1000 : null,
+          currentPromise: this.getPromise(),
+        })
+
+        continue
+      }
+
+      if (lock.expiration && lock.expiration <= now) {
+        lock.currentPromise?.resolve?.()
+        this.locks.set(key, {
+          ownerId: ownerId ?? null,
+          expiration: expire ? now + expire * 1000 : null,
+          currentPromise: this.getPromise(),
+        })
+
+        continue
+      }
+
+      if (lock.ownerId === ownerId) {
+        if (expire) {
+          lock.expiration = now + expire * 1000
+          this.locks.set(key, lock)
+        }
+
+        continue
+      }
+
+      if (lock.currentPromise && args?.awaitQueue) {
+        await lock.currentPromise.promise
+        return this.acquire(keys, args)
+      }
+
+      throw new Error(`"${key}" is already locked.`)
+    }
+  }
+
+  async release(
+    keys: string | string[],
+    args?: {
+      ownerId?: string | null
+    }
+  ): Promise<boolean> {
+    const { ownerId } = args ?? {}
+    keys = Array.isArray(keys) ? keys : [keys]
+
+    let success = true
+
+    for (const key of keys) {
+      const lock = this.locks.get(key)
+      if (!lock) {
+        success = false
+        continue
+      }
+
+      if (isDefined(ownerId) && lock.ownerId !== ownerId) {
+        success = false
+        continue
+      }
+
+      lock.currentPromise?.resolve?.()
+      this.locks.delete(key)
+    }
+
+    return success
+  }
+
+  async releaseAll(args?: { ownerId?: string | null }): Promise<void> {
+    const { ownerId } = args ?? {}
+
+    if (!isDefined(ownerId)) {
+      for (const [key, lock] of this.locks.entries()) {
+        lock.currentPromise?.resolve?.()
+        this.locks.delete(key)
+      }
+    } else {
+      for (const [key, lock] of this.locks.entries()) {
+        if (lock.ownerId === ownerId) {
+          lock.currentPromise?.resolve?.()
+          this.locks.delete(key)
+        }
+      }
+    }
+  }
+
+  private async getTimeout(seconds: number): Promise<void> {
+    return new Promise((_, reject) => {
+      setTimeout(() => {
+        reject(new Error("Timed-out acquiring lock."))
+      }, seconds * 1000)
+    })
+  }
+}

packages/modules/locking/src/services/index.ts

@@ -0,0 +1,2 @@
+export { default as LockingModuleService } from "./locking-module"
+export { default as LockingProviderService } from "./locking-provider"

packages/modules/locking/src/services/locking-module.ts

@@ -0,0 +1,90 @@
+import {
+  Context,
+  ILockingModule,
+  InternalModuleDeclaration,
+} from "@medusajs/types"
+import { EntityManager } from "@mikro-orm/core"
+import { LockingDefaultProvider } from "@types"
+import LockingProviderService from "./locking-provider"
+
+type InjectedDependencies = {
+  manager: EntityManager
+  lockingProviderService: LockingProviderService
+  [LockingDefaultProvider]: string
+}
+
+export default class LockingModuleService implements ILockingModule {
+  protected manager: EntityManager
+  protected providerService_: LockingProviderService
+  protected defaultProviderId: string
+
+  constructor(
+    container: InjectedDependencies,
+    protected readonly moduleDeclaration: InternalModuleDeclaration
+  ) {
+    this.manager = container.manager
+    this.providerService_ = container.lockingProviderService
+    this.defaultProviderId = container[LockingDefaultProvider]
+  }
+
+  async execute<T>(
+    keys: string | string[],
+    job: () => Promise<T>,
+    args?: {
+      timeout?: number
+      provider?: string
+    },
+    sharedContext: Context = {}
+  ): Promise<T> {
+    const providerId = args?.provider ?? this.defaultProviderId
+    const provider =
+      this.providerService_.retrieveProviderRegistration(providerId)
+
+    return provider.execute(keys, job, args, sharedContext)
+  }
+
+  async acquire(
+    keys: string | string[],
+    args?: {
+      ownerId?: string | null
+      expire?: number
+      provider?: string
+    },
+    sharedContext: Context = {}
+  ): Promise<void> {
+    const providerId = args?.provider ?? this.defaultProviderId
+    const provider =
+      this.providerService_.retrieveProviderRegistration(providerId)
+
+    await provider.acquire(keys, args, sharedContext)
+  }
+
+  async release(
+    keys: string | string[],
+    args?: {
+      ownerId?: string | null
+      provider?: string
+    },
+    sharedContext: Context = {}
+  ): Promise<boolean> {
+    const providerId = args?.provider ?? this.defaultProviderId
+    const provider =
+      this.providerService_.retrieveProviderRegistration(providerId)
+
+    return await provider.release(keys, args, sharedContext)
+  }
+
+  async releaseAll(
+    args?: {
+      ownerId?: string | null
+      provider?: string
+    },
+    sharedContext: Context = {}
+  ): Promise<void> {
+    const providerId = args?.provider ?? this.defaultProviderId
+    const provider =
+      this.providerService_.retrieveProviderRegistration(providerId)
+
+    return await provider.releaseAll(args, sharedContext)
+  }
+}

packages/modules/locking/src/services/locking-provider.ts

@@ -0,0 +1,40 @@
+import { Constructor, ILockingProvider } from "@medusajs/framework/types"
+import { MedusaError } from "@medusajs/framework/utils"
+import { LockingProviderRegistrationPrefix } from "../types"
+
+type InjectedDependencies = {
+  [key: `lp_${string}`]: ILockingProvider
+}
+
+export default class LockingProviderService {
+  protected __container__: InjectedDependencies
+
+  constructor(container: InjectedDependencies) {
+    this.__container__ = container
+  }
+
+  static getRegistrationIdentifier(
+    providerClass: Constructor<ILockingProvider>
+  ) {
+    if (!(providerClass as any).identifier) {
+      throw new MedusaError(
+        MedusaError.Types.INVALID_ARGUMENT,
+        `Trying to register a locking provider without an identifier.`
+      )
+    }
+    return `${(providerClass as any).identifier}`
+  }
+
+  public retrieveProviderRegistration(providerId: string): ILockingProvider {
+    try {
+      return this.__container__[
+        `${LockingProviderRegistrationPrefix}${providerId}`
+      ]
+    } catch (err) {
+      throw new MedusaError(
+        MedusaError.Types.NOT_FOUND,
+        `Could not find a locking provider with id: ${providerId}`
+      )
+    }
+  }
+}

packages/modules/locking/src/types/index.ts

@@ -0,0 +1,33 @@
+import {
+  ModuleProviderExports,
+  ModuleServiceInitializeOptions,
+} from "@medusajs/framework/types"
+
+export const LockingDefaultProvider = "default_provider"
+export const LockingIdentifiersRegistrationName = "locking_providers_identifier"
+
+export const LockingProviderRegistrationPrefix = "lp_"
+
+export type LockingModuleOptions = Partial<ModuleServiceInitializeOptions> & {
+  /**
+   * Providers to be registered
+   */
+  providers?: {
+    /**
+     * The module provider to be registered
+     */
+    resolve: string | ModuleProviderExports
+    /**
+     * If the provider is the default
+     */
+    is_default?: boolean
+    /**
+     * The id of the provider
+     */
+    id: string
+    /**
+     * key value pair of the configuration to be passed to the provider constructor
+     */
+    options?: Record<string, unknown>
+  }[]
+}

packages/modules/locking/tsconfig.json

@@ -0,0 +1,12 @@
+{
+  "extends": "../../../_tsconfig.base.json",
+  "compilerOptions": {
+    "paths": {
+      "@models": ["./src/models"],
+      "@services": ["./src/services"],
+      "@repositories": ["./src/repositories"],
+      "@types": ["./src/types"],
+      "@utils": ["./src/utils"]
+    }
+  }
+}