153 lines
3.0 KiB
Go
153 lines
3.0 KiB
Go
package main
|
||
|
||
import (
|
||
"errors"
|
||
"flag"
|
||
"io"
|
||
"log"
|
||
"net"
|
||
"os"
|
||
"os/exec"
|
||
"os/signal"
|
||
"syscall"
|
||
)
|
||
|
||
var (
|
||
c = flag.Int("c", 10, "concurrency level, default 10")
|
||
prefork = flag.Bool("prefork", false, "use prefork mode")
|
||
child = flag.Bool("child", false, "is child process")
|
||
)
|
||
|
||
func main() {
|
||
flag.Parse()
|
||
|
||
var ln net.Listener
|
||
var err error
|
||
|
||
if *prefork { // 如果要启动子进程模式
|
||
ln = doPrefork(*c)
|
||
} else {
|
||
ln, err = net.Listen("tcp", ":8972")
|
||
if err != nil {
|
||
panic(err)
|
||
}
|
||
}
|
||
|
||
start(ln) // 处理 net.Listener
|
||
}
|
||
|
||
func start(ln net.Listener) {
|
||
log.Println("started")
|
||
for {
|
||
conn, e := ln.Accept()
|
||
if e != nil {
|
||
var ne net.Error
|
||
if errors.As(e, &ne) && ne.Temporary() {
|
||
log.Printf("accept temp err: %v", ne)
|
||
continue
|
||
}
|
||
|
||
log.Printf("accept err: %v", e)
|
||
return
|
||
}
|
||
|
||
go func() {
|
||
_, err := io.Copy(conn, conn) // 实现 echo 协议，将收到的东西原样返回
|
||
if err != nil {
|
||
log.Printf("copy err: %v", err)
|
||
}
|
||
}()
|
||
}
|
||
}
|
||
|
||
func doPrefork(c int) net.Listener {
|
||
if *child {
|
||
// 子进程：从文件描述符恢复监听器
|
||
listener, err := net.FileListener(os.NewFile(3, ""))
|
||
if err != nil {
|
||
log.Fatal(err)
|
||
}
|
||
|
||
// 优雅关闭处理
|
||
sigCh := make(chan os.Signal, 1)
|
||
signal.Notify(sigCh, syscall.SIGTERM)
|
||
go func() {
|
||
<-sigCh
|
||
err := listener.Close()
|
||
if err != nil {
|
||
return
|
||
}
|
||
}()
|
||
|
||
return listener
|
||
}
|
||
|
||
// 主进程：创建监听器并将其传递给子进程
|
||
addr, err := net.ResolveTCPAddr("tcp", ":8972")
|
||
if err != nil {
|
||
log.Fatal(err)
|
||
}
|
||
tcpListener, err := net.ListenTCP("tcp", addr)
|
||
if err != nil {
|
||
log.Fatal(err)
|
||
}
|
||
fl, err := tcpListener.File()
|
||
if err != nil {
|
||
log.Fatal(err)
|
||
}
|
||
// 主进程不需要保持监听
|
||
if err := tcpListener.Close(); err != nil {
|
||
log.Fatal(err)
|
||
}
|
||
|
||
// 启动子进程并持续监控
|
||
children := make([]*exec.Cmd, c)
|
||
|
||
// 处理信号以优雅关闭
|
||
sigCh := make(chan os.Signal, 1)
|
||
signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
|
||
|
||
// 启动子进程管理协程
|
||
go func() {
|
||
<-sigCh
|
||
// 收到信号后通知所有子进程优雅关闭
|
||
for _, child := range children {
|
||
if child.Process != nil {
|
||
if err := child.Process.Signal(syscall.SIGTERM); err != nil {
|
||
log.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
os.Exit(0)
|
||
}()
|
||
|
||
// 声明子进程启动函数以便递归调用
|
||
var startChildProcess func(int)
|
||
startChildProcess = func(i int) {
|
||
cmd := exec.Command(os.Args[0], "-prefork", "-child")
|
||
cmd.Stdout = os.Stdout
|
||
cmd.Stderr = os.Stderr
|
||
cmd.ExtraFiles = []*os.File{fl}
|
||
children[i] = cmd
|
||
|
||
if err := cmd.Start(); err != nil {
|
||
log.Fatalf("启动子进程 %d 失败: %v", i, err)
|
||
}
|
||
|
||
// 在单独的协程中等待子进程结束并重启
|
||
go func() {
|
||
if err := cmd.Wait(); err != nil {
|
||
log.Printf("子进程 %d 异常退出: %v，正在重启...", i, err)
|
||
startChildProcess(i) // 递归重启
|
||
}
|
||
}()
|
||
}
|
||
// 启动和监控子进程
|
||
for i := range children {
|
||
startChildProcess(i)
|
||
}
|
||
|
||
// 主进程保持运行，但不处理连接
|
||
select {} // 阻塞主进程
|
||
}
|