crasa-transport/webtest/server.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"sync"
	"syscall"
	"time"

	"github.com/a-h/templ"
	"github.com/go-chi/chi/v5"
	"github.com/go-chi/chi/v5/middleware"
	"github.com/gorilla/websocket"
)

func run() error {
	log.Printf("startup orchestrator pid=%d", os.Getpid())

	root, err := findRepoRoot()
	if err != nil {
		return err
	}

	h := newHub()
	upgrader := websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool { return true },
	}

	r := chi.NewRouter()
	r.Use(middleware.RealIP)
	r.Use(middleware.RequestID)
	r.Use(middleware.Recoverer)

	r.Get("/", templ.Handler(indexPage()).ServeHTTP)

	r.Get("/ws", func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			log.Printf("ws upgrade failed remote=%s err=%v", r.RemoteAddr, err)
			return
		}
		h.add(conn)
		defer func() {
			h.remove(conn)
			_ = conn.Close()
		}()
		for {
			if _, _, err := conn.ReadMessage(); err != nil {
				return
			}
		}
	})

	staticServer := http.FileServer(http.FS(staticFiles))
	r.Handle("/static/*", http.StripPrefix("/", staticServer))

	httpServer := &http.Server{
		Addr:              "127.0.0.1:5050",
		Handler:           r,
		ReadHeaderTimeout: 5 * time.Second,
	}

	go func() {
		log.Printf("http server listen addr=http://127.0.0.1:5050")
		if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			log.Fatalf("http server failed: %v", err)
		}
	}()

	if err := runBuildInstall(root, h); err != nil {
		return err
	}

	serverBin, err := mustBinPath(root, "test-server")
	if err != nil {
		return err
	}
	clientBin, err := mustBinPath(root, "test-client")
	if err != nil {
		return err
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	serverCmd, err := launchProcess(ctx, root, "server", serverBin, []string{"--port", "5000"}, h)
	if err != nil {
		return err
	}

	var (
		clientMu     sync.Mutex
		nextClientID = 1
		clients      = map[int]*exec.Cmd{}
	)

	launchClient := func() (int, error) {
		clientMu.Lock()
		id := nextClientID
		nextClientID += 1
		clientMu.Unlock()

		name := fmt.Sprintf("client-%d", id)
		cmd, err := launchProcess(ctx, root, name, clientBin, []string{"--port=5000", "--host=localhost"}, h)
		if err != nil {
			return 0, err
		}

		clientMu.Lock()
		clients[id] = cmd
		clientMu.Unlock()

		go func(clientID int, clientName string, clientCmd *exec.Cmd) {
			err := clientCmd.Wait()
			clientMu.Lock()
			delete(clients, clientID)
			clientMu.Unlock()
			if err != nil {
				h.status("orchestrator", fmt.Sprintf("%s exited with error: %v", clientName, err))
				return
			}
			h.status("orchestrator", fmt.Sprintf("%s exited cleanly", clientName))
		}(id, name, cmd)

		return id, nil
	}

	stopClient := func(id int) bool {
		clientMu.Lock()
		cmd, ok := clients[id]
		if ok {
			delete(clients, id)
		}
		clientMu.Unlock()
		if !ok {
			return false
		}
		stopProcess(cmd)
		return true
	}

	stopAllClients := func() {
		clientMu.Lock()
		copied := make([]*exec.Cmd, 0, len(clients))
		for _, cmd := range clients {
			copied = append(copied, cmd)
		}
		clients = map[int]*exec.Cmd{}
		clientMu.Unlock()
		for _, cmd := range copied {
			stopProcess(cmd)
		}
	}

	r.Post("/api/clients", func(w http.ResponseWriter, r *http.Request) {
		id, err := launchClient()
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		name := fmt.Sprintf("client-%d", id)
		h.status("orchestrator", fmt.Sprintf("spawned %s", name))
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(map[string]any{
			"id":   id,
			"name": name,
		})
	})

	r.Delete("/api/clients/{id}", func(w http.ResponseWriter, r *http.Request) {
		idText := chi.URLParam(r, "id")
		id, err := strconv.Atoi(idText)
		if err != nil {
			http.Error(w, "invalid client id", http.StatusBadRequest)
			return
		}
		if !stopClient(id) {
			http.Error(w, "client not found", http.StatusNotFound)
			return
		}
		h.status("orchestrator", fmt.Sprintf("stopped client-%d", id))
		w.WriteHeader(http.StatusNoContent)
	})

	if _, err := launchClient(); err != nil {
		stopProcess(serverCmd)
		return err
	}

	h.status("orchestrator", "server and initial client launched")

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)

	waitCh := make(chan string, 1)
	go func() {
		if err := serverCmd.Wait(); err != nil {
			waitCh <- "server exited with error: " + err.Error()
			return
		}
		waitCh <- "server exited cleanly"
	}()

	for {
		select {
		case sig := <-sigCh:
			h.status("orchestrator", "signal: "+sig.String())
			goto shutdown
		case line := <-waitCh:
			h.status("orchestrator", line)
			h.status("orchestrator", "server exited; orchestrator staying alive until interrupted")
		}
	}

shutdown:
	cancel()
	stopProcess(serverCmd)
	stopAllClients()
	h.status("orchestrator", "shutdown complete")

	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer shutdownCancel()
	_ = httpServer.Shutdown(shutdownCtx)
	return nil
}