github.com/coyove/common@v0.0.0-20240403014525-f70e643f9de8/coward/base.go (about)

     1  package coward
     2  
     3  import (
     4  	"flag"
     5  	"io"
     6  	"math/rand"
     7  	"os"
     8  	"os/exec"
     9  	"strconv"
    10  	"strings"
    11  	"time"
    12  
    13  	"github.com/mitchellh/go-ps"
    14  	log "github.com/sirupsen/logrus"
    15  	"gopkg.in/natefinch/lumberjack.v2"
    16  )
    17  
    18  var (
    19  	coward    = flag.Bool("c", false, "")
    20  	slaves    = flag.Int("slaves", 0, "")
    21  	Slave     = flag.Int("s", 0, "internal use")
    22  	slaveMark = flag.Int("sm", 0, "internal use")
    23  )
    24  
    25  func Init(name string) bool {
    26  	flag.Parse()
    27  
    28  	log.SetReportCaller(true)
    29  	log.SetOutput(io.MultiWriter(os.Stdout, &lumberjack.Logger{
    30  		Filename:   "/var/log/" + name + ".log",
    31  		MaxSize:    100, // megabytes
    32  		MaxBackups: 16,
    33  		MaxAge:     28,   //days
    34  		Compress:   true, // disabled by default
    35  	}))
    36  
    37  	p, _ := ps.Processes()
    38  	for _, p := range p {
    39  		if strings.Contains(p.Executable(), name) && os.Getpid() != p.Pid() {
    40  			if *coward {
    41  				log.Info("coward mode, existing server: ", p.Pid(), ", exit quietly")
    42  				return false
    43  			}
    44  			if *Slave == 0 {
    45  				log.Info("terminate old server: ", p.Pid(), exec.Command("kill", "-9", strconv.Itoa(p.Pid())).Run())
    46  				time.Sleep(time.Second)
    47  			}
    48  		}
    49  	}
    50  
    51  	if *Slave == 0 {
    52  		rand.Seed(time.Now().Unix())
    53  		mark := rand.Intn(1024)
    54  		for i := 1; i <= *slaves; i++ {
    55  			log.Info("start slave", i, " ==== ", exec.Command(os.Args[0], "-s", strconv.Itoa(i), "-sm", strconv.Itoa(mark)).Start())
    56  		}
    57  	}
    58  
    59  	return true
    60  }