apps\hub\src\main\java\cn\nfarms\hub\task\Tasks.kt¶
源码路径:
apps\hub\src\main\java\cn\nfarms\hub\task\Tasks.kt· 200 行
package cn.nfarms.hub.task
import java.io.File
import java.util.UUID
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.Executors
import java.util.concurrent.atomic.AtomicBoolean
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.serialization.Serializable
/**
* HUB 的任务模型。
*
* 「一个任务」= 一次可取消、可看进度、有终态的作业。HUB 里所有重活
* (下载模型、跑 ffmpeg、拉 dsh)都走同一个调度器 —— 因为它们**必须串行**:
* 同时跑两个 ffmpeg 会互相抢 CPU,同时下两个 17 GB 权重会把磁盘写满。
* 所以调度器是「一条队列 + 单工作线程」,而不是每来一个任务开一个线程。
*/
@Serializable
data class TaskRecord(
val id: String,
val kind: String,
val title: String,
val detail: String = "",
val state: TaskState = TaskState.Queued,
/** 0..1;-1 表示不可计量 */
val progress: Float = -1f,
val createdAt: Long = System.currentTimeMillis(),
val startedAt: Long? = null,
val endedAt: Long? = null,
val exitCode: Int? = null,
/** 失败原因(人话,直接显示在界面上) */
val error: String? = null,
/** 产物路径(下载文件 / ffmpeg 输出) */
val artifact: String? = null,
/** 产物 SHA-256(下载类任务必填,校验通过才落盘) */
val sha256: String? = null,
) {
val finished: Boolean get() = state in setOf(TaskState.Done, TaskState.Failed, TaskState.Cancelled)
}
@Serializable
enum class TaskState { Queued, Running, Done, Failed, Cancelled }
/**
* 一个可执行作业。
*
* `run` 拿到一个 [JobCtl],它负责三件事:报进度、看是否被取消、把产物写回去。
* 为什么把「取消」做成信号而不是直接中断线程:ffmpeg / 下载器都要**自己**
* 决定怎么收场(ffmpeg 要发 `q` 或杀掉子进程;下载要删掉半截文件),
* 硬中断线程只会留下垃圾。
*/
interface HubJob {
val kind: String
val title: String
val detail: String
fun run(ctl: JobCtl)
}
class JobCtl(
val workDir: File,
private val cancelled: AtomicBoolean,
private val onProgress: (Float, String?) -> Unit,
) {
val isCancelled: Boolean get() = cancelled.get()
fun progress(fraction: Float, note: String? = null) {
onProgress(fraction.coerceIn(0f, 1f), note)
}
/** 抛这个异常表示「被用户取消」,调度器会记成 Cancelled 而不是 Failed。 */
fun throwIfCancelled() {
if (cancelled.get()) throw CancelledException()
}
}
class CancelledException : Exception("任务已取消")
/**
* 串行调度器。
*
* 状态全在内存里(界面实时看),落盘只落**历史**(见 [TaskHistory])——
* 「队列」这种东西进程一死就没意义了,不该持久化,否则下次启动会看到
* 一堆永远不会再跑的 Queued。
*/
class TaskScheduler(
private val workDir: File,
private val history: TaskHistory,
) {
private val worker = Executors.newSingleThreadExecutor { r ->
Thread(r, "hub-scheduler").apply { isDaemon = true }
}
private val queue = ArrayDeque<String>()
private val jobs = ConcurrentHashMap<String, HubJob>()
private val cancels = ConcurrentHashMap<String, AtomicBoolean>()
private val _tasks = MutableStateFlow<List<TaskRecord>>(emptyList())
val tasks: StateFlow<List<TaskRecord>> = _tasks.asStateFlow()
init {
_tasks.value = history.load().take(50)
}
/** 入队。返回任务 id。 */
fun submit(job: HubJob): String {
val id = UUID.randomUUID().toString().take(8)
jobs[id] = job
cancels[id] = AtomicBoolean(false)
val rec = TaskRecord(id = id, kind = job.kind, title = job.title, detail = job.detail)
publish { it + rec }
queue.addLast(id)
worker.submit { drain() }
return id
}
fun cancel(id: String) {
cancels[id]?.set(true)
// 还没开跑的直接标 Cancelled,不用等 worker
update(id) {
if (it.state == TaskState.Queued) it.copy(state = TaskState.Cancelled, endedAt = System.currentTimeMillis()) else it
}
}
fun clearFinished() {
val kept = _tasks.value.filter { !it.finished }
_tasks.value = kept.sortedByDescending { it.createdAt }
history.save(kept.filter { it.finished })
}
private fun drain() {
while (true) {
val id = synchronized(queue) { queue.removeFirstOrNull() } ?: return
val job = jobs[id] ?: continue
val cancel = cancels[id] ?: continue
update(id) { it.copy(state = TaskState.Running, startedAt = System.currentTimeMillis()) }
try {
job.run(
JobCtl(workDir, cancel) { f, note ->
update(id) { it.copy(progress = f, detail = note ?: it.detail) }
},
)
update(id) {
it.copy(
state = TaskState.Done,
progress = 1f,
endedAt = System.currentTimeMillis(),
)
}
} catch (e: CancelledException) {
update(id) { it.copy(state = TaskState.Cancelled, endedAt = System.currentTimeMillis()) }
} catch (e: Throwable) {
update(id) {
it.copy(
state = TaskState.Failed,
endedAt = System.currentTimeMillis(),
error = e.message ?: e.javaClass.simpleName,
)
}
} finally {
jobs.remove(id)
cancels.remove(id)
history.save(_tasks.value.filter { it.finished })
}
}
}
/** 产物/退出码由 job 自己回填 —— 用这个方法。 */
fun attachArtifact(id: String, path: String, sha256: String? = null, exitCode: Int? = null) {
update(id) { it.copy(artifact = path, sha256 = sha256, exitCode = exitCode) }
}
fun task(id: String): TaskRecord? = _tasks.value.firstOrNull { it.id == id }
private fun update(id: String, f: (TaskRecord) -> TaskRecord) {
publish { list -> list.map { if (it.id == id) f(it) else it } }
}
private fun publish(f: (List<TaskRecord>) -> List<TaskRecord>) {
_tasks.value = f(_tasks.value)
}
}
/** 只落**终态**的历史:队列不持久化(理由见 [TaskScheduler])。 */
class TaskHistory(private val file: File) {
private val store = cn.nfarms.core.store.JsonStore(file)
private val json = cn.nfarms.core.store.JsonStore.json()
fun load(): List<TaskRecord> =
store.read(cn.nfarms.core.store.ListSerializer(TaskRecord.serializer()).let { it }).orEmpty()
fun save(list: List<TaskRecord>) {
runCatching { store.write(cn.nfarms.core.store.ListSerializer(TaskRecord.serializer()), list.take(100)) }
}
@Suppress("unused")
private fun unused() = json
}