gpu-turnstile/cmd/gpu-turnstile/monitor.go

package main

import (
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"time"

	"gpu-turnstile/internal/control"
)

// ANSI colors for the monitor frame.
const (
	cReset  = "\x1b[0m"
	cDim    = "\x1b[2m"
	cRed    = "\x1b[31m"
	cGreen  = "\x1b[32m"
	cYellow = "\x1b[33m"
	cCyan   = "\x1b[36m"
)

// hotkeysLine is the monitor's footer.
const hotkeysLine = " " + cDim + "q quit · u update now" + cReset + "\x1b[K\n"

// monitorCommand renders a live status view of the running service,
// refreshed every second from the control channel. When the service
// reports a different version and the executable on disk changed (the
// updater replaced it), the monitor restarts itself onto the new binary.
// Hotkeys: q quits, u triggers an update check on the service.
func monitorCommand() int {
	if !stdoutIsTerminal() {
		fmt.Fprintln(os.Stderr, "gpu-turnstile: --monitor needs an interactive terminal")
		return 1
	}
	enableVirtualTerminal()
	restore := enableRawKeys()
	defer func() {
		if restore != nil {
			restore()
		}
	}()
	fmt.Print("\x1b[2J") // clear once; frames then redraw in place
	defer fmt.Print(cReset + "\n")
	exe, _ := os.Executable()
	var exeStamp time.Time
	if st, err := os.Stat(exe); err == nil {
		exeStamp = st.ModTime()
	}

	keys := make(chan byte, 8)
	go readKeys(keys)
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()

	var note string
	var noteAt time.Time
	noteCh := make(chan string, 1)
	updatePending := false

	poll := func() string {
		frame := renderWaiting()
		if reply, err := control.Ask(control.CmdStatus); err == nil {
			if msg, ok := strings.CutPrefix(reply, "OK "); ok {
				var snap statusSnapshot
				if json.Unmarshal([]byte(msg), &snap) == nil {
					if snap.Version != "" && snap.Version != version {
						if exeChanged(exe, exeStamp) {
							fmt.Print("\x1b[2J\x1b[H")
							fmt.Printf("gpu-turnstile: service updated to %s — restarting the monitor\n", snap.Version)
							restartSelf(exe, "--monitor")
							return "" // re-execed; this process exits below
						}
						snap.MonitorNote = fmt.Sprintf("note: the service runs %s, this monitor is %s", snap.Version, version)
					}
					if note != "" {
						snap.MonitorNote = note
					}
					frame = renderMonitor(snap, termWidth())
				}
			}
		}
		return frame
	}

	for {
		frame := poll()
		if frame == "" {
			return 0 // restartSelf fired
		}
		fmt.Print("\x1b[H" + frame + "\x1b[J") // home, frame, clear below
		select {
		case <-ticker.C:
			if note != "" && time.Since(noteAt) > 15*time.Second {
				note = ""
			}
		case k, ok := <-keys:
			if !ok {
				keys = nil
				continue
			}
			switch k {
			case 'q', 'Q', 3: // q or Ctrl+C (raw mode delivers it as a byte)
				return 0
			case 'u', 'U':
				if !updatePending {
					updatePending = true
					note, noteAt = "checking for updates…", time.Now()
					go func() {
						reply, err := control.Ask(control.CmdUpdateNow)
						if err != nil {
							noteCh <- "update: no answer from the service"
							return
						}
						msg := strings.TrimPrefix(reply, "OK ")
						msg = strings.TrimPrefix(msg, "ERR ")
						noteCh <- "update: " + msg
					}()
				}
			}
		case n := <-noteCh:
			updatePending = false
			note, noteAt = n, time.Now()
		}
	}
}

// readKeys reads single keypresses from stdin (raw mode was enabled by the
// caller) and delivers them until stdin fails.
func readKeys(keys chan<- byte) {
	defer close(keys)
	buf := make([]byte, 1)
	for {
		n, err := os.Stdin.Read(buf)
		if n > 0 {
			keys <- buf[0]
		}
		if err != nil {
			return
		}
	}
}

// exeChanged reports whether the executable on disk was replaced since the
// recorded stamp (the updater swaps it via rename, which changes ModTime).
func exeChanged(exe string, stamp time.Time) bool {
	if exe == "" || stamp.IsZero() {
		return false
	}
	st, err := os.Stat(exe)
	return err == nil && !st.ModTime().Equal(stamp)
}

// restartSelf starts a fresh copy of this executable with the given args on
// the same console; the caller exits right after.
func restartSelf(exe string, args ...string) {
	cmd := exec.Command(exe, args...)
	cmd.Stdin, cmd.Stdout, cmd.Stderr = os.Stdin, os.Stdout, os.Stderr
	cmd.Start() //nolint:errcheck // best effort: on failure we just exit
}

