go + gorm + mysql 唯一索引约束 实现分布式锁

基于数据库唯一索引实现

借助数据库的唯一约束特性来实现互斥访问，其核心在于插入记录时利用唯一索引的冲突检查机制。

package main

import (
	"errors"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"time"
)

// 定义锁结构
type DBLock struct {
	ID uint `gorm:"primaryKey"`
	Resource string `gorm:"uniqueIndex"` // 唯一索引，确保同一资源只有一个锁
	Holder string `gorm:"index"`
	ExpiresAt time.Time `gorm:"index"`
}

// 锁管理器
type LockManager struct {
	db *gorm.DB
	holderID string // 锁持有者标识
}

// 创建新的锁管理器
func NewLockManager(dsn string, holderID string) (*LockManager, error) {
	db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{})
	if err != nil {
		return nil, err
	}
	// 自动迁移表结构
	if err := db.AutoMigrate(&DBLock{}); err != nil {
		return nil, err
	}
	return &LockManager{
		db: db,
		holderID: holderID
	}
}

// 获取锁
func (lm *LockManager) Acquire(resource string, timeout time.Duration) (bool, error) {
	lock := DBLock{
		Resource: resource,
		Holder: lm.holderID,
		ExpiresAt: time.Now().Add(timeout),
	}

	// 尝试插入记录
	result := lm.db.Create(&lock)
	if result.Error != nil {
		// 检查是否是唯一索引冲突
		if errors.Is(result.Error, gorm.ErrEuplicatedKey) {
			return false, nil
		}
		return false, result.Error
	}

	return true, nil
}

// 释放锁
func (lm *LockManager) Release(resource string) error {
	// 删除锁记录，只删除当前持有者的锁
	result := lm.db.Where("resource = ? AND holder = ?", resource, lm.holderID).Delete(&DBLock{})
	if result.Error != nil {
		return result.Error
	}
	if result.RowsAffected == 0 {
		return errors.New("锁不存在或已被释放")
	}
	return nil
}

// 自动续约
func (lm *LockManager) Renew(resource string, duration time.Duration) error {
	tx := lm.db.Begin()
	defer func() {
		if r := recover(); r != nil {
			tx.Rollback()
		}
	}()

	// 查询锁记录
	var lock DBLock
	if err := tx.Where("resource = ? AND holder = ?", resource, lm.holderID).First(&lock).Error; err != nil {
		tx.Rollback()
		return err
	}

	// 更新锁的过期时间
	if err := tx.Model(&lock).Update("expires_at", time.Now().Add(duration)).Error; err != nil {
		tx.Rollback()
		return err
	}

	return tx.Commit().Error
}

// 检查锁状态
func (lm *LockManager) IsLocked(resource string) (bool, error) {
	var count int64
	err := lm.db.Model(&DBLock{}).Where("resource = ? AND expires_at > ?", resource, time.Now()).Count(&count).Error
	if err != nil {
		return false, err
	}
	return count > 0, nil
}