Tcp 长连接服务优雅重启的秘密

type gracefulTcpServer struct {
	listener     *net.TCPListener
	shutdownChan chan struct{}
	conns        map[net.Conn]struct{}

	servingConnCount atomic.Int32
	serveRunning     atomic.Bool
}


// 普通启动方式
func start(port int) (*gracefulTcpServer, error) {
	ln, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
    // handle error ignored

	s := &gracefulTcpServer{
		listener:         ln.(*net.TCPListener),
		shutdownChan:     make(chan struct{}, 1),
		conns:            make(map[net.Conn]struct{}, 16),
		servingConnCount: atomic.Int32{},
		serveRunning:     atomic.Bool{},
	}

	return s, nil
}

// 优雅重启启动方式
func startFromFork() (*gracefulTcpServer, error) {
    // ... ignored code

    // 从环境变量中获取 父进程的处理的连接数，用来恢复连接
	if nfdStr := os.Getenv(__GRACE_ENV_NFDS); nfdStr == "" {
		panic("not nfds env")
	} else if nfd, err = strconv.Atoi(nfdStr); err != nil {
		panic(err)
	}

	// restore conn fds, 0, 1, 2 has been used by os.Stdin, os.Stdout, os.Stderr
	lfd := os.NewFile(3, filepath.Join(tmpdir, "graceful"))
	ln, err := net.FileListener(lfd)
	// handle error ignored

	s := &gracefulTcpServer{
		listener:         ln.(*net.TCPListener),
		shutdownChan:     make(chan struct{}, 1),
		conns:            make(map[net.Conn]struct{}, 16),
		servingConnCount: atomic.Int32{},
		serveRunning:     atomic.Bool{},
	}

    // 从父进程继承的套接字中恢复连接
	for i := 0; i < nfd; i++ {
		fd := os.NewFile(uintptr(4+i), filepath.Join(tmpdir, strconv.Itoa(4+i)))
		conn, err := net.FileConn(fd)
		// handle error ignored
		go s.handleConn(conn)
	}

	return s, nil
}

func (s *gracefulTcpServer) gracefulRestart() {
	_ = s.listener.SetDeadline(time.Now())
	lfd, err := s.listener.File()

    // 给子进程设置 优雅重启 相关的环境变量
	os.Setenv(__GRACE_ENV_FLAG, "true")
	os.Setenv(__GRACE_ENV_NFDS, strconv.Itoa(len(s.conns)))

    // 将父进程的监听套接字传递给子进程
	files := make([]uintptr, 4, 3+1+len(s.conns))
	copy(files[:4], []uintptr{
		os.Stdin.Fd(),
		os.Stdout.Fd(),
		os.Stderr.Fd(),
		lfd.Fd(),
	})
    // 将父进程的套接字传递给子进程 
	for conn := range s.conns {
		connFd, _ := conn.(*net.TCPConn).File()
		files = append(files, connFd.Fd())
	}
	procAttr := &syscall.ProcAttr{
		Env:   os.Environ(),
		Files: files,
		Sys:   nil,
	}

    // 执行 fork + exec 调用
	childPid, err := syscall.ForkExec(os.Args[0], os.Args, procAttr)
}


func main() {
    // ...
    
    // 根据环境变量判断是 fork 还是新启动
	if v := os.Getenv(__GRACE_ENV_FLAG); v != "" {
		s, err = startFromFork()
	} else {
		s, err = start(*port)
	}

	go s.serve()

    // 处理信号，如果是 SIGHUP 信号，则执行 gracefulRestart 方法后再退出
	s.waitForSignals()
}

