Spaces:

a3216
/

ds2api

Running

204848

feat: 添加HF Space卡死自愈功能，增强代理健康检查逻辑，更新配置文件

8605295 4 months ago

7.83 kB

/**
 * Cloudflare Worker — ds2api 代理健康检查 + HF Space 卡死自愈
 *
 * 部署步骤:
 *   1. cd cf-worker && npm install
 *   2. 设置 secrets:
 *        wrangler secret put DS2API_BASE_URL
 *        wrangler secret put ADMIN_KEY
 *        wrangler secret put VERCEL_BYPASS_TOKEN   (如果开了 Vercel Deployment Protection)
 *        wrangler secret put HF_TOKEN              (HF API token，用于 Space 卡死自愈)
 *   3. wrangler deploy
 *
 * 两个 cron（见 wrangler.toml）：
 *   - 5 小时一次（"0 0,5,10,15,20 * * *" 等价写法）：代理健康检查（runCheck）
 *   - 10 分钟一次：HF Space 卡死检测+自动 restart（runSpaceCheck）
 *
 * 手动触发：
 *   GET/POST https://<worker>/check          → 代理健康检查
 *   GET/POST https://<worker>/check-space    → Space 卡死检测+自愈
 */

const LOGIN_PATH = "/admin/login"
const CHECK_ALL_PATH = "/admin/proxies/check-all"

// Space 卡死自愈相关常量。namespace/name 通过 env.HF_SPACE_REPO
// （格式 "a3216/ds2api"）配置；不配则用默认值。
const DEFAULT_HF_SPACE_REPO = "a3216/ds2api"
// 重启冷却：避免连续 restart 把 Space 打爆。重启后至少等 5 分钟才允许再次重启。
const RESTART_COOLDOWN_MS = 5 * 60 * 1000
// KV 中存上次重启时间的 key。
const LAST_RESTART_KV_KEY = "last_restart"

function vercelHeaders(env) {
  const headers = {}
  const bypass = env.VERCEL_BYPASS_TOKEN
  if (bypass) {
    headers["x-vercel-protection-bypass"] = bypass
  }
  return headers
}

async function login(baseURL, adminKey, extraHeaders) {
  const res = await fetch(`${baseURL}${LOGIN_PATH}`, {
    method: "POST",
    headers: { "Content-Type": "application/json", ...extraHeaders },
    body: JSON.stringify({ admin_key: adminKey, expire_hours: 1 }),
  })
  if (!res.ok) {
    const text = await res.text()
    throw new Error(`login failed (${res.status}): ${text}`)
  }
  const data = await res.json()
  return data.token
}

async function checkAll(baseURL, token, extraHeaders) {
  const res = await fetch(`${baseURL}${CHECK_ALL_PATH}`, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      Authorization: `Bearer ${token}`,
      ...extraHeaders,
    },
  })
  if (!res.ok) {
    const text = await res.text()
    throw new Error(`check-all failed (${res.status}): ${text}`)
  }
  return res.json()
}

async function runCheck(env) {
  const baseURL = (env.DS2API_BASE_URL || "").replace(/\/+$/, "")
  const adminKey = env.ADMIN_KEY || ""

  if (!baseURL || !adminKey) {
    throw new Error("DS2API_BASE_URL and ADMIN_KEY secrets must be set")
  }

  const extraHeaders = vercelHeaders(env)
  const token = await login(baseURL, adminKey, extraHeaders)
  const result = await checkAll(baseURL, token, extraHeaders)
  return result
}

// === HF Space 卡死自愈 ===

// 查询 Space runtime 状态。返回 { stage, hardwareCurrent }。
async function fetchSpaceRuntime(repo, token) {
  const url = `https://huggingface.co/api/spaces/${repo}/runtime`
  const res = await fetch(url, {
    headers: { Authorization: `Bearer ${token}` },
  })
  if (!res.ok) {
    const text = await res.text()
    throw new Error(`runtime query failed (${res.status}): ${text}`)
  }
  const data = await res.json()
  return {
    stage: data.stage,
    hardwareCurrent: data.hardware ? data.hardware.current : null,
  }
}

// 判定 Space 是否卡死：非 RUNNING 或未分配硬件。
function isSpaceStuck(runtime) {
  if (!runtime) return true
  if (runtime.stage !== "RUNNING") return true
  if (!runtime.hardwareCurrent) return true
  return false
}

