forge/modules/globallock/kv_locker.go

// Copyright 2026 Hanzo AI, Inc. All rights reserved.
// Copyright 2024 The Gitea Authors. All rights reserved.
// SPDX-License-Identifier: MIT

package globallock

import (
	"context"
	"errors"
	"sync"
	"sync/atomic"
	"time"

	"github.com/hanzoai/git/modules/nosql"

	"github.com/hanzokv/go/v9"
	"github.com/hanzokv/lock/v4"
	"github.com/hanzokv/lock/v4/kv/kvgo"
)

const kvLockKeyPrefix = "gitea:globallock:"

// kvLockExpiry is the default expiry time for a lock.
// Define it as a variable to make it possible to change it in tests.
var kvLockExpiry = 30 * time.Second

type kvLocker struct {
	conn kv.UniversalClient
	rs   *redsync.Redsync

	mutexM   sync.Map
	closed   atomic.Bool
	extendWg sync.WaitGroup
}

var _ Locker = &kvLocker{}

func NewKVLocker(connection string) Locker {
	conn := nosql.GetManager().GetKVClient(connection)
	l := &kvLocker{
		conn: conn,
		rs:   redsync.New(kvgo.NewPool(conn)),
	}
	l.extendWg.Add(1)
	l.startExtend()
	return l
}

func (l *kvLocker) Lock(ctx context.Context, key string) (ReleaseFunc, error) {
	return l.lock(ctx, key, 0)
}

func (l *kvLocker) TryLock(ctx context.Context, key string) (bool, ReleaseFunc, error) {
	f, err := l.lock(ctx, key, 1)

	var (
		errTaken     *redsync.ErrTaken
		errNodeTaken *redsync.ErrNodeTaken
	)
	if errors.As(err, &errTaken) || errors.As(err, &errNodeTaken) {
		return false, f, nil
	}
	return err == nil, f, err
}

// Close closes the locker.
// It will stop extending the locks and refuse to acquire new locks.
// In actual use, it is not necessary to call this function.
// But it's useful in tests to release resources.
// It could take some time since it waits for the extending goroutine to finish.
func (l *kvLocker) Close() error {
	l.closed.Store(true)
	l.extendWg.Wait()
	return nil
}

func (l *kvLocker) lock(ctx context.Context, key string, tries int) (ReleaseFunc, error) {
	if l.closed.Load() {
		return func() {}, errors.New("locker is closed")
	}

	options := []redsync.Option{
		redsync.WithExpiry(kvLockExpiry),
	}
	if tries > 0 {
		options = append(options, redsync.WithTries(tries))
	}
	mutex := l.rs.NewMutex(kvLockKeyPrefix+key, options...)
	if err := mutex.LockContext(ctx); err != nil {
		return func() {}, err
	}

	l.mutexM.Store(key, mutex)

	releaseOnce := sync.Once{}
	return func() {
		releaseOnce.Do(func() {
			l.mutexM.Delete(key)

			// It's safe to ignore the error here,
			// if it failed to unlock, it will be released automatically after the lock expires.
			// Do not call mutex.UnlockContext(ctx) here, or it will fail to release when ctx has timed out.
			_, _ = mutex.Unlock()
		})
	}, nil
}

func (l *kvLocker) startExtend() {
	if l.closed.Load() {
		l.extendWg.Done()
		return
	}

	toExtend := make([]*redsync.Mutex, 0)
	l.mutexM.Range(func(_, value any) bool {
		m := value.(*redsync.Mutex)

		// Extend the lock if it is not expired.
		// Although the mutex will be removed from the map before it is released,
		// it still can be expired because of a failed extension.
		// If it happens, it does not need to be extended anymore.
		if time.Now().After(m.Until()) {
			return true
		}

		toExtend = append(toExtend, m)
		return true
	})
	for _, v := range toExtend {
		// If it failed to extend, it will be released automatically after the lock expires.
		_, _ = v.Extend()
	}

	time.AfterFunc(kvLockExpiry/2, l.startExtend)
}