494 lines
16 KiB
Go
494 lines
16 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"encoding/xml"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"net"
|
|
"net/http"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/jmoiron/sqlx"
|
|
"vdi-cloud/platform/internal/model"
|
|
"vdi-cloud/platform/internal/store"
|
|
"vdi-cloud/platform/pkg/proxmox"
|
|
)
|
|
|
|
// This optional, single-host group is provisioned by the administrator. Device
|
|
// requests may select a bound desktop; they cannot choose VMIDs or PCI devices.
|
|
const exclusiveGPULock int64 = 7651092039
|
|
|
|
type gpuPrepareState struct {
|
|
Phase string `json:"phase"`
|
|
Node string `json:"node"`
|
|
Members []int `json:"members"`
|
|
Terminal string `json:"terminal"`
|
|
VMID int `json:"vmid,omitempty"`
|
|
UPID string `json:"upid,omitempty"`
|
|
}
|
|
|
|
func (s *Server) exclusiveGPU(d model.Desktop) bool {
|
|
if d.Provider != "proxmox" || d.Node != s.Config.ExclusiveGPUNode {
|
|
return false
|
|
}
|
|
for _, id := range s.Config.ExclusiveGPUVMIDs {
|
|
if d.VMID == id {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// Maintenance submissions take the same group lock as terminal preparations.
|
|
// Call inside their existing transaction; the durable reservation also covers
|
|
// the interval between enqueueing a preparation and claiming it in the worker.
|
|
func (s *Server) exclusiveGPUMaintenance(ctx context.Context, db sqlx.ExtContext, d model.Desktop) error {
|
|
if !s.exclusiveGPU(d) {
|
|
return nil
|
|
}
|
|
locked, err := tryGPULock(ctx, db)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
var busy bool
|
|
if locked {
|
|
err = sqlx.GetContext(ctx, db, &busy, `SELECT EXISTS(SELECT 1 FROM tasks WHERE kind='gpu-prepare' AND status IN ('queued','running','timeout'))`)
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !locked || busy {
|
|
return errors.New("显卡正在切换或结果未确认，请先核对切换任务")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *Server) gpuMembers(ctx context.Context, db sqlx.ExtContext) ([]model.Desktop, error) {
|
|
var members []model.Desktop
|
|
err := sqlx.SelectContext(ctx, db, &members, `SELECT * FROM desktops WHERE provider='proxmox' AND node=$1 AND vmid=ANY($2) ORDER BY vmid`, s.Config.ExclusiveGPUNode, s.Config.ExclusiveGPUVMIDs)
|
|
if err == nil && len(members) != len(s.Config.ExclusiveGPUVMIDs) {
|
|
err = errors.New("exclusive GPU group is incomplete")
|
|
}
|
|
return members, err
|
|
}
|
|
|
|
func tryGPULock(ctx context.Context, db sqlx.ExtContext) (bool, error) {
|
|
var locked bool
|
|
err := sqlx.GetContext(ctx, db, &locked, "SELECT pg_try_advisory_xact_lock($1)", exclusiveGPULock)
|
|
return locked, err
|
|
}
|
|
|
|
func (s *Server) gpuAllowsConnection(c *gin.Context, d model.Desktop) bool {
|
|
if !s.exclusiveGPU(d) {
|
|
return true
|
|
}
|
|
locked, err := tryGPULock(c.Request.Context(), deviceTx(c))
|
|
var busy bool
|
|
if err == nil && locked {
|
|
err = deviceTx(c).GetContext(c.Request.Context(), &busy, `SELECT EXISTS(SELECT 1 FROM tasks WHERE kind='gpu-prepare' AND status IN ('queued','running','timeout'))`)
|
|
}
|
|
if err != nil {
|
|
s.dbError(c, err)
|
|
return false
|
|
}
|
|
if !locked || busy {
|
|
s.fail(c, 409, "gpu_switch_busy", "显卡正在切换或上次切换尚未确认，请稍后重试或联系管理员")
|
|
return false
|
|
}
|
|
pve := s.currentPVE()
|
|
if pve == nil {
|
|
s.fail(c, 503, "pve_unavailable", "计算平台不可用")
|
|
return false
|
|
}
|
|
for _, id := range s.Config.ExclusiveGPUVMIDs {
|
|
power, e := pve.Power(c.Request.Context(), d.Node, id)
|
|
expected := "stopped"
|
|
if id == d.VMID {
|
|
expected = "running"
|
|
}
|
|
if e != nil || power != expected {
|
|
s.fail(c, 409, "gpu_owner_not_ready", "桌面电源或显卡归属尚未就绪，请先准备桌面")
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
func gpuGroupBusy(ctx context.Context, db sqlx.ExtContext, members []model.Desktop, taskID string) (bool, error) {
|
|
ids := make([]string, 0, len(members))
|
|
for _, d := range members {
|
|
ids = append(ids, d.ID)
|
|
}
|
|
var busy bool
|
|
err := sqlx.GetContext(ctx, db, &busy, `SELECT EXISTS(SELECT 1 FROM terminals WHERE session_desktop_id=ANY($1) AND last_seen>now()-interval '90 seconds') OR EXISTS(SELECT 1 FROM tasks WHERE target_id=ANY($1) AND id!=$2 AND kind!='terminal' AND status IN ('queued','running','timeout'))`, ids, taskID)
|
|
return busy, err
|
|
}
|
|
|
|
func (s *Server) prepareGPUDesktop(c *gin.Context) {
|
|
var in struct {
|
|
DesktopID string `json:"desktopId" binding:"required,max=100"`
|
|
}
|
|
if !s.decode(c, &in) {
|
|
return
|
|
}
|
|
tx := deviceTx(c)
|
|
ctx := c.Request.Context()
|
|
terminal := c.GetString("terminalId")
|
|
var d model.Desktop
|
|
err := tx.GetContext(ctx, &d, `SELECT d.* FROM desktops d JOIN bindings b ON b.desktop_id=d.id JOIN users u ON u.id=b.user_id JOIN terminals t ON t.id=u.terminal_id WHERE d.id=$1 AND t.id=$2 AND u.enabled AND t.enabled`, in.DesktopID, terminal)
|
|
if err == sql.ErrNoRows {
|
|
s.fail(c, 403, "desktop_forbidden", "该终端没有桌面使用权限")
|
|
return
|
|
}
|
|
if err != nil {
|
|
s.dbError(c, err)
|
|
return
|
|
}
|
|
if !s.exclusiveGPU(d) || d.SetupRequired || d.Endpoint == "" {
|
|
s.fail(c, 409, "desktop_not_preparable", "该桌面未启用显卡切换或连接配置尚未完成")
|
|
return
|
|
}
|
|
locked, err := tryGPULock(ctx, tx)
|
|
if err != nil {
|
|
s.dbError(c, err)
|
|
return
|
|
}
|
|
if !locked {
|
|
s.fail(c, 409, "gpu_switch_busy", "显卡正在切换，请稍后重试")
|
|
return
|
|
}
|
|
var previous model.Task
|
|
err = tx.GetContext(ctx, &previous, `SELECT * FROM tasks WHERE kind='gpu-prepare' AND status IN ('queued','running','timeout') ORDER BY created_at DESC LIMIT 1`)
|
|
if err == nil {
|
|
if previous.Actor == "device:"+terminal && previous.TargetID == d.ID && previous.Status != "timeout" {
|
|
s.deviceJSON(c, 202, gin.H{"taskId": previous.ID, "status": previous.Status, "message": previous.Message})
|
|
return
|
|
}
|
|
s.fail(c, 409, "gpu_switch_busy", "已有显卡切换任务或未确认结果，请稍后重试或联系管理员")
|
|
return
|
|
}
|
|
if err != sql.ErrNoRows {
|
|
s.dbError(c, err)
|
|
return
|
|
}
|
|
members, err := s.gpuMembers(ctx, tx)
|
|
if err != nil {
|
|
s.dbError(c, err)
|
|
return
|
|
}
|
|
for _, member := range members {
|
|
if !s.hardwareAllowsConnection(c, member.ID) {
|
|
return
|
|
}
|
|
}
|
|
busy, err := gpuGroupBusy(ctx, tx, members, "")
|
|
if err != nil {
|
|
s.dbError(c, err)
|
|
return
|
|
}
|
|
if busy {
|
|
s.fail(c, 409, "gpu_group_in_use", "请先退出正在使用的云桌面；其他终端使用中或管理任务未结束时无法切换")
|
|
return
|
|
}
|
|
state := gpuPrepareState{Phase: "queued", Node: d.Node, Members: s.Config.ExclusiveGPUVMIDs, Terminal: terminal}
|
|
encoded, _ := json.Marshal(state)
|
|
id := store.ID("task_")
|
|
_, err = tx.ExecContext(ctx, `INSERT INTO tasks(id,kind,action,target_id,target_name,actor,status,message,provider_ref,idempotency_key) VALUES($1,'gpu-prepare','prepare-desktop',$2,$3,$4,'queued','正在准备桌面，等待显卡切换',$5,$1)`, id, d.ID, d.Name, "device:"+terminal, string(encoded))
|
|
if err != nil {
|
|
s.dbError(c, err)
|
|
return
|
|
}
|
|
s.deviceJSON(c, 202, gin.H{"taskId": id, "status": "queued", "message": "正在准备桌面，首次启动可能需要两分钟"})
|
|
}
|
|
|
|
func (s *Server) gpuPrepareStatus(c *gin.Context) {
|
|
var task model.Task
|
|
err := deviceTx(c).GetContext(c.Request.Context(), &task, `SELECT * FROM tasks WHERE id=$1 AND kind='gpu-prepare' AND actor=$2`, c.Param("id"), "device:"+c.GetString("terminalId"))
|
|
if err == sql.ErrNoRows {
|
|
s.fail(c, 404, "task_not_found", "切换任务不存在")
|
|
return
|
|
}
|
|
if err != nil {
|
|
s.dbError(c, err)
|
|
return
|
|
}
|
|
s.deviceJSON(c, 200, gin.H{"taskId": task.ID, "status": task.Status, "message": task.Message})
|
|
}
|
|
|
|
func (s *Server) saveGPUStage(ctx context.Context, t model.Task, state gpuPrepareState, message string) error {
|
|
data, _ := json.Marshal(state)
|
|
result, err := s.DB.ExecContext(ctx, `UPDATE tasks SET provider_ref=$1,message=$2,updated_at=now() WHERE id=$3 AND status='running'`, string(data), message, t.ID)
|
|
if err == nil {
|
|
n, _ := result.RowsAffected()
|
|
if n != 1 {
|
|
err = errors.New("GPU task reservation lost")
|
|
}
|
|
}
|
|
return err
|
|
}
|
|
|
|
func sunshineReady(ctx context.Context, endpoint string) bool {
|
|
host := endpoint
|
|
if parsed, _, err := net.SplitHostPort(endpoint); err == nil {
|
|
host = parsed
|
|
}
|
|
// Platform endpoints are hostnames/IPs, never arbitrary URLs or paths.
|
|
if strings.ContainsAny(host, "/\\@?#") || host == "" {
|
|
return false
|
|
}
|
|
req, err := http.NewRequestWithContext(ctx, "GET", "http://"+net.JoinHostPort(host, "47989")+"/serverinfo", nil)
|
|
if err != nil {
|
|
return false
|
|
}
|
|
client := http.Client{Timeout: 2 * time.Second, Transport: &http.Transport{Proxy: nil}, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
|
|
defer client.CloseIdleConnections()
|
|
response, err := client.Do(req)
|
|
if err != nil {
|
|
return false
|
|
}
|
|
defer response.Body.Close()
|
|
var info struct {
|
|
XMLName xml.Name `xml:"root"`
|
|
Status string `xml:"status_code,attr"`
|
|
Hostname string `xml:"hostname"`
|
|
}
|
|
return response.StatusCode == 200 && xml.NewDecoder(io.LimitReader(response.Body, 65536)).Decode(&info) == nil && info.Status == "200" && info.Hostname != ""
|
|
}
|
|
|
|
func (s *Server) executeGPUPrepare(parent context.Context, t model.Task) {
|
|
ctx, cancel := context.WithTimeout(parent, 180*time.Second)
|
|
defer cancel()
|
|
state := gpuPrepareState{}
|
|
err := json.Unmarshal([]byte(t.ProviderRef), &state)
|
|
uncertain := false
|
|
if err == nil && state.Phase != "queued" && state.Phase != "" {
|
|
uncertain = true
|
|
}
|
|
if err == nil && (state.Phase != "queued" || state.Node != s.Config.ExclusiveGPUNode || fmt.Sprint(state.Members) != fmt.Sprint(s.Config.ExclusiveGPUVMIDs)) {
|
|
err = errors.New("GPU group changed or task already submitted")
|
|
}
|
|
var tx *sqlx.Tx
|
|
if err == nil {
|
|
tx, err = s.DB.BeginTxx(ctx, nil)
|
|
}
|
|
if tx != nil {
|
|
defer tx.Rollback()
|
|
}
|
|
// Hold this lock on a dedicated transaction throughout the switch. Tickets
|
|
// fail promptly instead of waiting behind a VM boot; no root credential leaves
|
|
// the platform. An API restart never automatically replays an uncertain POST.
|
|
if err == nil {
|
|
var locked bool
|
|
locked, err = tryGPULock(ctx, tx)
|
|
if err == nil && !locked {
|
|
err = errors.New("GPU group locked")
|
|
}
|
|
}
|
|
var members []model.Desktop
|
|
if err == nil {
|
|
members, err = s.gpuMembers(ctx, tx)
|
|
}
|
|
var target model.Desktop
|
|
if err == nil {
|
|
for _, d := range members {
|
|
if d.ID == t.TargetID {
|
|
target = d
|
|
}
|
|
}
|
|
var allowed bool
|
|
err = tx.GetContext(ctx, &allowed, `SELECT EXISTS(SELECT 1 FROM bindings b JOIN users u ON u.id=b.user_id JOIN terminals term ON term.id=u.terminal_id WHERE b.desktop_id=$1 AND u.terminal_id=$2 AND u.enabled AND term.enabled)`, t.TargetID, state.Terminal)
|
|
if err == nil && (!allowed || target.ID == "" || target.SetupRequired || target.Endpoint == "") {
|
|
err = errors.New("desktop authorization/configuration changed")
|
|
}
|
|
}
|
|
if err == nil {
|
|
var busy bool
|
|
busy, err = gpuGroupBusy(ctx, tx, members, t.ID)
|
|
if err == nil && busy {
|
|
err = errors.New("GPU group is in use")
|
|
}
|
|
}
|
|
if err == nil {
|
|
err = s.RefreshCompute(ctx)
|
|
}
|
|
pve := s.currentPVE()
|
|
if err == nil && pve == nil {
|
|
err = errors.New("PVE unavailable")
|
|
}
|
|
action := func(d model.Desktop, operation string) error {
|
|
state.Phase = operation + "-submitting"
|
|
state.VMID = d.VMID
|
|
state.UPID = ""
|
|
if err := s.saveGPUStage(ctx, t, state, "正在"+map[string]string{"start": "启动所选桌面", "shutdown": "正常关闭上一台桌面"}[operation]); err != nil {
|
|
return err
|
|
}
|
|
uncertain = true
|
|
upid, err := pve.Action(ctx, d.Node, d.VMID, operation)
|
|
if err != nil {
|
|
var rejected *proxmox.HTTPError
|
|
if errors.As(err, &rejected) && rejected.Status >= 400 && rejected.Status < 500 {
|
|
uncertain = false
|
|
}
|
|
return err
|
|
}
|
|
state.Phase = operation + "-waiting"
|
|
state.UPID = upid
|
|
if err = s.saveGPUStage(ctx, t, state, "正在等待虚拟机电源操作完成"); err != nil {
|
|
return err
|
|
}
|
|
if err = pve.Wait(ctx, d.Node, upid, time.Second); err != nil {
|
|
return err
|
|
}
|
|
power, err := pve.Power(ctx, d.Node, d.VMID)
|
|
expected := "running"
|
|
if operation == "shutdown" {
|
|
expected = "stopped"
|
|
}
|
|
if err == nil && power != expected {
|
|
err = errors.New("VM power transition not confirmed")
|
|
}
|
|
if err == nil {
|
|
uncertain = false
|
|
}
|
|
return err
|
|
}
|
|
if err == nil {
|
|
for _, d := range members {
|
|
if d.ID == target.ID {
|
|
continue
|
|
}
|
|
power, e := pve.Power(ctx, d.Node, d.VMID)
|
|
if e != nil {
|
|
err = e
|
|
break
|
|
}
|
|
if power == "running" {
|
|
err = action(d, "shutdown")
|
|
} else if power != "stopped" {
|
|
err = errors.New("unexpected VM power state")
|
|
}
|
|
if err != nil {
|
|
break
|
|
}
|
|
}
|
|
}
|
|
if err == nil {
|
|
power, e := pve.Power(ctx, target.Node, target.VMID)
|
|
err = e
|
|
if err == nil && power == "stopped" {
|
|
err = action(target, "start")
|
|
} else if err == nil && power != "running" {
|
|
err = errors.New("unexpected target state")
|
|
}
|
|
}
|
|
if err == nil {
|
|
state.Phase = "service-waiting"
|
|
err = s.saveGPUStage(ctx, t, state, "桌面已开机，正在等待连接服务")
|
|
for err == nil && !sunshineReady(ctx, target.Endpoint) {
|
|
select {
|
|
case <-ctx.Done():
|
|
err = ctx.Err()
|
|
case <-time.After(time.Second):
|
|
}
|
|
}
|
|
}
|
|
if err == nil {
|
|
// Recheck all owners immediately before releasing the reservation.
|
|
for _, d := range members {
|
|
power, e := pve.Power(ctx, d.Node, d.VMID)
|
|
expected := "stopped"
|
|
if d.ID == target.ID {
|
|
expected = "running"
|
|
}
|
|
if e != nil || power != expected {
|
|
err = errors.New("final GPU ownership check failed")
|
|
break
|
|
}
|
|
if _, e = s.DB.ExecContext(ctx, `UPDATE desktops SET power_state=$1,service_state='unknown',updated_at=now() WHERE id=$2`, power, d.ID); e != nil {
|
|
err = e
|
|
break
|
|
}
|
|
}
|
|
}
|
|
status, message := "success", "桌面已启动，连接服务已响应"
|
|
if err != nil {
|
|
log.Printf("GPU prepare task=%s phase=%s error=%v", t.ID, state.Phase, err)
|
|
status = "failed"
|
|
message = "桌面准备失败，请检查桌面授权、占用状态及 PVE"
|
|
if uncertain {
|
|
status = "timeout"
|
|
message = "电源操作或服务状态尚未确认，已阻止再次切换；请管理员核对 PVE 原任务"
|
|
}
|
|
}
|
|
finalCtx, finish := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer finish()
|
|
if err == nil {
|
|
state.Phase = "completed"
|
|
_ = s.saveGPUStage(finalCtx, t, state, message)
|
|
}
|
|
_ = s.complete(finalCtx, t, status, message)
|
|
}
|
|
|
|
// Reconcile only a known original PVE task. An unknown submission requires
|
|
// manual operator investigation; this endpoint never resubmits a power action.
|
|
func (s *Server) reconcileGPUPrepare(c *gin.Context) {
|
|
var task model.Task
|
|
if err := s.DB.GetContext(c.Request.Context(), &task, `SELECT * FROM tasks WHERE id=$1 AND kind='gpu-prepare' AND status='timeout'`, c.Param("id")); err != nil {
|
|
s.fail(c, 409, "not_uncertain", "该任务无需核对或不存在")
|
|
return
|
|
}
|
|
var state gpuPrepareState
|
|
if json.Unmarshal([]byte(task.ProviderRef), &state) != nil || state.UPID == "" {
|
|
s.fail(c, 409, "unknown_submission", "未记录原 PVE 任务编号，需要在 PVE 手工核对，不能自动重发")
|
|
return
|
|
}
|
|
tx, err := s.DB.BeginTxx(c.Request.Context(), nil)
|
|
if err != nil {
|
|
s.dbError(c, err)
|
|
return
|
|
}
|
|
defer tx.Rollback()
|
|
locked, err := tryGPULock(c.Request.Context(), tx)
|
|
if err != nil || !locked {
|
|
s.fail(c, 409, "gpu_switch_busy", "显卡正在切换")
|
|
return
|
|
}
|
|
pve := s.currentPVE()
|
|
if pve == nil {
|
|
s.fail(c, 503, "pve_unavailable", "计算平台不可用")
|
|
return
|
|
}
|
|
original, err := pve.Task(c.Request.Context(), state.Node, state.UPID)
|
|
if err != nil || original.Status != "stopped" {
|
|
s.fail(c, 409, "pve_task_pending", "原 PVE 任务未结束或无法核对")
|
|
return
|
|
}
|
|
for _, id := range state.Members {
|
|
power, e := pve.Power(c.Request.Context(), state.Node, id)
|
|
if e != nil || (power != "running" && power != "stopped") {
|
|
s.fail(c, 409, "power_unknown", "虚拟机实际电源状态未确认")
|
|
return
|
|
}
|
|
}
|
|
_, err = tx.ExecContext(c.Request.Context(), `UPDATE tasks SET status='failed',message='管理员已核对原 PVE 任务及实际电源状态，可重新准备桌面',updated_at=now() WHERE id=$1 AND status='timeout'`, task.ID)
|
|
if err == nil {
|
|
_, err = tx.ExecContext(c.Request.Context(), `INSERT INTO audit_logs VALUES($1,$2,'reconcile-gpu-switch',$3,'success',$4,now())`, store.ID("audit_"), c.MustGet("admin").(model.Administrator).Username, task.TargetName, task.ID)
|
|
}
|
|
if err == nil {
|
|
err = tx.Commit()
|
|
}
|
|
if err != nil {
|
|
s.dbError(c, err)
|
|
return
|
|
}
|
|
c.JSON(200, gin.H{"reconciled": true, "vmid": strconv.Itoa(state.VMID)})
|
|
}
|