func forkAndExecInChild(argv0 *byte, argv, envv []*byte, chroot, dir *byte, attr *ProcAttr, sys *SysProcAttr, pipe int) (pid int, err Errno) {
    // ... some ignored codes

	fd := make([]int, len(attr.Files))
	nextfd = len(attr.Files)
	for i, ufd := range attr.Files {
		if nextfd < int(ufd) {
			nextfd = int(ufd)
		}
		fd[i] = int(ufd)
	}
	nextfd++

    // 调用 fork
	runtime_BeforeFork()
	r1, _, err1 = rawSyscall(abi.FuncPCABI0(libc_fork_trampoline), 0, 0, 0)
	if err1 != 0 {
		runtime_AfterFork()
		return 0, err1
	}

    // fork 的函数原型: pid_t fork(void);
    // 调用成功后，父进程返回子进程的 pid，子进程返回 0。所以这里通过 r1 的值来判断是父进程还是子进程
    // 如果是父进程，那么在这里就返回了，子进程会继续执行下面的代码。
	if r1 != 0 {
		runtime_AfterFork()
		return int(r1), 0
	}

    // 调用系统调用来设置 子进程 的各种属性：
    // 会话ID、进程组ID、用户ID、组ID、工作目录、是否前台运行等等。
    
    // 找到 fd[i] < i 并且把他们移到 len(fd) 之后，这样在后面的 dup2 中就不会被 “踩踏”。
    // 踩踏：一个变量或者资源被无意中（通常是意外地）覆盖或修改的情况。
	for i = 0; i < len(fd); i++ {
		if fd[i] >= 0 && fd[i] < i {
			if nextfd == pipe { // don't stomp on pipe
				nextfd++
			}
			if runtime.GOOS == "openbsd" {
				_, _, err1 = rawSyscall(dupTrampoline, uintptr(fd[i]), uintptr(nextfd), O_CLOEXEC)
			} else {
				_, _, err1 = rawSyscall(dupTrampoline, uintptr(fd[i]), uintptr(nextfd), 0)
				if err1 != 0 {
					goto childerror
				}
				_, _, err1 = rawSyscall(abi.FuncPCABI0(libc_fcntl_trampoline), uintptr(nextfd), F_SETFD, FD_CLOEXEC)
			}
			if err1 != 0 {
				goto childerror
			}
			fd[i] = nextfd
			nextfd++
		}
	}

	// 调用 dup2 来复制父进程的文件描述符，也就是 gracefulTcpServer.gracefulRestart 中传递给子进程的文件描述符
	for i = 0; i < len(fd); i++ {
        // ... some ignored codes

        // 基本只有 0，1，2 会触发这里的逻辑
		if fd[i] == i {
			// dup2(i, i) won't clear close-on-exec flag on Linux,
			// probably not elsewhere either.
			_, _, err1 = rawSyscall(abi.FuncPCABI0(libc_fcntl_trampoline), uintptr(fd[i]), F_SETFD, 0)
			if err1 != 0 {
				goto childerror
			}
			continue
		}
	
        // 其他的文件描述符，监听套接字 + 连接套接字采用 dup2 来复制
        // 这样父进程退出后，子进程还可以继续使用这些套接字
		_, _, err1 = rawSyscall(abi.FuncPCABI0(libc_dup2_trampoline), uintptr(fd[i]), uintptr(i), 0)
		if err1 != 0 {
			goto childerror
		}
	}

    // ... some ignored codes

	// 调用 execve 系统调用来执行新的程序
	_, _, err1 = rawSyscall(abi.FuncPCABI0(libc_execve_trampoline),
		uintptr(unsafe.Pointer(argv0)),
		uintptr(unsafe.Pointer(&argv[0])),
		uintptr(unsafe.Pointer(&envv[0])))

childerror:
	// send error code on pipe
	rawSyscall(abi.FuncPCABI0(libc_write_trampoline), uintptr(pipe), uintptr(unsafe.Pointer(&err1)), unsafe.Sizeof(err1))
	for {
		rawSyscall(abi.FuncPCABI0(libc_exit_trampoline), 253, 0, 0)
	}
}

func main() {
    upg, _ := tableflip.New(tableflip.Options{})
    defer upg.Stop()

    go func() {
        sig := make(chan os.Signal, 1)
        signal.Notify(sig, syscall.SIGHUP)
        for range sig {
            upg.Upgrade()
        }
    }()

    // Listen must be called before Ready
    ln, _ := upg.Listen("tcp", "localhost:8080")
    defer ln.Close()

    // 运行服务
    go http.Serve(ln, nil)

    if err := upg.Ready(); err != nil {
        panic(err)
    }

    <-upg.Exit()
}

