优雅的关闭Go TCP Server

本文将讨论如何在Go中优雅地关闭TCP服务器。虽然服务器通常不会停止运行(直到进程终止),但在某些情况下(例如在测试中),以有序的方式关闭它们是有用的。

Go TCP server的高级结构

我们先快速回顾一下Go中实现的TCP服务器的高层结构。Go在套接字的顶层提供了一些方便的抽象。下面是典型服务器的伪代码:

listener := net.Listen("tcp", ... address ...)
for {
  conn := listener.Accept()
  go handler(conn)
}

其中handler是一个阻塞函数,它等待来自客户端的命令,执行所需的处理,并将响应发送回。

鉴于这种结构,我们应该明白“关闭服务器”的含义。服务器在任何给定时间都执行两种不同的功能:

  1. 它监听新的连接
  2. 它处理已有连接

很明显,我们可以停止监听新的连接,从而处理(1);但是现有的连接呢?

不幸的是,这里没有简单的答案。TCP协议级别太低,无法最终解决此问题。如果我们想设计一个广泛适用的解决方案,我们必须保守。具体来说,最安全的方法是关闭服务器以等待客户端关闭其连接。这是我们将首先研究的方法。

步骤1:等待客户端连接关闭

在这个解决方案中,我们将显式地关闭侦听器(停止接受新连接),等待客户端结束其连接。这是一种保守的方法,但它在许多实际需要关闭服务器的场景(如测试)中非常有效。在测试中,使所有客户端在预期服务器关闭之前关闭连接很容易。

我将一段一段地展示代码,但是这里提供了完整的可运行代码示例。让我们从服务器类型和构造函数开始:

type Server struct {
  listener net.Listener
  quit     chan interface{}
  wg       sync.WaitGroup
}

func NewServer(addr string) *Server {
  s := &Server{
    quit: make(chan interface{}),
  }
  l, err := net.Listen("tcp", addr)
  if err != nil {
    log.Fatal(err)
  }
  s.listener = l
  s.wg.Add(1)
  go s.serve()
  return s
}

NewServer创建一个新的Server实例,该服务器在后台goroutine中侦听新连接。除了net.Listener之外,Server结构还包含一个用于发出关闭信号的通道和一个等待组,等待服务器的所有goroutine实际完成。

以下是构造函数调用的服务方法:

func (s *Server) serve() {
  defer s.wg.Done()

  for {
    conn, err := s.listener.Accept()
    if err != nil {
      select {
      case <-s.quit:
        return
      default:
        log.Println("accept error", err)
      }
    } else {
      s.wg.Add(1)
      go func() {
        s.handleConection(conn)
        s.wg.Done()
      }()
    }
  }
}

这是一个标准的Accept循环,除了select。此select所做的是在接受错误输出时检查(以非阻塞方式)s.quit通道上是否存在事件(例如发送或关闭)。如果有,则意味着错误是由我们关闭侦听器引起的,并且服务将安静地返回。如果Accept返回时没有错误,则运行连接处理程序[1]。

下面是告诉服务器正常关闭的Stop方法:

func (s *Server) Stop() {
  close(s.quit)
  s.listener.Close()
  s.wg.Wait()
}

首先关闭s.quit通道。然后它关闭监听器。这将导致服务中的Accept调用返回错误。由于s.quit此时已关闭,服务将返回,不再处理。

Stop方法的最后一行是在s.wg,这也是关键。注意,serve通知等待组它在返回时完成。但这不是我们等待的唯一一次。对handleConnection的每个调用也由wg add/done包装。因此,Stop将阻塞直到所有处理程序都返回,而serve将停止接受新连接。这是一个安全的关闭点。

为了完整起见,这里是handleConnection;这里的handleConnection只读取客户端数据并将其记录下来,而不发送任何数据。当然,这部分代码对于每个服务器都是不同的:

func (s *Server) handleConection(conn net.Conn) {
  defer conn.Close()
  buf := make([]byte, 2048)
  for {
    n, err := conn.Read(buf)
    if err != nil && err != io.EOF {
      log.Println("read error", err)
      return
    }
    if n == 0 {
      return
    }
    log.Printf("received from %v: %s", conn.RemoteAddr(), string(buf[:n]))
  }
}

使用此服务器很简单:

