#node下基于MySQL实现分布式锁

import { EntityModel } from '@midwayjs/orm';
import { BaseEntity, Column, PrimaryColumn } from 'typeorm';

@EntityModel('fem_shedlock')
export class ShedlockEntity extends BaseEntity {
  @PrimaryColumn({ type: 'varchar', length: 64, name: 'name' })
  name: string;

  @Column({ type: 'timestamp', name: 'lock_until' })
  lockUntil: Date;

  @Column({ type: 'timestamp', name: 'locked_at' })
  lockedAt: Date;

  @Column({ type: 'varchar', length: 255, name: 'locked_by' })
  lockedBy: string;
}

import { Logger, Provide } from '@midwayjs/decorator';
import { InjectEntityModel } from '@midwayjs/orm';
import { Repository, LessThan, MoreThanOrEqual } from 'typeorm';
import { v4 as uuidv4 } from 'uuid';
import { ILogger } from '@midwayjs/logger';
import * as dayjs from 'dayjs';
import type { ManipulateType } from 'dayjs';

import { ShedlockEntity } from '@/entity/shedlock.entity';

/**
 * 生成锁创建者编号，由环境变量中的容器名称、容器IP和一个随机的uuid字符串组成。
 */
const getLockCreator = () => {
  const name = process.env.POD_NAME || 'noname';
  const ip = process.env.POD_IP || 'noip';
  const suffix = uuidv4();
  return `${name}_${ip}_${suffix}`;
};

// 锁创建者编号
const LOCK_CREATOR = getLockCreator();

@Provide()
export default class ShedlockService {
  @Logger()
  logger: ILogger;

  @InjectEntityModel(ShedlockEntity)
  shedlockRepository: Repository<ShedlockEntity>;

  /**
   * 查看锁是否已经存在，不管锁有没有过期
   * @param lockName 琐名称
   * @private
   */
  private async find(lockName: string) {
    return await this.shedlockRepository.findOne({
      where: {
        name: lockName,
      },
    });
  }

  /**
   * 插入锁，返回是否插入成功
   * @param lockName 琐名称
   * @param lockTime 锁定时间
   * @param lockTimeUnit 锁定时间的时间单位
   * @private
   */
  private async insert(
    lockName: string,
    lockTime: number,
    lockTimeUnit: ManipulateType
  ): Promise<boolean> {
    const current = dayjs();
    try {
      await this.shedlockRepository.insert({
        name: lockName,
        lockUntil: current.add(lockTime, lockTimeUnit).toDate(),
        lockedAt: current.toDate(),
        lockedBy: LOCK_CREATOR,
      });
    } catch (e) {
      this.logger.info(
        `[Shedlock]Insert lock failed. lockName=${lockName} creator=${LOCK_CREATOR}`,
        e
      );
      return false;
    }
    return true;
  }

  /**
   * 尝试更新锁
   * @param lockName 琐名称
   * @param lockTime 锁定时间
   * @param lockTimeUnit 锁定时间单位
   * @private
   */
  private async update(
    lockName: string,
    lockTime: number,
    lockTimeUnit: ManipulateType
  ): Promise<boolean> {
    const current = dayjs();
    try {
      const updateEffect = await this.shedlockRepository.update(
        {
          name: lockName,
          lockUntil: LessThan(current.toDate()),
        },
        {
          lockUntil: current.add(lockTime, lockTimeUnit).toDate(),
          lockedAt: current.toDate(),
          lockedBy: LOCK_CREATOR,
        }
      );
      return updateEffect.affected > 0;
    } catch (e) {
      this.logger.info(
        `[Shedlock]Update lock failed. lockName=${lockName} creator=${LOCK_CREATOR}`,
        e
      );
      return false;
    }
  }

  /**
   * 更新锁的持续时间，只能更新自己的锁
   * @param lockName 琐名称
   * @param untilTime 新的持续时间
   * @private
   */
  private async updateUntil(lockName: string, untilTime: Date) {
    const updateEffect = await this.shedlockRepository.update(
      {
        name: lockName,
        lockedBy: LOCK_CREATOR,
        lockUntil: MoreThanOrEqual(dayjs().toDate()),
      },
      {
        lockUntil: untilTime,
      }
    );
    return updateEffect.affected > 0;
  }

  /**
   * 尝试获取锁
   * @param lockName 琐名称
   * @param lockTime 锁时间
   * @param lockTimeUnit 锁时间单位
   */
  async lock(
    lockName: string,
    lockTime: number,
    lockTimeUnit: ManipulateType
  ): Promise<boolean> {
    if (!(await this.find(lockName))) {
      const insertSuccess = await this.insert(lockName, lockTime, lockTimeUnit);
      // 插入失败，尝试执行后续的更新锁逻辑
      if (insertSuccess) {
        return true;
      }
    }
    return await this.update(lockName, lockTime, lockTimeUnit);
  }

  /**
   * 锁续期，更新锁的持续时间，只有拥有锁的时候才可以延长持续时间
   * @param lockName 琐名称
   * @param lockTime 持续时间
   * @param lockTimeUnit 持续时间单位
   */
  async extend(
    lockName: string,
    lockTime: number,
    lockTimeUnit: ManipulateType
  ): Promise<boolean> {
    if (!(await this.find(lockName))) {
      return false;
    }
    const current = dayjs();
    return await this.updateUntil(
      lockName,
      current.add(lockTime, lockTimeUnit).toDate()
    );
  }

  /**
   * 释放锁
   * @param lockName 琐名称
   */
  async unlock(lockName: string): Promise<boolean> {
    return await this.updateUntil(lockName, dayjs().toDate());
  }
}

// 获取LOCK_NAME的锁，锁持有10分钟
await this.shedlockService.lock(LOCK_NAME, 10, 'minutes')

// 释放锁
await this.shedlockService.unlock(LOCK_NAME);