type Upgrader struct {
	// ...
	// Upgrader.Listen 调用使用的嵌入字段 Fds 的方法
	*Fds
	// ...
}

// Listen 返回从父进程继承来的文件描述符，如果没有继承来的文件描述符，那么就创建一个新的文件描述符。
func (f *Fds) Listen(network, addr string) (net.Listener, error) {
	return f.ListenWithCallback(network, addr, f.newListener)
}

func (f *Fds) ListenWithCallback(network, addr string, callback func(network, addr string) (net.Listener, error)) (net.Listener, error) {
	// 从父进程继承来的文件描述符尝试获取监听套接字
	ln, err := f.listenerLocked(network, addr)
	if err != nil {
		return nil, err
	}
	// 如果找到那么就返回
	if ln != nil {
		return ln, nil
	}

	// 调用回调函数 cllback(f.newListener) 来创建新的监听套接字
	// 这里的 f.newListener = func(network, addr string) (net.Listener, error) { 
	//  f.lc 默认是 空的 net.ListenConfig, 等价于 net.Listen(network, addr)
	//	return f.lc.Listen(context.Background(), network, addr)
	// }
	ln, err = callback(network, addr)
	if err != nil {
		return nil, fmt.Errorf("can't create new listener: %s", err)
	}

	if _, ok := ln.(Listener); !ok {
		ln.Close()
		return nil, fmt.Errorf("%T doesn't implement tableflip.Listener", ln)
	}

	// 加入到文件描述符集合中
	err = f.addListenerLocked(network, addr, ln.(Listener))
	if err != nil {
		ln.Close()
		return nil, err
	}

	return ln, nil
}

func (u *Upgrader) Upgrade() error {
	// 需要注意的是 upgradeC 是一个 chan chan error 类型的，也就是说它是一个 chan 的 chan
	response := make(chan error, 1)

	// 这里事先检查下当前进程是否处于 退出/升级 状态
	select {
	case <-u.stopC:
		return errors.New("terminating")
	case <-u.exitC:
		return errors.New("already upgraded")
	case u.upgradeC <- response:
	}

	// 阻塞在这里，等待升级结果
	return <-response
}

func (u *Upgrader) run() {
	// ... some ignored codes

	for {
		select {
		case <-parentExited:
			// ...
		case <-processReady:
			// ...
		case <-u.stopC:
			u.Fds.closeAndRemoveUsed()
			return
		case request := <-u.upgradeC:
			// 一些异常情况检测

			// 执行升级
			file, err := u.doUpgrade()
			request <- err

			if err == nil {
				// 告诉子进程，父进程已经退出
				u.exitFd <- neverCloseThisFile{file}
				u.Fds.closeUsed()
				return
			}
		}
	}
}

func (u *Upgrader) doUpgrade() (*os.File, error) {
	// 启动子进程
	// u.env 主要是对 startChild 需要的一些调用的抽象集合，为了适应不同的平台
	// u.Fds.copy() 是复制一份父进程的文件描述符集合，这样子进程就可以直接使用父进程的文件描述符
	child, err := startChild(u.env, u.Fds.copy())
	if err != nil {
		return nil, fmt.Errorf("can't start child: %s", err)
	}

	// 更多的信号处理，退出/升级 等等
}

