212 lines
7.4 KiB
Go
212 lines
7.4 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
|
|
"vdi-cloud/platform/internal/model"
|
|
"vdi-cloud/platform/internal/store"
|
|
"vdi-cloud/platform/pkg/proxmox"
|
|
)
|
|
|
|
type exclusiveTestPVE struct {
|
|
mu sync.Mutex
|
|
power map[int]string
|
|
actions []string
|
|
mode string
|
|
confirmed bool
|
|
}
|
|
|
|
func exclusiveFixture(t *testing.T, mode string) (*fixture, *exclusiveTestPVE) {
|
|
t.Helper()
|
|
f := newFixtureWorker(t, false)
|
|
f.server.Config.ExclusiveGPUNode = "gpu-node"
|
|
f.server.Config.ExclusiveGPUVMIDs = []int{100, 101, 102}
|
|
for i, id := range []string{"desk_demo_windows", "desk_demo_kylin", "desk_demo_ubuntu"} {
|
|
if _, err := f.db.Exec(`UPDATE desktops SET provider='proxmox',node='gpu-node',vmid=$1,endpoint='127.0.0.1',setup_required=false,power_state='stopped' WHERE id=$2`, 100+i, id); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
if _, err := f.db.Exec(`INSERT INTO terminals(id,sn,name,bootstrap_hash,token_hash,last_seen) VALUES('gpu-box','GPU-BOX','GPU test','',$1,now());`, store.Hash("gpu-device")); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := f.db.Exec(`UPDATE users SET terminal_id='gpu-box' WHERE id='user_demo_wang'; INSERT INTO bindings(user_id,desktop_id) SELECT 'user_demo_wang',id FROM desktops ON CONFLICT DO NOTHING`); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
p := &exclusiveTestPVE{power: map[int]string{100: "running", 101: "stopped", 102: "stopped"}, mode: mode}
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
path := strings.TrimPrefix(r.URL.Path, "/api2/json")
|
|
var data any
|
|
var id int
|
|
switch {
|
|
case strings.Contains(path, "/tasks/"):
|
|
if p.mode == "poll-failure" {
|
|
w.WriteHeader(503)
|
|
return
|
|
}
|
|
for _, action := range p.actions {
|
|
if action == "100:shutdown" {
|
|
p.power[100] = "stopped"
|
|
p.confirmed = true
|
|
}
|
|
if action == "102:start" {
|
|
p.power[102] = "running"
|
|
}
|
|
}
|
|
data = map[string]string{"status": "stopped", "exitstatus": "OK"}
|
|
case strings.HasSuffix(path, "/status/current"):
|
|
fmt.Sscanf(path, "/nodes/gpu-node/qemu/%d/status/current", &id)
|
|
data = map[string]string{"status": p.power[id]}
|
|
case r.Method == "POST" && strings.Contains(path, "/status/"):
|
|
fmt.Sscanf(path, "/nodes/gpu-node/qemu/%d/status/", &id)
|
|
action := path[strings.LastIndex(path, "/")+1:]
|
|
if action == "start" && (!p.confirmed || p.power[100] != "stopped") {
|
|
t.Error("start before confirmed shutdown")
|
|
}
|
|
p.actions = append(p.actions, fmt.Sprintf("%d:%s", id, action))
|
|
if p.mode == "ambiguous" {
|
|
w.WriteHeader(500)
|
|
return
|
|
}
|
|
if p.mode == "rejected" {
|
|
w.WriteHeader(403)
|
|
return
|
|
}
|
|
data = "UPID:gpu-test"
|
|
default:
|
|
t.Error("unexpected GPU switch request", r.Method, path)
|
|
http.NotFound(w, r)
|
|
return
|
|
}
|
|
json.NewEncoder(w).Encode(map[string]any{"data": data})
|
|
}))
|
|
t.Cleanup(srv.Close)
|
|
f.server.PVE = &proxmox.Client{BaseURL: srv.URL, HTTP: srv.Client()}
|
|
return f, p
|
|
}
|
|
|
|
func prepareExclusive(t *testing.T, f *fixture) model.Task {
|
|
t.Helper()
|
|
code, out := f.request("POST", "/terminal/desktop-prepare", "gpu-device", map[string]any{"desktopId": "desk_demo_ubuntu"}, "")
|
|
expect(t, 202, code, out)
|
|
var task model.Task
|
|
if err := f.db.Get(&task, `UPDATE tasks SET status='running' WHERE id=$1 RETURNING *`, out["taskId"]); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return task
|
|
}
|
|
|
|
func TestExclusiveGPUAuthorizationAndOccupiedGroup(t *testing.T) {
|
|
f, p := exclusiveFixture(t, "normal")
|
|
code, out := f.request("POST", "/terminal/desktop-prepare", "gpu-device", map[string]any{"desktopId": "unbound"}, "")
|
|
expect(t, 403, code, out)
|
|
if _, err := f.db.Exec(`UPDATE terminals SET session_desktop_id='desk_demo_windows' WHERE id='gpu-box'`); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
code, out = f.request("POST", "/terminal/desktop-prepare", "gpu-device", map[string]any{"desktopId": "desk_demo_ubuntu"}, "")
|
|
expect(t, 409, code, out)
|
|
if len(p.actions) != 0 {
|
|
t.Fatal("occupied group changed power", p.actions)
|
|
}
|
|
}
|
|
|
|
func TestExclusiveGPUReservationBlocksAdministrativePower(t *testing.T) {
|
|
f, p := exclusiveFixture(t, "normal")
|
|
prepareExclusive(t, f)
|
|
admin := f.login("admin")
|
|
for _, action := range []string{"start", "shutdown", "reboot"} {
|
|
code, out := f.request("POST", "/desktops/desk_demo_windows/actions", admin, map[string]any{"action": action}, "gpu-admin-"+action)
|
|
expect(t, 409, code, out)
|
|
}
|
|
_, check, err := f.server.powerBatchPreflight(context.Background(), f.db, "desk_demo_windows", "start", "")
|
|
if err != nil || check.Code != "exclusive_gpu_power" || len(p.actions) != 0 {
|
|
t.Fatal(check, err, p.actions)
|
|
}
|
|
}
|
|
|
|
func TestExclusiveGPUShutdownThenStartAndStaleCacheTicket(t *testing.T) {
|
|
f, p := exclusiveFixture(t, "normal")
|
|
listener, err := net.Listen("tcp", "127.0.0.1:47989")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
ready := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
fmt.Fprint(w, `<root status_code="200"><hostname>test-sunshine</hostname></root>`)
|
|
})}
|
|
go ready.Serve(listener)
|
|
t.Cleanup(func() { ready.Close() })
|
|
task := prepareExclusive(t, f)
|
|
f.server.executeGPUPrepare(context.Background(), task)
|
|
var status string
|
|
f.db.Get(&status, "SELECT status FROM tasks WHERE id=$1", task.ID)
|
|
if status != "success" {
|
|
t.Fatal("prepare failed", status)
|
|
}
|
|
if fmt.Sprint(p.actions) != "[100:shutdown 102:start]" {
|
|
t.Fatal(p.actions)
|
|
}
|
|
// Sync cache may lag behind real PVE state after an immediate start.
|
|
f.db.Exec(`UPDATE desktops SET power_state='stopped' WHERE id='desk_demo_ubuntu'`)
|
|
code, out := f.request("POST", "/terminal/connect-ticket", "gpu-device", map[string]any{"desktopId": "desk_demo_ubuntu"}, "")
|
|
expect(t, 201, code, out)
|
|
code, out = f.request("POST", "/terminal/connect-ticket/consume", "gpu-device", map[string]any{"ticket": out["ticket"]}, "")
|
|
expect(t, 200, code, out)
|
|
}
|
|
|
|
func TestExclusiveGPUUnknownSubmissionNeverReplayed(t *testing.T) {
|
|
for _, mode := range []string{"ambiguous", "poll-failure", "rejected"} {
|
|
t.Run(mode, func(t *testing.T) {
|
|
f, p := exclusiveFixture(t, mode)
|
|
task := prepareExclusive(t, f)
|
|
f.server.executeGPUPrepare(context.Background(), task)
|
|
var result model.Task
|
|
f.db.Get(&result, "SELECT * FROM tasks WHERE id=$1", task.ID)
|
|
want := "timeout"
|
|
if mode == "rejected" {
|
|
want = "failed"
|
|
}
|
|
if result.Status != want || len(p.actions) != 1 {
|
|
t.Fatal(result.Status, p.actions)
|
|
}
|
|
if mode == "rejected" {
|
|
return
|
|
}
|
|
f.server.executeGPUPrepare(context.Background(), result)
|
|
if len(p.actions) != 1 {
|
|
t.Fatal("replayed uncertain power action", p.actions)
|
|
}
|
|
code, out := f.request("POST", "/terminal/desktop-prepare", "gpu-device", map[string]any{"desktopId": "desk_demo_ubuntu"}, "")
|
|
expect(t, 409, code, out)
|
|
code, out = f.request("POST", "/tasks/"+task.ID+"/gpu-reconcile", f.login("admin"), nil, "")
|
|
if mode == "ambiguous" {
|
|
expect(t, 409, code, out)
|
|
} else {
|
|
expect(t, 409, code, out) /* original task cannot be read */
|
|
}
|
|
if len(p.actions) != 1 {
|
|
t.Fatal("reconciliation changed power")
|
|
}
|
|
if mode == "poll-failure" {
|
|
p.mu.Lock()
|
|
p.mode = "normal"
|
|
p.mu.Unlock()
|
|
code, out = f.request("POST", "/tasks/"+task.ID+"/gpu-reconcile", f.login("admin"), nil, "")
|
|
expect(t, 200, code, out)
|
|
f.db.Get(&result, "SELECT * FROM tasks WHERE id=$1", task.ID)
|
|
if result.Status != "failed" || len(p.actions) != 1 {
|
|
t.Fatal("reconciliation must release reservation without replay", result.Status, p.actions)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|