package api
import (
"fmt"
"html/template"
"log"
"net/http"
"nvidiadashboard/pkg/nvidia"
"os"
"path/filepath"
"strings"
"github.com/go-chi/chi/v5"
"github.com/go-chi/chi/v5/middleware"
"github.com/inhies/go-bytesize"
"golang.org/x/net/http2"
"golang.org/x/net/http2/h2c"
)
type (
NVIDIARepository interface {
GetGPU(ID int) (nvidia.GPUDetail, error)
GetGPUs() ([]nvidia.GPU, error)
DriverVersion() string
CUDAVersion() string
}
Server struct {
repo NVIDIARepository
mux *chi.Mux
errLog *log.Logger
}
WebPack struct {
GPUs []nvidia.GPU
GPU nvidia.GPUDetail
Username string
DriverVersion string
CUDAVersion string
}
)
var (
templateFuncMap = template.FuncMap{
"ConvertByteSize": convertByteSize,
"PercentageRounded": percentageRounded,
}
)
const internalServerErrorPage string = `
Internal Server Error
Internal Server Error
%v
nvidia-smi dashboard
`
func New(repo NVIDIARepository) *Server {
s := &Server{
mux: chi.NewRouter(),
repo: repo,
errLog: log.New(os.Stderr, log.Prefix(), log.Flags()),
}
s.mux.Use(middleware.RequestID)
s.mux.Use(middleware.Logger)
s.mux.Use(middleware.Recoverer)
s.mux.Use(middleware.URLFormat)
s.mux.Get("/", s.handleRoot)
s.mux.Get("/{uuid:GPU-(.*)}", s.handleGPU)
s.mux.Get("/{path:(.*).html}", func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusNotFound)
w.Write([]byte("404 page not found"))
})
workDir, _ := os.Getwd()
filesDir := http.Dir(filepath.Join(workDir, "static"))
fileServer(s.mux, "/", filesDir)
return s
}
func (s *Server) Serve(port uint) error {
h2s := &http2.Server{}
srv := &http.Server{
Addr: fmt.Sprintf("0.0.0.0:%v", port),
Handler: h2c.NewHandler(s.mux, h2s), // HTTP/2 Cleartext handler
}
return srv.ListenAndServe()
}
func (s *Server) handleRoot(w http.ResponseWriter, _ *http.Request) {
gpus, err := s.repo.GetGPUs()
if err != nil {
sendInternalServerErrorHTML(w, err)
return
}
if len(gpus) > 0 {
w.Header().Add("Location", "/"+gpus[0].UUID)
w.WriteHeader(http.StatusTemporaryRedirect)
return
}
wp := WebPack{
Username: "anonymous",
}
t, err := template.ParseFiles("static/no_gpu.html")
if err != nil {
sendInternalServerErrorHTML(w, err)
}
err = t.Execute(w, wp)
if err != nil {
s.errLog.Println("[ERROR]", err)
}
}
func (s *Server) handleGPU(w http.ResponseWriter, r *http.Request) {
uuid := chi.URLParam(r, "uuid")
gpus, err := s.repo.GetGPUs()
if err != nil {
sendInternalServerErrorHTML(w, err)
return
}
i := -1
for _, gpu := range gpus {
if gpu.UUID == uuid {
i = gpu.Index
}
}
if i == -1 {
w.WriteHeader(http.StatusNotFound)
w.Write([]byte("404 page not found"))
}
gpu, err := s.repo.GetGPU(i)
if err != nil {
sendInternalServerErrorHTML(w, err)
return
}
wp := WebPack{
Username: "anonymous",
GPUs: gpus,
GPU: gpu,
DriverVersion: s.repo.DriverVersion(),
CUDAVersion: s.repo.CUDAVersion(),
}
t, err := template.New("index").Funcs(templateFuncMap).ParseFiles("static/index.html")
if err != nil {
sendInternalServerErrorHTML(w, err)
}
err = t.ExecuteTemplate(w, "index.html", wp)
if err != nil {
s.errLog.Println("[ERROR]", err)
}
}
// FileServer conveniently sets up a http.FileServer handler to serve
// static files from a http.FileSystem.
func fileServer(r chi.Router, path string, root http.FileSystem) {
if strings.ContainsAny(path, "{}*") {
panic("FileServer does not permit any URL parameters.")
}
if path != "/" && path[len(path)-1] != '/' {
r.Get(path, http.RedirectHandler(path+"/", http.StatusMovedPermanently).ServeHTTP)
path += "/"
}
path += "*"
r.Get(path, func(w http.ResponseWriter, r *http.Request) {
rctx := chi.RouteContext(r.Context())
pathPrefix := strings.TrimSuffix(rctx.RoutePattern(), "/*")
fs := http.StripPrefix(pathPrefix, http.FileServer(root))
fs.ServeHTTP(w, r)
})
}
func sendInternalServerErrorHTML(w http.ResponseWriter, v any) {
w.WriteHeader(http.StatusInternalServerError)
body := fmt.Sprintf(internalServerErrorPage, v)
w.Write([]byte(body))
}
func convertByteSize(v int) string {
bs := bytesize.New(float64(v))
return bs.String()
}
func percentageRounded(a, b int) int {
p := (a / b) * 100
return p
}