func startChild(env *env, passedFiles map[fileName]*file) (*child, error) {
	// 开启一个管道，用来通知父进程子进程已经准备好了
	// 所以 readyR 是父进程用来读取的，readyW 是子进程用来写入的，因此 readyW 也需要传递给子进程
	readyR, readyW, err := os.Pipe()
	if err != nil {
		return nil, fmt.Errorf("pipe failed: %s", err)
	}

	// 传递文件描述符的名字
	namesR, namesW, err := os.Pipe()
	if err != nil {
		readyR.Close()
		readyW.Close()
		return nil, fmt.Errorf("pipe failed: %s", err)
	}

	// fds 和 fdNames 保证相同的顺序记录 passedFiles 中的文件描述符
	fds := []*os.File{os.Stdin, os.Stdout, os.Stderr, readyW, namesR}
	var fdNames [][]string
	for name, file := range passedFiles {
		nameSlice := make([]string, len(name))
		copy(nameSlice, name[:])
		fdNames = append(fdNames, nameSlice)
		fds = append(fds, file.File)
	}

	// 传递环境变量，sentinelEnvVar=yes 也就是告诉新的子进程，这是一个升级的子进程
	// 在 tableflip.New > newUpgrader > newParent 的开始处可以看到
	// 这个变量决定了 Upgrader 中的 parent 字段是否有值，如果有值，那么就是一个优雅升级的子进程
	sentinel := fmt.Sprintf("%s=yes", sentinelEnvVar)
	var environ []string
	for _, val := range env.environ() {
		if val != sentinel {
			environ = append(environ, val)
		}
	}
	environ = append(environ, sentinel)

	// 准备好新的子进程的 command, args 文件描述符，环境变量后就可以启动子进程了
	// env.newProc 是一个抽象方法，用来创建子进程
	// 默认使用 newOSProcess 其中核心还是调用 syscall.StartProcess 来创建子进程	
	// syscall.StartProcess 与 syscall.ForkExec 是一样的
	proc, err := env.newProc(os.Args[0], os.Args[1:], fds, environ)
	if err != nil {
		// fork/exec 执行失败，那么释放之前创建的资源
		// ...
		return
	}

	// 后续，子进程启动成功后就不会执行到这里了，因此这里的代码都是在父进程中执行的
	// 也就是说，父进程会想子进程传递一些数据（fdNames）然后等待了来自子进程的信号
	exited := make(chan struct{})
	result := make(chan error, 1)
	ready := make(chan *os.File, 1)

	c := &child{
		// ...
	}
	go c.writeNames(fdNames)
	go c.waitExit(result, exited)
	go c.waitReady(ready)
	return c, nil
}

func (u *Upgrader) Ready() error {
	u.readyOnce.Do(func() {
		u.Fds.closeInherited()
		close(u.readyC)
	})
	// some ignored codes

	// 如果没有父进程，那么就直接返回（parent == nil 意味着优雅升级进程来的）
	if u.parent == nil {
		return nil
	}

	// 通知父进程子进程已经准备好了, 父进程可以退出了
	// 这里其实就是创建子进程时开辟的 ready 管道的 write 端写入特定的信号数据 (notifyReady)
	return u.parent.sendReady()
}

#include <unistd.h>
pid_t fork(void);

#include<sys/types.h>
#include<unistd.h>
#include<stdio.h>
#include<stdlib.h>

int main()
{
	pid_t pid = fork(); // 创建子进程
	if (pid < 0) {
		perror("fork error");
		exit(1);
	}

	if (pid != 0) 
	{
		// 父进程
		printf("this is parent process, PID is %d.\n", getpid());
		exit(0);
	}

	// 子进程
	printf("this is child process %d, PID is %d.\n", i, getpid());
}

// linux-6.4/include/linux/sched.h#739
struct task_struct {
	//...

	/* Filesystem information: */
	struct fs_struct		*fs;
	/* 记录了进程打开的文件 */
	struct files_struct		*files;

	//...
}

struct files_struct {
	//...
	// 保存了进程打开的文件描述符, 是一个数组，数组的下标就是文件描述符 fd
	struct file __rcu * fd_array[NR_OPEN_DEFAULT];
}

// linux-6.4/kernel/fork.c#2246
__latent_entropy struct task_struct *copy_process(
					struct pid *pid,
					int trace,
					int node,
					struct kernel_clone_args *args)
{
	struct task_struct *p;

	// 拷贝进程打开的文件描述符
	retval = copy_files(clone_flags, p, args->no_files);
}

listenFile := os.NewFile(3, "/tmp/graceful")
ln, err := net.FileListener(listenFile)

#include <unistd.h>
int execve(const char *pathname, char *const _Nullable argv[],
           char *const _Nullable envp[]);