s := NewServer(addr)
// do whatever here...
s.Stop()

回想一下,NewServer返回一个服务器,但不阻塞。s.Stop确实会阻塞。在测试中,您要做的是:

  1. 确保与服务器交互的所有客户端都已关闭其连接。
  2. 等待s.Stop。

步骤2:主动关闭打开的客户端连接

在步骤1中,我们希望所有客户端在声明关闭进程成功之前关闭其连接。在这里,我们将看到一种更激进的方法,在Stop()中,服务器将主动尝试关闭打开的客户端连接。首先,我将介绍一种既简单又健壮的技术,以牺牲一些性能为代价。之后,我们将讨论一些替代方案。

此步骤的完整代码。与步骤1相同,只是handleConection的代码:

func (s *Server) handleConection(conn net.Conn) {
  defer conn.Close()
  buf := make([]byte, 2048)
ReadLoop:
  for {
    select {
    case <-s.quit:
      return
    default:
      conn.SetDeadline(time.Now().Add(200 * time.Millisecond))
      n, err := conn.Read(buf)
      if err != nil {
        if opErr, ok := err.(*net.OpError); ok && opErr.Timeout() {
          continue ReadLoop
        } else if err != io.EOF {
          log.Println("read error", err)
          return
        }
      }
      if n == 0 {
        return
      }
      log.Printf("received from %v: %s", conn.RemoteAddr(), string(buf[:n]))
    }
  }
}

此处理程序为每个套接字读取设置一个截止日期。这里的截止时间是200毫秒,但可以设置为对您的特定应用程序有意义的任何其他时间。如果读取返回超时,则表示客户端在超时期间处于空闲状态,连接可以安全关闭。所以循环的每次迭代都会检查s.quit并返回是否存在事件。

这种方法是健壮的,因为我们(很可能)不会在客户端主动发送消息时关闭连接。它也很简单,因为它将所有额外的逻辑限制为handleConnection

当然,这里还有一些性能损耗。首先,每200毫秒发出一次conn.Read调用,这比单个阻塞调用稍慢;不过,我认为这可以忽略不计。更严重的是,每一个Stop请求都会延迟200毫秒。在大多数情况下,如果我们想关闭服务器,这可能是可以的,但是可以根据特定的协议需要调整截止时间。

这种设计的另一种方法是跟踪handleconaction外部所有打开的连接,并在调用Stop时强制关闭它们。这将可能是更高效的,以实现复杂性和一些缺乏鲁棒性为代价。这样的Stop很容易在客户端主动发送数据时关闭连接,从而导致客户端错误。

为了获得正确路径上的灵感,我们可以查看标准库的http.Server.Shutdown方法,其文档如下:

Shutdown gracefully shuts down the server without interrupting any active connections. Shutdown works by first closing all open listeners, then closing all idle connections, and then waiting indefinitely for connections to return to idle and then shut down

“idle”在这里是什么意思?大致上,客户端已经有一段时间没有发送任何请求了。HTTP服务器比一般的TCP服务器有优势,因为它是一个更高级别的协议,所以它知道客户端通信模式。在不同的协议中,不同的关闭策略可能是有意义的。

另一个例子是服务端发起消息的协议,或者至少其中一些消息。例如,给定的连接可能处于客户端等待服务器发送某个事件的状态。服务端关闭此连接通常是安全的,无需等待任何东西。

结论

我将用两个一般准则来总结这篇文章:

  1. 尽可能安全地关闭
  2. 考虑更高层次的协议

我通常在编写测试时遇到关闭TCP服务器的需要。我希望每个测试都是独立的,并在测试完成后进行清理,包括所有的客户端-服务器连接和监听服务器。对于这个场景,步骤1非常有效。关闭所有客户端连接后,Server.Stop将立即返回。

[^1]: 注意使用WaitGroup的模式:wg.Add(1)是在go语句启动goroutine之前调用的。这是在启动go s.serve()之前在构造函数中完成的。这种方式对安全很重要。如果我们在goroutine内部调用wg.Add(1),在goroutine有机会运行之前调用wg.Wait()的执行序列可能会发生;因为在这种情况下,Wait组中尚未添加任何内容,Wait将返回;这显然不是我们想要的。

原文链接

Graceful shutdown of a TCP server in Go