好问题，逐一分析：

1. closeCallbacksCollected 命名

这个名字的歧义在于："collected" 既可以理解为"已收集到回调"，也可以理解为"已完成收集动作"。实际语义是后者——"close 流程已完成回调的收集与清空"。

更好的选项：

// 推荐：强调 close 流程已完成，回调将不再被缓冲
private var closeFinalized = false

// 或者，强调"任何新注册的回调应立即调用"这个行为意图
private var invokeOnCloseImmediately = false

// 或者，从状态机角度命名（close 已进入最终阶段）
private var closeCallbacksDrained = false  // "已排干"，比 collected 更精确

我倾向于 closeCallbacksDrained，因为 drain（排干/清空并处理）在并发编程里是个惯用词，语义非常精确。

2. 用 ConcurrentHashMap 替换 states

直接用 ConcurrentHashMap 不够，原因如下：

// 这三个方法存在复合操作，ConcurrentHashMap 无法保证原子性
fun hasState(key: String): Boolean  // containsKey
fun getState(key: String): T        // containsKey + get  ← 两步！
fun getStateOrNull(key: String): T? // containsKey + get  ← 两步！

getState 里的 containsKey + get 是两个独立操作，中间可能被其他线程的 setState 打断，产生 TOCTOU 问题。所以保留锁是必要的，但可以把锁收敛得更干净。

3. states.toMutableMap() → ConcurrentHashMap

// 普通 MutableMap 转 ConcurrentHashMap
val connection = WsConnection(
    channel, pathParams, queryParams, app,
    ConcurrentHashMap(states),  // 直接用拷贝构造函数
    headers
)

但如上所述，内部仍需显式锁，所以这个改动意义不大。不如在 WsConnection 内部统一用 HashMap + statesLock，对外不暴露实现细节。

4. 其他 BUG 与可优化点

① closeReason 的 AtomicReference 用法有隐患

// 当前代码
closeReason.compareAndSet(WsCloseReason.Normal, reason)

如果调用方传入的 reason 本身就是 Normal，这个 CAS 不会生效——但这是合法调用，不是 bug。真正的问题是：closeReason 同时被 close() 写、被 onClose() 读，而读操作在锁外：

// onClose() 里，此时已在锁内，closeReason 已经被设置，这里是安全的
runCatching { callback.accept(closeReason.get()) }

// close() 里，锁外读取
val finalReason = closeReason.get()  // ← 在 synchronized 块外面

closeReason 是 AtomicReference 所以读本身是安全的，但可以简化为普通 @Volatile 变量，语义更清晰：

@Volatile private var closeReason: WsCloseReason = WsCloseReason.Normal
// 在 synchronized 块内直接赋值即可，锁已经保证可见性

② dispatchMessage / dispatchError 在连接关闭后仍可能执行回调

internal fun dispatchMessage(message: WsMessage) {
    if (isClosed) return        // 检查
    val callbacks = ...
    synchronized(messageCallbacks) { ... }
    callbacks.forEach { ... }   // ← 执行时连接可能已关闭
}

isClosed 检查和回调执行之间存在窗口。通常问题不大（回调本身应该是幂等的），但如果追求严格，可以在 close() 里给回调列表设一个"已关闭"哨兵，或者在回调执行前二次检查。

③ messageCallbacks / errorCallbacks 使用 ArrayList + synchronized，可以统一风格

当前 statesLock 用 Any() 对象，其他几个回调列表直接 synchronized(list) 用列表自身作锁——风格不统一，容易混淆。

综合重构建议

// ① 重命名
private var closeCallbacksDrained = false

// ② closeReason 简化为 @Volatile
@Volatile private var closeReason: WsCloseReason = WsCloseReason.Normal

// ③ 统一锁对象风格（或全部用伴生对象风格的私有锁）
private val messageCallbacksLock = Any()
private val closeCallbacksLock = Any()
private val errorCallbacksLock = Any()
// statesLock 已有，保持一致

// ④ close() 里消除锁外读取
synchronized(closeCallbacksLock) {
    closeReason = reason           // 直接赋值，@Volatile + 锁双重保证
    callbacks = ArrayList(closeCallbacks)
    closeCallbacks.clear()
    closeCallbacksDrained = true
}
// 不再需要锁外的 closeReason.get()
callbacks.forEach { cb -> runCatching { cb.accept(closeReason) } }

总结一下优先级：