overflow_probes/service/ContainerService.go

package service

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"reflect"
	"strconv"
	"time"

	cda "git.loafle.net/commons_go/di/annotation"
	cdr "git.loafle.net/commons_go/di/registry"
	"git.loafle.net/commons_go/logging"
	crc "git.loafle.net/commons_go/rpc/client"
	oopcc "git.loafle.net/overflow/overflow_probe_container/client"
	"git.loafle.net/overflow/overflow_probes/config"
)

func init() {
	cdr.RegisterType(reflect.TypeOf((*ContainerService)(nil)), &cda.ComponentAnnotation{})
}

type ContainerService struct {
	clients map[string]*containerState
}

type containerState struct {
	pid    int
	port   int
	client crc.Client
}

func (cs *ContainerService) Call(name string, result interface{}, method string, params ...interface{}) error {
	c, err := cs.GetClient(name)
	if nil != err {
		return err
	}
	return c.Call(result, method, params...)
}

func (cs *ContainerService) Send(name string, method string, params ...interface{}) error {
	c, err := cs.GetClient(name)
	if nil != err {
		return err
	}
	return c.Send(method, params...)
}

func (cs *ContainerService) GetClient(name string) (crc.Client, error) {
	cState, ok := cs.clients[name]
	if !ok {
		if err := cs.runProcess(name); nil != err {
			return nil, err
		}
		cState, _ = cs.clients[name]
	}

	return cState.client, nil
}

func (cs *ContainerService) Connect(name string) error {
	return nil
}

func (cs *ContainerService) runProcess(name string) error {
	ok := cs.checkProcess(name)

	if ok {
		return nil
	}

	cmdPath := path.Join(config.Config.Paths["root"], config.Config.Paths["bin"], name+".sh")
	pidPath := path.Join(config.Config.Paths["root"], config.Config.Paths["pid"], name+".pid")

	cmd := exec.Command(cmdPath, pidPath)
	if err := cmd.Start(); nil != err {
		logging.Logger().Error(fmt.Sprintf("Probe: To run container(%s) failed err %v", name, err))
		return err
	}

	port, err := watchPidFileCreate(pidPath, time.Duration(time.Second*2))
	if nil != err {
		return err
	}

	clientAddr := fmt.Sprintf("localhost:%d", port)
	client := oopcc.New(clientAddr, nil)
	if err := client.Connect(); nil != err {
		cmd.Process.Kill()
		return err
	}

	cState := &containerState{
		pid:    cmd.Process.Pid,
		port:   port,
		client: client,
	}

	cs.clients[name] = cState

	return nil
}

func (cs *ContainerService) checkProcess(name string) bool {
	cState, ok := cs.clients[name]
	if !ok || nil == cState || nil == cState.client {
		return false
	}

	_, err := os.FindProcess(cState.pid)
	if nil != err {
		return false
	}

	return true
}

func (cs *ContainerService) killProcess(name string) error {
	cState, ok := cs.clients[name]
	if !ok || nil == cState || nil == cState.client {
		return fmt.Errorf("Probe: Container[%s] is not exist", name)
	}

	p, err := os.FindProcess(cState.pid)
	if nil != err {
		return err
	}

	if err = p.Kill(); nil != err {
		return err
	}

	return nil
}

func watchPidFileCreate(pidFilePath string, waitTime time.Duration) (int, error) {
	startTime := time.Now()

	for {
		if _, err := os.Stat(pidFilePath); err == nil {
			buf, err := ioutil.ReadFile(pidFilePath)
			if nil != err {
				return 0, err
			}
			portNumber, err := strconv.ParseInt(string(buf), 10, 32)
			if nil != err {
				return 0, err
			}
			return int(portNumber), nil
		}

		if time.Since(startTime) > waitTime {
			return 0, fmt.Errorf("Probe: pid file not exist")
		}

		time.Sleep(time.Duration(time.Millisecond * 100))
	}
}