// 触发 HF Space restart。返回 true 表示请求已被接受。
async function restartSpace(repo, token) {
  const url = `https://huggingface.co/api/spaces/${repo}/restart`
  const res = await fetch(url, {
    method: "POST",
    headers: { Authorization: `Bearer ${token}` },
  })
  if (!res.ok) {
    const text = await res.text()
    throw new Error(`restart failed (${res.status}): ${text}`)
  }
  return true
}

// 从 KV 读上次重启时间戳（ms）。KV 未绑定或无值返回 0。
async function getLastRestart(kv) {
  if (!kv) return 0
  try {
    const v = await kv.get(LAST_RESTART_KV_KEY)
    const n = parseInt(v, 10)
    return Number.isFinite(n) ? n : 0
  } catch {
    return 0
  }
}

async function setLastRestart(kv) {
  if (!kv) return
  try {
    await kv.put(LAST_RESTART_KV_KEY, String(Date.now()))
  } catch {
    // KV 写失败不阻断主流程。
  }
}

// Space 卡死检测+自愈主逻辑。
async function runSpaceCheck(env) {
  const token = env.HF_TOKEN || ""
  if (!token) {
    throw new Error("HF_TOKEN secret must be set for space self-healing")
  }
  const repo = (env.HF_SPACE_REPO || DEFAULT_HF_SPACE_REPO).trim()

  const runtime = await fetchSpaceRuntime(repo, token)
  if (!isSpaceStuck(runtime)) {
    return {
      ok: true,
      action: "none",
      repo,
      stage: runtime.stage,
      hardwareCurrent: runtime.hardwareCurrent,
      message: "space healthy, no action",
    }
  }

  // 卡死，检查冷却。
  const last = await getLastRestart(env.SPACE_STATE)
  const elapsed = Date.now() - last
  if (elapsed < RESTART_COOLDOWN_MS) {
    const waitMs = RESTART_COOLDOWN_MS - elapsed
    return {
      ok: true,
      action: "cooldown",
      repo,
      stage: runtime.stage,
      hardwareCurrent: runtime.hardwareCurrent,
      message: `space stuck but in restart cooldown, wait ${Math.ceil(waitMs / 1000)}s`,
    }
  }

  // 触发重启。
  await restartSpace(repo, token)
  await setLastRestart(env.SPACE_STATE)
  return {
    ok: true,
    action: "restarted",
    repo,
    stage: runtime.stage,
    hardwareCurrent: runtime.hardwareCurrent,
    message: "space was stuck, restart triggered",
  }
}

// 区分两种 cron：5 小时一次的代理检查 vs 10 分钟一次的 Space 检查。
// scheduled.cron 即触发的 cron 表达式（字符串）。
function isSpaceCheckCron(cron) {
  return cron === "*/10 * * * *"
}

