apps/app-runner/server.go

package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"

	"golang.org/x/crypto/ssh"
)

type CommandState string

type CommandStatus struct {
	Command string       `json:"command"`
	State   CommandState `json:"state"`
}

type Status struct {
	Commit   *Commit         `json:"commit"`
	Commands []CommandStatus `json:"commands"`
}

type Server struct {
	l           sync.Locker
	port        int
	appId       string
	service     string
	id          string
	ready       bool
	cmd         *exec.Cmd
	repoAddr    string
	branch      string
	rootDir     string
	signer      ssh.Signer
	appDir      string
	runCommands []Command
	self        string
	managerAddr string
	logs        *Log
	currDir     string
	status      *Status
}

func NewServer(port int, appId, service, id, repoAddr, branch, rootDir string, signer ssh.Signer, appDir string, runCommands []Command, self string, manager string) *Server {
	return &Server{
		l:           &sync.Mutex{},
		port:        port,
		ready:       false,
		appId:       appId,
		service:     service,
		id:          id,
		repoAddr:    repoAddr,
		branch:      branch,
		rootDir:     rootDir,
		signer:      signer,
		appDir:      appDir,
		runCommands: runCommands,
		self:        self,
		managerAddr: manager,
		logs:        &Log{},
		currDir:     "",
		status:      nil,
	}
}

func (s *Server) Start() error {
	http.HandleFunc("/update", s.handleUpdate)
	http.HandleFunc("/ready", s.handleReady)
	http.HandleFunc("/logs", s.handleLogs)
	if s.managerAddr != "" && s.appId != "" {
		go s.pingManager()
	}
	if err := s.run(); err != nil {
		return err
	}
	return http.ListenAndServe(fmt.Sprintf(":%d", s.port), nil)
}

func (s *Server) handleLogs(w http.ResponseWriter, r *http.Request) {
	fmt.Fprint(w, s.logs.Contents())
}

func (s *Server) handleReady(w http.ResponseWriter, r *http.Request) {
	s.l.Lock()
	defer s.l.Unlock()
	if s.ready {
		fmt.Fprintln(w, "ok")
	} else {
		http.Error(w, "not ready", http.StatusInternalServerError)
	}
}

func (s *Server) handleUpdate(w http.ResponseWriter, r *http.Request) {
	fmt.Println("update")
	s.l.Lock()
	s.ready = false
	s.l.Unlock()
	if err := s.run(); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	s.l.Lock()
	s.ready = true
	s.l.Unlock()
}

func (s *Server) run() error {
	newDir, err := os.MkdirTemp(s.appDir, "code-*")
	if err != nil {
		return err
	}
	commit, err := CloneRepositoryBranch(s.repoAddr, s.branch, s.rootDir, s.signer, newDir)
	if err != nil {
		s.status = &Status{
			Commit: nil,
		}
		return err
	}
	s.status = &Status{
		Commit:   commit,
		Commands: []CommandStatus{},
	}
	commands := []string{}
	for _, c := range s.runCommands {
		args := []string{c.Bin}
		args = append(args, c.Args...)
		cmd := strings.Join(args, " ")
		commands = append(commands, cmd)
		s.status.Commands = append(s.status.Commands, CommandStatus{
			Command: cmd,
			State:   "waiting",
		})
	}
	logM := io.MultiWriter(os.Stdout, s.logs)
	for i, c := range commands {
		if i > 0 {
			s.status.Commands[i-1].State = "success"
		}
		cmd := &exec.Cmd{
			Dir:    filepath.Join(newDir, s.rootDir),
			Path:   "/bin/sh",
			Args:   []string{"/bin/sh", "-c", c},
			Env:    append(os.Environ(), s.runCommands[i].Env...),
			Stdout: logM,
			Stderr: logM,
		}
		cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
		fmt.Printf("Running: %s\n", c)
		s.status.Commands[i].State = "running"
		if i < len(s.runCommands)-1 {
			if err := cmd.Run(); err != nil {
				return err
			}
		} else {
			if s.cmd != nil {
				// TODO(gio): make this point configurable
				if err := s.kill(); err != nil {
					return err
				}
				if s.currDir != "" {
					if err := os.RemoveAll(s.currDir); err != nil {
						return err
					}
				}
			}
			if err := cmd.Start(); err != nil {
				return err
			}
			s.cmd = cmd
		}
	}
	s.currDir = newDir
	return nil
}

type pingReq struct {
	Id      string  `json:"id"`
	Service string  `json:"service"`
	Address string  `json:"address"`
	Status  *Status `json:"status,omitempty"`
	Logs    string  `json:"logs"`
}

func (s *Server) pingManager() {
	defer func() {
		go func() {
			time.Sleep(5 * time.Second)
			s.pingManager()
		}()
	}()
	buf, err := json.Marshal(pingReq{
		Id:      s.id,
		Service: s.service,
		Address: fmt.Sprintf("http://%s:%d", s.self, s.port),
		Status:  s.status,
		Logs:    s.logs.Contents(),
	})
	if err != nil {
		return
	}
	registerWorkerAddr := fmt.Sprintf("%s/api/project/%s/workers", s.managerAddr, s.appId)
	resp, err := http.Post(registerWorkerAddr, "application/json", bytes.NewReader(buf))
	if err != nil {
		fmt.Println(err)
	} else {
		// check resp code
		io.Copy(os.Stdout, resp.Body)
	}
}

func (s *Server) kill() error {
	if s.cmd == nil {
		return nil
	}

	err := syscall.Kill(-s.cmd.Process.Pid, syscall.SIGKILL)
	if err != nil {
		return err
	}
	// NOTE(gio): No need to check err as we just killed the process
	s.cmd.Wait()
	s.cmd = nil
	return nil
}