mirror of
https://gitee.com/samwaf/SamWaf.git
synced 2026-11-01 10:41:43 +08:00
439 lines
12 KiB
Go
439 lines
12 KiB
Go
package wafupdate
|
|
|
|
import (
|
|
"errors"
|
|
"io"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
// 升级阶段 key。数组顺序即执行顺序，权重之和恒为 100（progress_test.go 有守护用例）。
|
|
const (
|
|
StageCheck = "check"
|
|
StageBackup = "backup"
|
|
StageManifest = "manifest"
|
|
StageDownload = "download"
|
|
StageVerify = "verify"
|
|
StageExtract = "extract"
|
|
StageReplace = "replace"
|
|
StageRestart = "restart"
|
|
)
|
|
|
|
// 升级整体状态
|
|
const (
|
|
UpdateStateIdle = "idle"
|
|
UpdateStateRunning = "running"
|
|
UpdateStateRestarting = "restarting"
|
|
UpdateStateSuccess = "success"
|
|
UpdateStateFailed = "failed"
|
|
UpdateStateCanceled = "canceled"
|
|
)
|
|
|
|
// 单个阶段的状态
|
|
const (
|
|
StageStatePending = "pending"
|
|
StageStateRunning = "running"
|
|
StageStateDone = "done"
|
|
StageStateWarn = "warn" // 出了问题但不阻断升级，例如备份失败
|
|
StageStateFailed = "failed"
|
|
StageStateSkipped = "skipped"
|
|
)
|
|
|
|
// ErrUpdateCanceled 用户在下载阶段主动取消。只有下载阶段可取消：
|
|
// 进入二进制替换后中断反而危险。
|
|
var ErrUpdateCanceled = errors.New("升级已被用户取消")
|
|
|
|
type stageDef struct {
|
|
Key string
|
|
Weight int
|
|
}
|
|
|
|
// 权重是经验值，只保证"下载占大头且是真百分比"。其余阶段没有阶段内进度，
|
|
// 到达即停在本档起点，完成才跳下一档——不做假的匀速插值。
|
|
var stageDefs = []stageDef{
|
|
{StageCheck, 5},
|
|
{StageBackup, 10},
|
|
{StageManifest, 5},
|
|
{StageDownload, 50},
|
|
{StageVerify, 10},
|
|
{StageExtract, 10},
|
|
{StageReplace, 5},
|
|
{StageRestart, 5},
|
|
}
|
|
|
|
// StageSnapshot 单阶段对外快照
|
|
type StageSnapshot struct {
|
|
Key string `json:"key"`
|
|
State string `json:"state"`
|
|
CostMs int64 `json:"cost_ms"`
|
|
Detail string `json:"detail"`
|
|
}
|
|
|
|
// ProgressSnapshot 升级进度对外快照
|
|
type ProgressSnapshot struct {
|
|
State string `json:"state"`
|
|
Stage string `json:"stage"`
|
|
StageIndex int `json:"stage_index"`
|
|
Percent int `json:"percent"`
|
|
Downloaded int64 `json:"downloaded"`
|
|
Total int64 `json:"total"` // 0 = 服务端未给 Content-Length
|
|
Speed int64 `json:"speed"` // 字节/秒
|
|
FromVersion string `json:"from_version"`
|
|
ToVersion string `json:"to_version"`
|
|
Channel string `json:"channel"`
|
|
StartedAt int64 `json:"started_at"`
|
|
Error string `json:"error"`
|
|
ErrorStage string `json:"error_stage"`
|
|
Stages []StageSnapshot `json:"stages"`
|
|
}
|
|
|
|
type stageState struct {
|
|
state string
|
|
detail string
|
|
startAt time.Time
|
|
costMs int64
|
|
}
|
|
|
|
// ProgressTracker 升级进度快照。刻意只放内存不落库：
|
|
// 升级是单实例的瞬时行为，重启之后本就该从"探测新版本号"接续，落库没有意义。
|
|
type ProgressTracker struct {
|
|
mu sync.RWMutex
|
|
state string
|
|
from string
|
|
to string
|
|
channel string
|
|
startedAt time.Time
|
|
stages []stageState
|
|
curIdx int // -1 表示尚未进入任何阶段
|
|
downloaded int64
|
|
total int64
|
|
speed int64
|
|
dlStartedAt time.Time
|
|
lastPercent int
|
|
errMsg string
|
|
errStage string
|
|
canceled atomic.Bool
|
|
}
|
|
|
|
// GlobalUpdateProgress 全局唯一的升级进度。与 global.GWAF_RUNTIME_IS_UPDATETING
|
|
// 同步置位：后者仍是并发互斥闸，这里只负责"说清楚跑到哪一步了"。
|
|
var GlobalUpdateProgress = NewProgressTracker()
|
|
|
|
func NewProgressTracker() *ProgressTracker {
|
|
t := &ProgressTracker{curIdx: -1, state: UpdateStateIdle}
|
|
t.resetStages()
|
|
return t
|
|
}
|
|
|
|
func (t *ProgressTracker) resetStages() {
|
|
t.stages = make([]stageState, len(stageDefs))
|
|
for i := range t.stages {
|
|
t.stages[i].state = StageStatePending
|
|
}
|
|
}
|
|
|
|
func stageIndexOf(key string) int {
|
|
for i, d := range stageDefs {
|
|
if d.Key == key {
|
|
return i
|
|
}
|
|
}
|
|
return -1
|
|
}
|
|
|
|
// Begin 开启一轮升级，清空上一轮的残留。
|
|
func (t *ProgressTracker) Begin(from, to, channel string) {
|
|
t.mu.Lock()
|
|
defer t.mu.Unlock()
|
|
t.state = UpdateStateRunning
|
|
t.from, t.to, t.channel = from, to, channel
|
|
t.startedAt = time.Now()
|
|
t.resetStages()
|
|
t.curIdx = -1
|
|
t.downloaded, t.total, t.speed = 0, 0, 0
|
|
t.lastPercent = 0
|
|
t.errMsg, t.errStage = "", ""
|
|
t.canceled.Store(false)
|
|
}
|
|
|
|
// StageStart 进入某阶段。nil receiver 安全，方便调用方不做判空。
|
|
func (t *ProgressTracker) StageStart(key string) {
|
|
if t == nil {
|
|
return
|
|
}
|
|
t.mu.Lock()
|
|
defer t.mu.Unlock()
|
|
i := stageIndexOf(key)
|
|
if i < 0 {
|
|
return
|
|
}
|
|
// 中间被跳过的阶段（如 GitHub 渠道不做 SHA 校验）标 skipped，避免一直挂在 pending
|
|
for j := 0; j < i; j++ {
|
|
if t.stages[j].state == StageStatePending {
|
|
t.stages[j].state = StageStateSkipped
|
|
}
|
|
}
|
|
t.curIdx = i
|
|
t.stages[i].state = StageStateRunning
|
|
t.stages[i].startAt = time.Now()
|
|
if key == StageDownload {
|
|
t.dlStartedAt = time.Now()
|
|
}
|
|
}
|
|
|
|
// StageDone 结束某阶段，detail 为空则保留原有 detail。
|
|
func (t *ProgressTracker) StageDone(key, detail string) {
|
|
t.finishStage(key, StageStateDone, detail)
|
|
}
|
|
|
|
// StageWarn 阶段有问题但不阻断升级（当前只有备份失败会走这里）。
|
|
func (t *ProgressTracker) StageWarn(key, detail string) {
|
|
t.finishStage(key, StageStateWarn, detail)
|
|
}
|
|
|
|
func (t *ProgressTracker) finishStage(key, state, detail string) {
|
|
if t == nil {
|
|
return
|
|
}
|
|
t.mu.Lock()
|
|
defer t.mu.Unlock()
|
|
i := stageIndexOf(key)
|
|
if i < 0 {
|
|
return
|
|
}
|
|
t.stages[i].state = state
|
|
if detail != "" {
|
|
t.stages[i].detail = detail
|
|
}
|
|
if !t.stages[i].startAt.IsZero() {
|
|
t.stages[i].costMs = time.Since(t.stages[i].startAt).Milliseconds()
|
|
}
|
|
if t.curIdx == i && i+1 < len(stageDefs) {
|
|
t.curIdx = i + 1
|
|
}
|
|
}
|
|
|
|
// SetDownload 刷新下载字节数与速率。
|
|
func (t *ProgressTracker) SetDownload(downloaded, total int64) {
|
|
if t == nil {
|
|
return
|
|
}
|
|
t.mu.Lock()
|
|
defer t.mu.Unlock()
|
|
t.downloaded, t.total = downloaded, total
|
|
if !t.dlStartedAt.IsZero() {
|
|
if elapsed := time.Since(t.dlStartedAt).Seconds(); elapsed > 0.2 {
|
|
t.speed = int64(float64(downloaded) / elapsed)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Fail 以失败收尾，记录卡在哪个阶段。
|
|
func (t *ProgressTracker) Fail(err error) {
|
|
if t == nil || err == nil {
|
|
return
|
|
}
|
|
t.mu.Lock()
|
|
defer t.mu.Unlock()
|
|
// 取消位也参与判定：下载流经 gzip/tar 等中间层时，ErrUpdateCanceled 可能被包成别的错误
|
|
if errors.Is(err, ErrUpdateCanceled) || t.canceled.Load() {
|
|
t.state = UpdateStateCanceled
|
|
} else {
|
|
t.state = UpdateStateFailed
|
|
}
|
|
t.errMsg = err.Error()
|
|
if t.curIdx >= 0 && t.curIdx < len(stageDefs) {
|
|
t.errStage = stageDefs[t.curIdx].Key
|
|
if t.stages[t.curIdx].state == StageStateRunning {
|
|
t.stages[t.curIdx].state = StageStateFailed
|
|
if !t.stages[t.curIdx].startAt.IsZero() {
|
|
t.stages[t.curIdx].costMs = time.Since(t.stages[t.curIdx].startAt).Milliseconds()
|
|
}
|
|
}
|
|
// 失败之后的阶段保持 pending：它们是"没走到"，不是"被绕过"。
|
|
// 若标成 skipped 会被计入权重，界面上的进度会在失败瞬间反而往前跳。
|
|
}
|
|
}
|
|
|
|
// MarkRestarting 二进制已替换完成，进入等待服务重启阶段。
|
|
// 此后 WebSocket 必断，界面靠轮询版本号判断是否就绪。
|
|
func (t *ProgressTracker) MarkRestarting() {
|
|
if t == nil {
|
|
return
|
|
}
|
|
t.mu.Lock()
|
|
defer t.mu.Unlock()
|
|
t.state = UpdateStateRestarting
|
|
i := stageIndexOf(StageRestart)
|
|
t.curIdx = i
|
|
t.stages[i].state = StageStateRunning
|
|
t.stages[i].startAt = time.Now()
|
|
}
|
|
|
|
// ReplaceDone 判断二进制替换是否真的完成。
|
|
// 用于识别"升级链路返回 nil 但其实什么都没做"的情况。
|
|
func (t *ProgressTracker) ReplaceDone() bool {
|
|
t.mu.RLock()
|
|
defer t.mu.RUnlock()
|
|
i := stageIndexOf(StageReplace)
|
|
return t.stages[i].state == StageStateDone
|
|
}
|
|
|
|
// Cancel 请求取消。只在下载阶段有意义：计数 Reader 每次 Read 检查该标志。
|
|
// 返回 false 表示当前状态不允许取消。
|
|
func (t *ProgressTracker) Cancel() bool {
|
|
t.mu.RLock()
|
|
canCancel := t.state == UpdateStateRunning && t.curIdx == stageIndexOf(StageDownload)
|
|
t.mu.RUnlock()
|
|
if !canCancel {
|
|
return false
|
|
}
|
|
t.canceled.Store(true)
|
|
return true
|
|
}
|
|
|
|
func (t *ProgressTracker) IsCanceled() bool { return t != nil && t.canceled.Load() }
|
|
|
|
func (t *ProgressTracker) State() string {
|
|
t.mu.RLock()
|
|
defer t.mu.RUnlock()
|
|
return t.state
|
|
}
|
|
|
|
// Snapshot 输出对外快照，百分比在此计算。
|
|
func (t *ProgressTracker) Snapshot() ProgressSnapshot {
|
|
t.mu.Lock()
|
|
defer t.mu.Unlock()
|
|
|
|
snap := ProgressSnapshot{
|
|
State: t.state,
|
|
StageIndex: t.curIdx + 1, // 对外从 1 开始计
|
|
Downloaded: t.downloaded,
|
|
Total: t.total,
|
|
Speed: t.speed,
|
|
FromVersion: t.from,
|
|
ToVersion: t.to,
|
|
Channel: t.channel,
|
|
Error: t.errMsg,
|
|
ErrorStage: t.errStage,
|
|
Stages: make([]StageSnapshot, 0, len(stageDefs)),
|
|
}
|
|
if !t.startedAt.IsZero() {
|
|
snap.StartedAt = t.startedAt.Unix()
|
|
}
|
|
if t.curIdx >= 0 && t.curIdx < len(stageDefs) {
|
|
snap.Stage = stageDefs[t.curIdx].Key
|
|
}
|
|
for i, d := range stageDefs {
|
|
snap.Stages = append(snap.Stages, StageSnapshot{
|
|
Key: d.Key,
|
|
State: t.stages[i].state,
|
|
CostMs: t.stages[i].costMs,
|
|
Detail: t.stages[i].detail,
|
|
})
|
|
}
|
|
|
|
switch t.state {
|
|
case UpdateStateIdle:
|
|
return snap
|
|
case UpdateStateSuccess:
|
|
snap.Percent = 100
|
|
t.lastPercent = 100
|
|
return snap
|
|
}
|
|
|
|
percent := 0
|
|
for i, d := range stageDefs {
|
|
switch t.stages[i].state {
|
|
case StageStateDone, StageStateWarn, StageStateSkipped:
|
|
percent += d.Weight
|
|
case StageStateRunning:
|
|
// 只有下载阶段有阶段内进度
|
|
if d.Key == StageDownload && t.total > 0 {
|
|
ratio := float64(t.downloaded) / float64(t.total)
|
|
if ratio > 1 {
|
|
ratio = 1
|
|
}
|
|
percent += int(float64(d.Weight) * ratio)
|
|
}
|
|
}
|
|
}
|
|
if percent > 99 && t.state != UpdateStateSuccess {
|
|
percent = 99 // 只有确认就绪才给 100
|
|
}
|
|
// 百分比只进不退：阶段被标 skipped 之类的边界不该让进度倒回去
|
|
if percent < t.lastPercent {
|
|
percent = t.lastPercent
|
|
}
|
|
t.lastPercent = percent
|
|
snap.Percent = percent
|
|
return snap
|
|
}
|
|
|
|
// Success 升级完成（新版本已就绪）。当前由前端探测到新版本号后自然呈现，
|
|
// 后端进程此时已被替换重启，保留该方法供测试与后续扩展使用。
|
|
func (t *ProgressTracker) Success() {
|
|
t.mu.Lock()
|
|
defer t.mu.Unlock()
|
|
t.state = UpdateStateSuccess
|
|
i := stageIndexOf(StageRestart)
|
|
t.stages[i].state = StageStateDone
|
|
}
|
|
|
|
// Reset 回到空闲态。
|
|
func (t *ProgressTracker) Reset() {
|
|
t.mu.Lock()
|
|
defer t.mu.Unlock()
|
|
t.state = UpdateStateIdle
|
|
t.resetStages()
|
|
t.curIdx = -1
|
|
t.downloaded, t.total, t.speed, t.lastPercent = 0, 0, 0, 0
|
|
t.errMsg, t.errStage = "", ""
|
|
t.canceled.Store(false)
|
|
}
|
|
|
|
// progressReader 包住下载流做两件事：统计字节数、响应取消。
|
|
// 更新做了限频（≥200ms 或每 1%），避免 32KB 一块就写一次锁。
|
|
type progressReader struct {
|
|
r io.Reader
|
|
t *ProgressTracker
|
|
total int64
|
|
read int64
|
|
lastNotify time.Time
|
|
lastPct int64
|
|
}
|
|
|
|
func newProgressReader(r io.Reader, t *ProgressTracker, total int64) io.Reader {
|
|
if t == nil {
|
|
return r
|
|
}
|
|
t.SetDownload(0, total)
|
|
return &progressReader{r: r, t: t, total: total, lastNotify: time.Now()}
|
|
}
|
|
|
|
func (p *progressReader) Read(b []byte) (int, error) {
|
|
if p.t.IsCanceled() {
|
|
return 0, ErrUpdateCanceled
|
|
}
|
|
n, err := p.r.Read(b)
|
|
if n > 0 {
|
|
p.read += int64(n)
|
|
notify := time.Since(p.lastNotify) >= 200*time.Millisecond
|
|
if !notify && p.total > 0 {
|
|
if pct := p.read * 100 / p.total; pct != p.lastPct {
|
|
p.lastPct = pct
|
|
notify = true
|
|
}
|
|
}
|
|
if notify || err == io.EOF {
|
|
p.lastNotify = time.Now()
|
|
p.t.SetDownload(p.read, p.total)
|
|
}
|
|
}
|
|
if err == io.EOF {
|
|
p.t.SetDownload(p.read, p.total)
|
|
}
|
|
return n, err
|
|
}
|