VEU.WathBody/3-app/1-lib/Kit.Helpers/Service/DistributedPg/PgTableDistributedLockService.cs

using Kit.Helpers.Postgres;

namespace Kit.Helpers;

/// <summary>
/// Провайдер блокировок через самописную схему (и её элементы) distributed. <br />
/// Все активные блокировки хрянятся в таблице в явном виде. <br />
/// У активной блокировки заполянется время начала (acquired_at_utc timestamptz) и обновляется раз в <see cref="_refreshIntervalMs"/>. <br />
/// Если приложение упадёт и перестанет обновлять время начала, через (timeout_ms int default) миллисекунд блокировка будет считаться недействительной. <br />
/// При следующей операции, связанной с любой блокировкой, запись о недействительной блокировке будет удалена.
/// </summary>
public class PgTableDistributedLockService : IDistributedLockService
{
    private readonly string _connectionString;

    private const int _sleepIntervalMs = 10;
    private const int _refreshIntervalMs = 100;
    private readonly object _ownersLock = new object();
    private readonly IDictionary<int, LockOwner> _owners = new Dictionary<int, LockOwner>();

    public PgTableDistributedLockService(IConfigurationDistributedPg configurationPgLock)
    {
        _connectionString = configurationPgLock.DistributedConnectionString;
    }

    private class LockOwner
    {
        public string Id { get; set; }
        public int ThreadId { get; set; }
        /// <summary>
        /// Уровень вложенности блокировок в рамках одного потока и объекта. 
        /// 0 - активных блокировок нет.
        /// 1 - одна блокировка.
        /// >1 - есть вложенные блокировки.
        /// </summary>
        public IDictionary<string, int> LockedObjects { get; set; }
        public LockOwner(int threadId)
        {
            ThreadId = threadId;
            Id = Guid.NewGuid().ToString();
            LockedObjects = new Dictionary<string, int>();
        }
    }

    private IList<KeyValuePair<string, object>> GetParameters(string type, string key, string ownerId)
    {
        return new List<KeyValuePair<string, object>>
        {
            new KeyValuePair<string, object>("_type", type),
            new KeyValuePair<string, object>("_key", key),
            new KeyValuePair<string, object>("_owner_id", ownerId),
        };
    }

    private bool TryAcquireLock(string type, string key, string ownerId)
    {        
        return _connectionString.ExecuteScalar<bool>("distributed.try_acquire_lock", 
            isStoredProcedure: true, this.GetParameters(type, key, ownerId));
    }

    private void Tick(string type, string key, string ownerId)
    {
        _connectionString.ExecuteNonQuery("distributed.lock_tick_for_owner_id",
            isStoredProcedure: true, this.GetParameters(type, key, ownerId));
    }

    private void Free(string type, string key, string ownerId)
    {
        _connectionString.ExecuteNonQuery("distributed.lock_free_for_owner_id",
            isStoredProcedure: true, this.GetParameters(type, key, ownerId));
    }

    private int Count(string type, string? key)
    {
        string sql;
        var parameters = new List<KeyValuePair<string, object>>
        {
            new KeyValuePair<string, object>("_type", type)
        };
        if (key == null)
        {
            sql = "distributed.lock_count_by_type";
        }
        else
        {
            sql = "distributed.lock_count_by_type_key";
            parameters.Add(new KeyValuePair<string, object>("_key", key));
        }
        return _connectionString.ExecuteScalar<int>(sql, isStoredProcedure: true, parameters);
    }

    private LockOwner GetLockOwner(string lockedObjectKey)
    {
        lock (_ownersLock)
        {
            LockOwner lockOwner;
            int threadId = Environment.CurrentManagedThreadId;
            if (_owners.ContainsKey(threadId))
            {
                lockOwner = _owners[threadId];
            }
            else
            {
                lockOwner = new LockOwner(threadId);
                _owners.Add(threadId, lockOwner);
            }

            if (lockOwner.LockedObjects.ContainsKey(lockedObjectKey))
            {
                lockOwner.LockedObjects[lockedObjectKey]++;
            }
            else
            {
                lockOwner.LockedObjects.Add(lockedObjectKey, 1);
            }

            return lockOwner;
        }
    }

    public void ExecuteOnLock(string type, string key, Action action) => ExecuteOnLock(type, key, () =>
    {
        action();
        return true;
    });

    public TResult ExecuteOnLock<TResult>(string type, string key, Func<TResult> action)
    {
        string lockedObjectKey = $"{type};;;{key}";
        // Формируем Id владельца блокировки
        LockOwner owner = this.GetLockOwner(lockedObjectKey);
        try
        {
            // Ждём момента, когда наш запрос займёт блокировку
            while (TryAcquireLock(type, key, owner.Id) == false)
            {
                Thread.Sleep(_sleepIntervalMs);
            }
            // Периодически обновляя дату начала блокировки, выполняем операцию
            using (var timer = new Timer(_ => Tick(type, key, owner.Id), null, 0, _refreshIntervalMs))
            {
                return action();
            }
        }
        finally
        {
            lock (_ownersLock)
            {
                //Если это не вложенная блокировка
                if (owner.LockedObjects[lockedObjectKey] < 2)
                {
                    // Освобождаем блокировку для владельца
                    Free(type, key, owner.Id);
                    owner.LockedObjects.Remove(lockedObjectKey);
                    if (owner.LockedObjects.Count < 1)
                    {
                        _owners.Remove(owner.ThreadId);
                    }
                }
                else
                {
                    owner.LockedObjects[lockedObjectKey]--;
                }
            }
        }
    }

    public int GetLockCount(string type) => Count(type, null);
    public int GetLockCount(string type, string key) => Count(type, key);
}