export default {
  async scheduled(event, env, ctx) {
    if (isSpaceCheckCron(event.cron)) {
      ctx.waitUntil(
        runSpaceCheck(env)
          .then((r) => {
            if (r.action === "restarted") {
              console.warn(`[space] ${r.message} repo=${r.repo} stage=${r.stage} hardware=${r.hardwareCurrent}`)
            } else {
              console.log(`[space] ${r.message} repo=${r.repo} stage=${r.stage} hardware=${r.hardwareCurrent}`)
            }
          })
          .catch((err) => {
            console.error(`[space] check failed: ${err.message}`)
          })
      )
      return
    }

    // 默认：代理健康检查 cron。
    ctx.waitUntil(
      runCheck(env)
        .then((result) => {
          const items = result.items || []
          const healthy = items.filter((i) => i.healthy && !i.disabled).length
          const banned = items.filter((i) => i.disabled).length
          console.log(
            `proxy health check done: ${healthy}/${items.length} healthy, ${banned} banned`
          )
        })
        .catch((err) => {
          console.error(`proxy health check failed: ${err.message}`)
        })
    )
  },

  async fetch(request, env) {
    const url = new URL(request.url)

    if (url.pathname === "/check") {
      try {
        const result = await runCheck(env)
        return Response.json(result)
      } catch (err) {
        return Response.json({ error: err.message }, { status: 500 })
      }
    }

    if (url.pathname === "/check-space") {
      try {
        const result = await runSpaceCheck(env)
        return Response.json(result)
      } catch (err) {
        return Response.json({ error: err.message }, { status: 500 })
      }
    }

    return new Response("not found", { status: 404 })
  },
}

	/**
	* Cloudflare Worker — ds2api 代理健康检查 + HF Space 卡死自愈
	*
	* 部署步骤:
	* 1. cd cf-worker && npm install
	* 2. 设置 secrets:
	* wrangler secret put DS2API_BASE_URL
	* wrangler secret put ADMIN_KEY
	* wrangler secret put VERCEL_BYPASS_TOKEN (如果开了 Vercel Deployment Protection)
	* wrangler secret put HF_TOKEN (HF API token，用于 Space 卡死自愈)
	* 3. wrangler deploy
	*
	* 两个 cron（见 wrangler.toml）：
	* - 5 小时一次（"0 0,5,10,15,20 * * *" 等价写法）：代理健康检查（runCheck）
	* - 10 分钟一次：HF Space 卡死检测+自动 restart（runSpaceCheck）
	*
	* 手动触发：
	* GET/POST https://<worker>/check → 代理健康检查
	* GET/POST https://<worker>/check-space → Space 卡死检测+自愈
	*/

	const LOGIN_PATH = "/admin/login"
	const CHECK_ALL_PATH = "/admin/proxies/check-all"

	// Space 卡死自愈相关常量。namespace/name 通过 env.HF_SPACE_REPO
	// （格式 "a3216/ds2api"）配置；不配则用默认值。
	const DEFAULT_HF_SPACE_REPO = "a3216/ds2api"
	// 重启冷却：避免连续 restart 把 Space 打爆。重启后至少等 5 分钟才允许再次重启。
	const RESTART_COOLDOWN_MS = 5 * 60 * 1000
	// KV 中存上次重启时间的 key。
	const LAST_RESTART_KV_KEY = "last_restart"

	function vercelHeaders(env) {
	const headers = {}
	const bypass = env.VERCEL_BYPASS_TOKEN
	if (bypass) {
	headers["x-vercel-protection-bypass"] = bypass
	}
	return headers
	}

	async function login(baseURL, adminKey, extraHeaders) {
	const res = await fetch(`${baseURL}${LOGIN_PATH}`, {
	method: "POST",
	headers: { "Content-Type": "application/json", ...extraHeaders },
	body: JSON.stringify({ admin_key: adminKey, expire_hours: 1 }),
	})
	if (!res.ok) {
	const text = await res.text()
	throw new Error(`login failed (${res.status}): ${text}`)
	}
	const data = await res.json()
	return data.token
	}

	async function checkAll(baseURL, token, extraHeaders) {
	const res = await fetch(`${baseURL}${CHECK_ALL_PATH}`, {
	method: "POST",
	headers: {
	"Content-Type": "application/json",
	Authorization: `Bearer ${token}`,
	...extraHeaders,
	},
	})
	if (!res.ok) {
	const text = await res.text()
	throw new Error(`check-all failed (${res.status}): ${text}`)
	}
	return res.json()
	}

	async function runCheck(env) {
	const baseURL = (env.DS2API_BASE_URL \|\| "").replace(/\/+$/, "")
	const adminKey = env.ADMIN_KEY \|\| ""

	if (!baseURL \|\| !adminKey) {
	throw new Error("DS2API_BASE_URL and ADMIN_KEY secrets must be set")
	}

	const extraHeaders = vercelHeaders(env)
	const token = await login(baseURL, adminKey, extraHeaders)
	const result = await checkAll(baseURL, token, extraHeaders)
	return result
	}

	// === HF Space 卡死自愈 ===

	// 查询 Space runtime 状态。返回 { stage, hardwareCurrent }。
	async function fetchSpaceRuntime(repo, token) {
	const url = `https://huggingface.co/api/spaces/${repo}/runtime`
	const res = await fetch(url, {
	headers: { Authorization: `Bearer ${token}` },
	})
	if (!res.ok) {
	const text = await res.text()
	throw new Error(`runtime query failed (${res.status}): ${text}`)
	}
	const data = await res.json()
	return {
	stage: data.stage,
	hardwareCurrent: data.hardware ? data.hardware.current : null,
	}
	}

	// 判定 Space 是否卡死：非 RUNNING 或未分配硬件。
	function isSpaceStuck(runtime) {
	if (!runtime) return true
	if (runtime.stage !== "RUNNING") return true
	if (!runtime.hardwareCurrent) return true
	return false
	}

	// 触发 HF Space restart。返回 true 表示请求已被接受。
	async function restartSpace(repo, token) {
	const url = `https://huggingface.co/api/spaces/${repo}/restart`
	const res = await fetch(url, {
	method: "POST",
	headers: { Authorization: `Bearer ${token}` },
	})
	if (!res.ok) {
	const text = await res.text()
	throw new Error(`restart failed (${res.status}): ${text}`)
	}
	return true
	}

	// 从 KV 读上次重启时间戳（ms）。KV 未绑定或无值返回 0。
	async function getLastRestart(kv) {
	if (!kv) return 0
	try {
	const v = await kv.get(LAST_RESTART_KV_KEY)
	const n = parseInt(v, 10)
	return Number.isFinite(n) ? n : 0
	} catch {
	return 0
	}
	}

	async function setLastRestart(kv) {
	if (!kv) return
	try {
	await kv.put(LAST_RESTART_KV_KEY, String(Date.now()))
	} catch {
	// KV 写失败不阻断主流程。
	}
	}

	// Space 卡死检测+自愈主逻辑。
	async function runSpaceCheck(env) {
	const token = env.HF_TOKEN \|\| ""
	if (!token) {
	throw new Error("HF_TOKEN secret must be set for space self-healing")
	}
	const repo = (env.HF_SPACE_REPO \|\| DEFAULT_HF_SPACE_REPO).trim()

	const runtime = await fetchSpaceRuntime(repo, token)
	if (!isSpaceStuck(runtime)) {
	return {
	ok: true,
	action: "none",
	repo,
	stage: runtime.stage,
	hardwareCurrent: runtime.hardwareCurrent,
	message: "space healthy, no action",
	}
	}

	// 卡死，检查冷却。
	const last = await getLastRestart(env.SPACE_STATE)
	const elapsed = Date.now() - last
	if (elapsed < RESTART_COOLDOWN_MS) {
	const waitMs = RESTART_COOLDOWN_MS - elapsed
	return {
	ok: true,
	action: "cooldown",
	repo,
	stage: runtime.stage,
	hardwareCurrent: runtime.hardwareCurrent,
	message: `space stuck but in restart cooldown, wait ${Math.ceil(waitMs / 1000)}s`,
	}
	}

	// 触发重启。
	await restartSpace(repo, token)
	await setLastRestart(env.SPACE_STATE)
	return {
	ok: true,
	action: "restarted",
	repo,
	stage: runtime.stage,
	hardwareCurrent: runtime.hardwareCurrent,
	message: "space was stuck, restart triggered",
	}
	}

	// 区分两种 cron：5 小时一次的代理检查 vs 10 分钟一次的 Space 检查。
	// scheduled.cron 即触发的 cron 表达式（字符串）。
	function isSpaceCheckCron(cron) {
	return cron === "/10 * * *"
	}

	export default {
	async scheduled(event, env, ctx) {
	if (isSpaceCheckCron(event.cron)) {
	ctx.waitUntil(
	runSpaceCheck(env)
	.then((r) => {
	if (r.action === "restarted") {
	console.warn(`[space] ${r.message} repo=${r.repo} stage=${r.stage} hardware=${r.hardwareCurrent}`)
	} else {
	console.log(`[space] ${r.message} repo=${r.repo} stage=${r.stage} hardware=${r.hardwareCurrent}`)
	}
	})
	.catch((err) => {
	console.error(`[space] check failed: ${err.message}`)
	})
	)
	return
	}

	// 默认：代理健康检查 cron。
	ctx.waitUntil(
	runCheck(env)
	.then((result) => {
	const items = result.items \|\| []
	const healthy = items.filter((i) => i.healthy && !i.disabled).length
	const banned = items.filter((i) => i.disabled).length
	console.log(
	`proxy health check done: ${healthy}/${items.length} healthy, ${banned} banned`
	)
	})
	.catch((err) => {
	console.error(`proxy health check failed: ${err.message}`)
	})
	)
	},

	async fetch(request, env) {
	const url = new URL(request.url)

	if (url.pathname === "/check") {
	try {
	const result = await runCheck(env)
	return Response.json(result)
	} catch (err) {
	return Response.json({ error: err.message }, { status: 500 })
	}
	}

	if (url.pathname === "/check-space") {
	try {
	const result = await runSpaceCheck(env)
	return Response.json(result)
	} catch (err) {
	return Response.json({ error: err.message }, { status: 500 })
	}
	}

	return new Response("not found", { status: 404 })
	},
	}