houston/client/operation.go

package client

import (
	"archive/zip"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"go-ayo/logger"
	"houston/shared"
	"houston/shared/protos"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync/atomic"
	"syscall"
	"time"
)

type parsedVersionString = []string

func parseVersionString(ver string) parsedVersionString {
	return strings.Split(ver, ".")
}

func lastExecutionArgs(verpath string) []string {
	argf, err := os.Open(path.Join(verpath, "@args"))
	if os.IsNotExist(err) {
		argf, err = os.Open(path.Clean(path.Join(verpath, "..", "@args")))
		if os.IsNotExist(err) {
			return nil
		}
	}
	defer argf.Close()

	var out []string
	dec := json.NewDecoder(argf)
	dec.Decode(&out)
	return out
}

func compareVersionString(lhs, rhs parsedVersionString) int {
	minlen := len(lhs)
	if minlen > len(rhs) {
		minlen = len(rhs)
	}

	for i := 0; i < minlen; i++ {
		if len(lhs[i]) < len(rhs[i]) {
			return -1
		}

		if len(lhs[i]) > len(rhs[i]) {
			return 1
		}

		if lhs[i] < rhs[i] {
			return -1
		}

		if lhs[i] > rhs[i] {
			return 1
		}
	}

	return len(lhs) - len(rhs)
}

func findLastestVersion(root string) (string, error) {
	// 최신 버전을 찾음
	entries, err := os.ReadDir(root)
	if err != nil {
		return "", err
	}
	if len(entries) == 0 {
		return "", nil
	}
	latest := parseVersionString(entries[0].Name())
	for i := 1; i < len(entries); i++ {
		next := parseVersionString(entries[i].Name())
		if compareVersionString(latest, next) < 0 {
			latest = next
		}
	}
	return strings.Join(latest, "."), nil
}

func (meta *procmeta) launch(args []string, exitChan chan<- *exec.Cmd) error {
	exepath := args[0]
	verpath := path.Dir(exepath)
	args[0] = path.Base(exepath)

	cmd := exec.Command("./"+args[0], args[1:]...)
	cmd.Dir = verpath

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return err
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return err
	}

	err = os.MkdirAll(path.Join(cmd.Dir, "logs"), os.ModePerm)
	if err != nil {
		return err
	}

	now := time.Now().UTC()
	ext := path.Ext(cmd.Args[0])
	nameonly := path.Base(cmd.Args[0])
	if len(ext) > 0 {
		nameonly = nameonly[:len(nameonly)-len(ext)]
	}
	ts := now.Format("2006-01-02T15-04-05")
	stdPrefix := path.Join(cmd.Dir, "logs", fmt.Sprintf("%s_%s", nameonly, ts))
	errfile, err := os.Create(stdPrefix + ".stderr.log")
	if err != nil {
		return err
	}
	outfile, err := os.Create(stdPrefix + ".stdout.log")
	if err != nil {
		return err
	}

	go func() {
		defer func() {
			recover()
			stdout.Close()
			errfile.Close()
		}()

		buff := make([]byte, 1024)
		for {
			size, err := stderr.Read(buff)
			if err != nil {
				exitChan <- cmd
				errfile.Close()
				break
			}
			errfile.Write(buff[:size])
			new := atomic.AddInt32(&meta.stderrSize, int32(size))
			logger.Println("stderrSize :", new)
		}
	}()

	go func() {
		defer func() {
			recover()
			stderr.Close()
			outfile.Close()
		}()

		buff := make([]byte, 1024)
		for {
			size, err := stdout.Read(buff)
			if err != nil {
				exitChan <- cmd
				break
			}
			outfile.Write(buff[:size])
			new := atomic.AddInt32(&meta.stdoutSize, int32(size))
			logger.Println("stdoutSize :", new)
		}
	}()

	err = cmd.Start()
	if err != nil {
		return err
	}

	meta.cmd = cmd
	meta.stdin = stdin
	meta.stdPrefix = stdPrefix
	meta.state = protos.ProcessState_Running

	return nil
}

func (hc *houstonClient) startChildProcess(req *shared.StartProcessRequest) error {
	if req.Version == "latest" {
		// 최신 버전을 찾음
		latest, err := findLastestVersion(path.Join("./", req.Name))
		if err != nil {
			return err
		}

		req.Version = latest
	}

	meta := &procmeta{
		name:    req.Name,
		version: req.Version,
		state:   protos.ProcessState_Error,
	}

	verpath := path.Join("./", req.Name, req.Version)
	fi, err := os.Stat(verpath)
	if err == nil && fi.IsDir() {
		// Define regular expression
		re := regexp.MustCompile(`[^\s"']+|"([^"]*)"|'([^']*)`)

		// Split input string into array of strings
		result := re.FindAllString(req.Args, -1)
		for i := range result {
			result[i] = strings.Trim(result[i], "\"'")
		}
		result[0] = path.Join(verpath, result[0])
		err := meta.launch(result, hc.exitChan)
		if err != nil {
			return err
		}

		// launch가 성공하면 args 저장. this and parent folder
		if argfile, err := os.Create(path.Join(req.Name, "@args")); err == nil {
			enc := json.NewEncoder(argfile)
			enc.Encode(result)
			argfile.Close()
		}
		if argfile, err := os.Create(path.Join(verpath, "@args")); err == nil {
			enc := json.NewEncoder(argfile)
			enc.Encode(result)
			argfile.Close()
		}

		hc.childProcs = append(hc.childProcs, meta)

		op := protos.NewOperationClient(hc.client)
		op.Refresh(context.Background(), hc.makeOperationQueryRequest())

		return nil
	}

	return err
}

