基于cats-effect的Worker如何避免竞态条件？

核心问题分析 #

1. 锁机制的非原子性竞态 #

你实现的锁逻辑if (!locked.get() && jobsTrieMap.nonEmpty)后直接执行locked.set(true)，这两步并非原子操作。多个worker线程可能同时通过if判断，导致：

• 多线程进入临界区，重复获取同一个jobId
• 锁状态被频繁覆盖，完全失去互斥作用

2. JobId生成的竞态依赖 #

用jobsTrieMap.size生成jobId的方式存在严重问题：

• size()在并发场景下不是原子操作，多线程调用addJob会生成重复jobId，导致任务被覆盖丢失
• worker移除任务后map尺寸减小，后续addJob会复用已使用过的jobId，进一步加剧任务丢失

3. 任务获取与移除的非原子性 #

先通过jobsTrieMap(jobId)获取任务再执行移除的操作，中间存在时间窗口：其他线程可能已移除该jobId，导致apply方法抛出NoSuchElementException，而未捕获的异常会静默终止worker线程，最终导致系统挂起。

4. 阻塞式代码与cats-effect的不兼容 #

worker中使用while(true)+Thread.sleep(100)的阻塞逻辑，违背了cats-effect的非阻塞设计原则。这种写法会持续占用线程池资源，当大量线程被阻塞时，会导致线程池耗尽，系统无法继续处理任务。

修复方案 #

步骤1：原子化JobId生成 #

使用独立的AtomicLong生成唯一jobId，彻底摆脱对map尺寸的依赖：

private val nextJobId: AtomicLong = AtomicLong(0L)

def addJob(jobFn: () => Any): Unit =
  val jobId = nextJobId.getAndIncrement()
  jobsTrieMap.put(jobId, jobFn)

步骤2：原子化锁抢占 #

用AtomicBoolean.compareAndSet实现原子锁获取，确保同一时间只有一个线程进入临界区：

if (locked.compareAndSet(false, true) && jobsTrieMap.nonEmpty) {
  // 临界区逻辑
}

步骤3：原子化任务获取与移除 #

使用TrieMap.remove原子性地完成任务获取与移除，避免中间竞态：

jobsTrieMap.remove(jobId) match {
  case Some(toDo) => // 执行任务
  case None => // 任务已被其他线程处理，跳过
}

步骤4：改用cats-effect非阻塞循环 #

用IO递归循环替代阻塞式while(true)，充分利用cats-effect的线程调度能力：

private def worker: IO[Unit] =
  IO.uncancelable { _ =>
    IO(locked.compareAndSet(false, true)).flatMap { lockedSuccess =>
      if (lockedSuccess) {
        IO {
          val jobId = assignedTillJobId.getAndIncrement()
          jobsTrieMap.remove(jobId)
        }.flatMap {
          case Some(toDo) =>
            IO(toDo()) >> IO(count.incrementAndGet()) >> IO(locked.set(false))
          case None =>
            IO(locked.set(false))
        }
      } else {
        IO.sleep(100.millis)
      }
    }
  }.foreverM

完整修复代码 #

import cats.effect.IO
import cats.effect.std.Dispatcher
import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong}
import scala.collection.concurrent.TrieMap
import scala.concurrent.duration._

object ThreadingError extends App:
  val jobIdsAdded = (0L until 10000L).toList
  val dispatcher = Dispatcher[IO].unsafeRunSync()
  
  jobIdsAdded.foreach(_ => dispatcher.unsafeRunAndForget(IO(ParallelJobs.addJob(() => {}))))
  
  def waitForCompletion: IO[Unit] =
    IO(ParallelJobs.count.get()).flatMap { cnt =>
      if (cnt < 10000L) {
        IO(print(s"${cnt}\r")) >> IO.sleep(200.millis) >> waitForCompletion
      } else {
        IO.println("\nAll jobs completed!")
      }
    }
  
  waitForCompletion.unsafeRunSync()
  dispatcher.unsafeShutdown()

object ParallelJobs:
  private val allCores = Runtime.getRuntime.availableProcessors()
  private val availableCores = allCores - 1
  private val assignedTillJobId: AtomicLong = AtomicLong(0L)
  private val nextJobId: AtomicLong = AtomicLong(0L)
  private val jobsTrieMap: TrieMap[Long, () => Any] = TrieMap.empty[Long, () => Any]

  private val locked: AtomicBoolean = AtomicBoolean(false)
  val count: AtomicLong = AtomicLong(0L)

  private val workerGroup = (0 until availableCores).map(_ => worker).reduce(_ >> _)
  workerGroup.unsafeRunAsync {
    case Left(e) =>
      println(s"Worker error: ${e.getMessage}")
      e.printStackTrace()
    case Right(_) => ()
  }

  def addJob(jobFn: () => Any): Unit =
    val jobId = nextJobId.getAndIncrement()
    jobsTrieMap.put(jobId, jobFn)

  private def worker: IO[Unit] =
    IO.uncancelable { _ =>
      IO(locked.compareAndSet(false, true)).flatMap { lockedSuccess =>
        if (lockedSuccess) {
          IO {
            val jobId = assignedTillJobId.getAndIncrement()
            jobsTrieMap.remove(jobId)
          }.flatMap {
            case Some(toDo) =>
              IO(toDo()) >> IO(count.incrementAndGet()) >> IO(locked.set(false))
            case None =>
              IO(locked.set(false))
          }
        } else {
          IO.sleep(100.millis)
        }
      }
    }.foreverM

额外说明 #

• 用Dispatcher安全地从主线程调用addJob，避免在IO上下文外直接操作共享状态
• 替换Thread.sleep为IO.sleep，符合cats-effect非阻塞模型
• 用IO.uncancelable确保临界区操作不会被取消，避免锁泄漏

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