feat: add proxmox task polling job
This commit is contained in:
@@ -0,0 +1,285 @@
|
||||
package tasks
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"time"
|
||||
|
||||
"github.com/hibiken/asynq"
|
||||
)
|
||||
|
||||
const TypeProxmoxTaskPoll = "proxmox.task.poll"
|
||||
|
||||
type ProxmoxTaskPollPayload struct {
|
||||
ClusterID string `json:"cluster_id"`
|
||||
Node string `json:"node"`
|
||||
UPID string `json:"upid"`
|
||||
TargetType string `json:"target_type"`
|
||||
TargetID string `json:"target_id"`
|
||||
TenantID string `json:"tenant_id"`
|
||||
ProfileID string `json:"profile_id,omitempty"`
|
||||
Action string `json:"action"`
|
||||
SuccessStatus string `json:"success_status"`
|
||||
}
|
||||
|
||||
type ProxmoxTaskStatus struct {
|
||||
Status string
|
||||
ExitStatus string
|
||||
}
|
||||
|
||||
type ProxmoxTaskClient interface {
|
||||
GetTaskStatus(ctx context.Context, node string, upid string) (ProxmoxTaskStatus, error)
|
||||
}
|
||||
|
||||
type ProxmoxTaskClientResolver interface {
|
||||
ResolveTaskClient(ctx context.Context, clusterID string) (ProxmoxTaskClient, error)
|
||||
}
|
||||
|
||||
type VMStatusStore interface {
|
||||
SetVMStatus(ctx context.Context, vmID string, status string) error
|
||||
}
|
||||
|
||||
type AuditWriter interface {
|
||||
WriteAudit(ctx context.Context, event AuditEvent) error
|
||||
}
|
||||
|
||||
type AuditEvent struct {
|
||||
TenantID string
|
||||
ProfileID string
|
||||
Action string
|
||||
TargetType string
|
||||
TargetID string
|
||||
Metadata map[string]any
|
||||
}
|
||||
|
||||
type ProxmoxTaskPollHandler struct {
|
||||
resolver ProxmoxTaskClientResolver
|
||||
vmStore VMStatusStore
|
||||
auditWriter AuditWriter
|
||||
logger *slog.Logger
|
||||
maxPolls int
|
||||
pollDelay time.Duration
|
||||
}
|
||||
|
||||
func NewProxmoxTaskPollHandler(
|
||||
resolver ProxmoxTaskClientResolver,
|
||||
vmStore VMStatusStore,
|
||||
auditWriter AuditWriter,
|
||||
logger *slog.Logger,
|
||||
) ProxmoxTaskPollHandler {
|
||||
return ProxmoxTaskPollHandler{
|
||||
resolver: resolver,
|
||||
vmStore: vmStore,
|
||||
auditWriter: auditWriter,
|
||||
logger: logger,
|
||||
maxPolls: 60,
|
||||
pollDelay: 2 * time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
func (h ProxmoxTaskPollHandler) WithPolling(maxPolls int, pollDelay time.Duration) ProxmoxTaskPollHandler {
|
||||
h.maxPolls = maxPolls
|
||||
h.pollDelay = pollDelay
|
||||
return h
|
||||
}
|
||||
|
||||
func NewProxmoxTaskPollTask(payload ProxmoxTaskPollPayload) (*asynq.Task, error) {
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return asynq.NewTask(TypeProxmoxTaskPoll, body), nil
|
||||
}
|
||||
|
||||
func (h ProxmoxTaskPollHandler) ProcessTask(ctx context.Context, task *asynq.Task) error {
|
||||
payload, err := decodeProxmoxTaskPollPayload(task.Payload())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
client, err := h.resolver.ResolveTaskClient(ctx, payload.ClusterID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("resolve proxmox task client: %w", err)
|
||||
}
|
||||
|
||||
for attempt := 0; attempt < h.maxPolls; attempt++ {
|
||||
status, err := client.GetTaskStatus(ctx, payload.Node, payload.UPID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("poll proxmox task status: %w", err)
|
||||
}
|
||||
|
||||
if isTaskRunning(status) {
|
||||
if err := sleepContext(ctx, h.pollDelay); err != nil {
|
||||
return err
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if isTaskSuccessful(status) {
|
||||
return h.completeTarget(ctx, payload, payload.SuccessStatus, status)
|
||||
}
|
||||
|
||||
return h.completeTarget(ctx, payload, "failed", status)
|
||||
}
|
||||
|
||||
return fmt.Errorf("proxmox task %s still running after %d polls", payload.UPID, h.maxPolls)
|
||||
}
|
||||
|
||||
func (h ProxmoxTaskPollHandler) completeTarget(ctx context.Context, payload ProxmoxTaskPollPayload, finalStatus string, taskStatus ProxmoxTaskStatus) error {
|
||||
if payload.TargetType != "vm" {
|
||||
return fmt.Errorf("unsupported target type %q", payload.TargetType)
|
||||
}
|
||||
|
||||
if err := h.vmStore.SetVMStatus(ctx, payload.TargetID, finalStatus); err != nil {
|
||||
return fmt.Errorf("set vm status: %w", err)
|
||||
}
|
||||
|
||||
if err := h.auditWriter.WriteAudit(ctx, AuditEvent{
|
||||
TenantID: payload.TenantID,
|
||||
ProfileID: payload.ProfileID,
|
||||
Action: payload.Action,
|
||||
TargetType: payload.TargetType,
|
||||
TargetID: payload.TargetID,
|
||||
Metadata: map[string]any{
|
||||
"cluster_id": payload.ClusterID,
|
||||
"node": payload.Node,
|
||||
"upid": payload.UPID,
|
||||
"task_status": taskStatus.Status,
|
||||
"exit_status": taskStatus.ExitStatus,
|
||||
"target_status": finalStatus,
|
||||
"success_status": payload.SuccessStatus,
|
||||
},
|
||||
}); err != nil {
|
||||
return fmt.Errorf("write audit event: %w", err)
|
||||
}
|
||||
|
||||
h.logger.Info("completed proxmox task poll", "target_type", payload.TargetType, "target_id", payload.TargetID, "status", finalStatus)
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodeProxmoxTaskPollPayload(body []byte) (ProxmoxTaskPollPayload, error) {
|
||||
var payload ProxmoxTaskPollPayload
|
||||
if err := json.Unmarshal(body, &payload); err != nil {
|
||||
return ProxmoxTaskPollPayload{}, fmt.Errorf("decode proxmox task poll payload: %w", err)
|
||||
}
|
||||
if payload.ClusterID == "" {
|
||||
return ProxmoxTaskPollPayload{}, fmt.Errorf("cluster_id is required")
|
||||
}
|
||||
if payload.Node == "" {
|
||||
return ProxmoxTaskPollPayload{}, fmt.Errorf("node is required")
|
||||
}
|
||||
if payload.UPID == "" {
|
||||
return ProxmoxTaskPollPayload{}, fmt.Errorf("upid is required")
|
||||
}
|
||||
if payload.TargetType == "" {
|
||||
return ProxmoxTaskPollPayload{}, fmt.Errorf("target_type is required")
|
||||
}
|
||||
if payload.TargetID == "" {
|
||||
return ProxmoxTaskPollPayload{}, fmt.Errorf("target_id is required")
|
||||
}
|
||||
if payload.TenantID == "" {
|
||||
return ProxmoxTaskPollPayload{}, fmt.Errorf("tenant_id is required")
|
||||
}
|
||||
if payload.Action == "" {
|
||||
return ProxmoxTaskPollPayload{}, fmt.Errorf("action is required")
|
||||
}
|
||||
if payload.SuccessStatus == "" {
|
||||
return ProxmoxTaskPollPayload{}, fmt.Errorf("success_status is required")
|
||||
}
|
||||
return payload, nil
|
||||
}
|
||||
|
||||
func isTaskRunning(status ProxmoxTaskStatus) bool {
|
||||
return status.Status == "" || status.Status == "running"
|
||||
}
|
||||
|
||||
func isTaskSuccessful(status ProxmoxTaskStatus) bool {
|
||||
return status.Status == "stopped" && status.ExitStatus == "OK"
|
||||
}
|
||||
|
||||
func sleepContext(ctx context.Context, delay time.Duration) error {
|
||||
if delay <= 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
timer := time.NewTimer(delay)
|
||||
defer timer.Stop()
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-timer.C:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
type SQLVMStatusStore struct {
|
||||
db *sql.DB
|
||||
}
|
||||
|
||||
func NewSQLVMStatusStore(db *sql.DB) SQLVMStatusStore {
|
||||
return SQLVMStatusStore{db: db}
|
||||
}
|
||||
|
||||
func (s SQLVMStatusStore) SetVMStatus(ctx context.Context, vmID string, status string) error {
|
||||
result, err := s.db.ExecContext(ctx, `
|
||||
update public.vms
|
||||
set status = $2,
|
||||
updated_at = now()
|
||||
where id = $1
|
||||
`, vmID, status)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
rows, err := result.RowsAffected()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if rows == 0 {
|
||||
return fmt.Errorf("vm not found")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type SQLAuditWriter struct {
|
||||
db *sql.DB
|
||||
}
|
||||
|
||||
func NewSQLAuditWriter(db *sql.DB) SQLAuditWriter {
|
||||
return SQLAuditWriter{db: db}
|
||||
}
|
||||
|
||||
func (w SQLAuditWriter) WriteAudit(ctx context.Context, event AuditEvent) error {
|
||||
metadata, err := json.Marshal(event.Metadata)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
var profileID any
|
||||
if event.ProfileID != "" {
|
||||
profileID = event.ProfileID
|
||||
}
|
||||
|
||||
_, err = w.db.ExecContext(ctx, `
|
||||
insert into public.audit_log (
|
||||
tenant_id,
|
||||
profile_id,
|
||||
action,
|
||||
target_type,
|
||||
target_id,
|
||||
metadata
|
||||
)
|
||||
values ($1, $2, $3, $4, $5, $6::jsonb)
|
||||
`, event.TenantID, profileID, event.Action, event.TargetType, event.TargetID, string(metadata))
|
||||
return err
|
||||
}
|
||||
|
||||
type UnconfiguredProxmoxTaskClientResolver struct{}
|
||||
|
||||
func (UnconfiguredProxmoxTaskClientResolver) ResolveTaskClient(context.Context, string) (ProxmoxTaskClient, error) {
|
||||
return nil, fmt.Errorf("proxmox task client resolver is not configured")
|
||||
}
|
||||
Reference in New Issue
Block a user