var errNoRunningProcess = errors.New("no running processed")

func (hc *houstonClient) stopChildProcess(req *shared.StopProcessRequest) error {
	if req.Version == "latest" {
		// 최신 버전을 찾음
		latest, err := findLastestVersion(path.Join("./", req.Name))
		if err != nil {
			return err
		}

		req.Version = latest
	}

	var remains []*procmeta
	var killing []*procmeta
	for _, proc := range hc.childProcs {
		if proc.state != protos.ProcessState_Running {
			continue
		}

		if req.Pid != 0 {
			if req.Pid == int32(proc.cmd.Process.Pid) {
				// 해당 pid만 제거
				killing = append(killing, proc)
			} else {
				remains = append(remains, proc)
			}
		} else if proc.name == req.Name {
			if len(req.Version) == 0 {
				// program 다 정지
				killing = append(killing, proc)
			} else if req.Version == proc.version {
				// program의 특정 버전만 정지
				killing = append(killing, proc)
			} else {
				// 해당 사항 없음
				remains = append(remains, proc)
			}
		} else {
			// 해당 사항 없음
			remains = append(remains, proc)
		}
	}

	if len(killing) > 0 {
		for _, proc := range killing {
			if err := proc.cmd.Process.Signal(syscall.SIGTERM); err != nil {
				proc.cmd.Process.Signal(os.Kill)
				proc.state = protos.ProcessState_Stopping
			}
		}

		op := protos.NewOperationClient(hc.client)
		op.Refresh(context.Background(), hc.makeOperationQueryRequest())

		for _, proc := range killing {
			proc.cmd.Wait()
		}

		hc.childProcs = remains

		op.Refresh(context.Background(), hc.makeOperationQueryRequest())

		return nil
	}

	return errNoRunningProcess
}

func (hc *houstonClient) restartChildProcess(req *shared.RestartProcessRequest) error {
	if req.Version == "latest" {
		// 최신 버전을 찾음
		latest, err := findLastestVersion(path.Join("./", req.Name))
		if err != nil {
			return err
		}

		req.Version = latest
	}

	var restarts []*procmeta
	for _, proc := range hc.childProcs {
		if proc.name == req.Name {
			if len(req.Version) == 0 {
				restarts = append(restarts, proc)
			} else if req.Version == proc.version {
				restarts = append(restarts, proc)
			}
		}
	}

	if len(restarts) == 0 {
		return errNoRunningProcess
	}

	for _, proc := range restarts {
		if err := proc.cmd.Process.Signal(syscall.SIGTERM); err != nil {
			proc.cmd.Process.Signal(os.Kill)
		}
		proc.state = protos.ProcessState_Stopping
	}

	op := protos.NewOperationClient(hc.client)
	op.Refresh(context.Background(), hc.makeOperationQueryRequest())

	for _, proc := range restarts {
		proc.cmd.Wait()
		proc.state = protos.ProcessState_Stopped
	}
	op.Refresh(context.Background(), hc.makeOperationQueryRequest())

	for _, proc := range restarts {
		args := proc.cmd.Args
		args[0] = path.Join(proc.cmd.Dir, args[0])

		if err := proc.launch(args, hc.exitChan); err != nil {
			return err
		}
	}
	op.Refresh(context.Background(), hc.makeOperationQueryRequest())

	return nil
}

func (hc *houstonClient) uploadFiles(req *shared.UploadRequest) error {
	if req.Version == "latest" {
		// 최신 버전을 찾음
		latest, err := findLastestVersion(path.Join("./", req.Name))
		if err != nil {
			return err
		}

		req.Version = latest
	}

	root := path.Join(req.Name, req.Version)
	matches, err := filepath.Glob(path.Join(root, req.Filter))
	if err != nil {
		return err
	}

	if len(matches) == 0 {
		resp, err := http.Post(req.Url, "application/zip", bytes.NewBuffer([]byte{}))
		if err != nil {
			return err
		}
		resp.Body.Close()
		return nil
	}

	// Create a file to write the archive to.
	f, err := os.CreateTemp("", "")
	if err != nil {
		return err
	}
	go func(f *os.File) {
		defer func() {
			tempname := f.Name()
			f.Close()

			resp, _ := http.Post(req.Url, "application/zip", f)
			if resp != nil && resp.Body != nil {
				resp.Body.Close()
			}
			os.Remove(tempname)
			if del, err := strconv.ParseBool(req.DeleteAfterUploaded); del && err == nil {
				for _, f := range matches {
					os.Remove(f)
				}
			}
		}()

		// Create a new zip archive.
		w := zip.NewWriter(f)
		defer w.Close()

		for _, file := range matches {
			relative := file[len(root)+1:]
			fw, err := w.Create(relative)
			if err != nil {
				logger.Error(err)
				return
			}

			src, err := os.Open(file)
			if err != nil {
				logger.Error(err)
				return
			}
			defer src.Close()

			if _, err = io.Copy(fw, src); err != nil {
				logger.Error(err)
				return
			}

		}
	}(f)

	return nil
}