func renderWaiting() string {
	return cDim + " gpu-turnstile — waiting for a running service…" + cReset + "\x1b[K\n\x1b[K\n" + hotkeysLine
}

// renderMonitor draws one full frame. Each line ends with \x1b[K (clear to
// end of line) so shrinking content leaves no residue.
func renderMonitor(snap statusSnapshot, width int) string {
	if width < 40 {
		width = 80
	}
	var b strings.Builder

	left := " gpu-turnstile"
	if snap.UptimeS > 0 {
		left += "  " + cDim + "up " + fmtDur(snap.UptimeS) + cReset
	}
	right := snap.Version
	pad := width - printableLen(" gpu-turnstile  up "+fmtDur(snap.UptimeS)) - len(right) - 1
	if snap.UptimeS == 0 {
		pad = width - len(" gpu-turnstile") - len(right) - 1
	}
	if pad < 1 {
		pad = 1
	}
	b.WriteString(cDim + left + strings.Repeat(" ", pad) + right + cReset + "\x1b[K\n")
	b.WriteString(cDim + " " + strings.Repeat("─", width-2) + cReset + "\x1b[K\n")

	for _, d := range snap.Downstreams {
		b.WriteString(renderDownstream(d) + "\x1b[K\n")
	}
	b.WriteString("\x1b[K\n")
	b.WriteString(renderLock(snap.Lock) + "\x1b[K\n")
	if snap.Lock.ImageQueue > 0 {
		b.WriteString(fmt.Sprintf(" Queue: %s%d image job(s) waiting%s\x1b[K\n",
			cYellow, snap.Lock.ImageQueue, cReset))
	}
	if snap.GPU.Known {
		b.WriteString(" GPU: " + renderVRAM(snap.GPU.UsedMB, snap.GPU.TotalMB) + "\x1b[K\n")
	}
	if snap.MonitorNote != "" {
		b.WriteString(" " + cYellow + snap.MonitorNote + cReset + "\x1b[K\n")
	}
	b.WriteString("\x1b[K\n")
	b.WriteString(hotkeysLine)
	return b.String()
}

// renderVRAM renders "4.2 / 16.0 GiB used" (or MiB below 1 GiB).
func renderVRAM(used, total int) string {
	format := func(mb int) string {
		if mb >= 1024 {
			return fmt.Sprintf("%.1f GiB", float64(mb)/1024)
		}
		return fmt.Sprintf("%d MiB", mb)
	}
	if total > 0 {
		return format(used) + " / " + format(total) + " used"
	}
	return format(used) + " used"
}

// printableLen counts characters without ANSI escapes (ASCII-only content).
func printableLen(s string) int { return len(s) }

func renderDownstream(d statusDownstream) string {
	url := cDim + d.URL + cReset
	switch d.Managed {
	case "stopped":
		return fmt.Sprintf(" %s○%s %-8s %sstopped (managed — starts on demand)%s  %s",
			cDim, cReset, d.Name, cDim, cReset, url)
	case "starting":
		return fmt.Sprintf(" %s◌%s %-8s %sstarting…%s  %s",
			cYellow, cReset, d.Name, cYellow, cReset, url)
	}
	suffix := ""
	if d.Managed == "external" {
		suffix = " (external)"
	}
	if d.Up {
		return fmt.Sprintf(" %s●%s %-8s %sUP%s%s   %s", cGreen, cReset, d.Name, cGreen, cReset, suffix, url)
	}
	return fmt.Sprintf(" %s●%s %-8s %sDOWN%s %s", cRed, cReset, d.Name, cRed, cReset, url)
}

func renderLock(l statusLock) string {
	dur := cDim + "(" + fmtDur(l.SinceS) + ")" + cReset
	switch l.State {
	case "idle":
		return fmt.Sprintf(" Lock: %sidle%s %s", cGreen, cReset, dur)
	case "llm":
		s := fmt.Sprintf(" Lock: %sLLM%s — %d in flight", cCyan, cReset, l.LLMInflight)
		if l.LLMWaiting > 0 {
			s += fmt.Sprintf(", %d waiting", l.LLMWaiting)
		}
		if l.Detail != "" {
			s += " — " + l.Detail
		}
		return s + " " + dur
	case "image":
		s := fmt.Sprintf(" Lock: %sIMAGE%s", cYellow, cReset)
		if l.Detail != "" {
			s += " — " + l.Detail
		}
		return s + " " + dur
	case "external":
		return fmt.Sprintf(" Lock: %sEXTERNAL%s — %s %s", cRed, cReset, l.External, dur)
	}
	return " Lock: unknown"
}

// fmtDur renders seconds as a compact duration ("1m32s", "2h07m").
func fmtDur(s int64) string {
	if s < 0 {
		s = 0
	}
	d := time.Duration(s) * time.Second
	if d >= time.Hour {
		return fmt.Sprintf("%dh%02dm", int(d.Hours()), int(d.Minutes())%60)
	}
	return d.Round(time.Second).String()
}