ASP.NET Core中WebSocket管理类Add方法锁流程优化咨询

原代码锁逻辑的问题分析 #

原Add方法的锁设计存在以下性能与安全性隐患：

• 全局锁粒度过粗：使用全局ReaderWriterLockSlim写锁覆盖整个添加流程，即使是不同key的WebSocket添加操作也会互相阻塞，严重限制并发能力。
• 双重锁冗余嵌套：在全局写锁内部调用CustomServerWebSockets.Add时，又会获取子对象的写锁，增加锁开销的同时，存在潜在的嵌套锁风险。
• 字典操作冗余：先调用ContainsKey再调用TryAdd，两次字典查找，且依赖全局锁避免竞态条件，效率低下。

优化思路 #

1. 替换为线程安全字典：用ConcurrentDictionary<TKey, CustomServerWebSockets>替代普通Dictionary，利用其原子操作GetOrAdd安全完成"检查存在-创建实例"的逻辑，无需全局锁保护字典。
2. 缩小锁范围：仅在操作子对象内部列表时，由子对象自身的锁负责线程安全，全局层面不再持有写锁，不同key的添加操作可完全并行。
3. 原子化字典操作：通过GetOrAdd原子完成对象的获取与创建，避免冗余查找和竞态条件。

优化后的代码实现 #

WebSocketsOrganizer类调整 #

public sealed class WebSocketsOrganizer<TKey> : IDisposable where TKey : notnull
{
    // 替换为线程安全的ConcurrentDictionary，移除全局ReaderWriterLockSlim
    private readonly ConcurrentDictionary<TKey, CustomServerWebSockets> _customWebSockets = new();
    private bool _disposed;

    public void Add(TKey key, CustomServerWebSocket customWebSocket)
    {
        // 原子性获取或创建WebSocket分组，无需全局锁
        var webSocketGroup = _customWebSockets.GetOrAdd(key, _ => new CustomServerWebSockets());
        // 交由子对象自身的锁处理内部列表添加
        webSocketGroup.Add(customWebSocket);
    }

    // 同步优化Remove方法（参考）
    public void Remove(TKey key, CustomServerWebSocket customWebSocket)
    {
        if (!_customWebSockets.TryGetValue(key, out var webSocketGroup))
        {
            return;
        }

        webSocketGroup.Remove(customWebSocket);

        // 分组为空时尝试原子移除，避免竞态
        if (!webSocketGroup.Any())
        {
            _customWebSockets.TryRemove(key, out _);
        }
    }

    public void Dispose()
    {
        if (_disposed)
        {
            return;
        }

        _disposed = true;
        // 释放所有子分组资源
        foreach (var group in _customWebSockets.Values)
        {
            group.Dispose();
        }
    }

    //...
}

CustomServerWebSockets类可选优化 #

如果无需严格的列表顺序或索引操作，可替换为ConcurrentBag简化锁逻辑；若需精确移除操作，保留原锁方案即可：

public sealed class CustomServerWebSockets : IDisposable
{
    // 无需手动锁的线程安全容器，适合只增删不依赖索引的场景
    private readonly ConcurrentBag<CustomServerWebSocket> _customWebSockets = new();
    private bool _disposed;

    public void Add(CustomServerWebSocket customWebSocket)
    {
        _customWebSockets.Add(customWebSocket);
    }

    public void Remove(CustomServerWebSocket socket)
    {
        // 若需精确移除，保留原List+ReaderWriterLockSlim方案即可
        // _customWebSocketsLock.EnterWriteLock();
        // try
        // {
        //     _customWebSockets.Remove(socket);
        // }
        // finally
        // {
        //     _customWebSocketsLock.ExitWriteLock();
        // }
    }

    public bool Any()
    {
        return _customWebSockets.Any();
    }

    public void Dispose()
    {
        if (_disposed)
        {
            return;
        }

        _disposed = true;
        // 若使用了锁，在此释放锁资源
        // _customWebSocketsLock.Dispose();
    }

    //...
}

优化效果 #

• 并发能力提升：不同key的添加操作完全并行，同一key的操作仅在子分组层面排队，锁粒度极大缩小。
• 锁开销降低：移除全局写锁，避免双重锁嵌套，减少线程阻塞概率。
• 代码更安全简洁：依赖ConcurrentDictionary的原子操作，无需手动处理字典的线程安全问题，降低出错风险。

内容的提问来源于stack exchange，提问作者Mehdi