server: Use explicit context arguments

Replace the Server.Context field with explicit context.Context arguments
to most Server functions.

server.go

@@ -47,16 +47,98 @@ type Server struct {
 	// If nil, logging is done via the log package's standard logger.
 	ErrorLog *log.Logger
 
-	// Context is the base context to use.
+	listeners map[*net.Listener]context.CancelFunc
-	// If nil, context.Background is used.
+	conns     map[*net.Conn]context.CancelFunc
-	Context context.Context
+	doneChan  chan struct{}
-
+	closed    int32
-	listeners map[*net.Listener]struct{}
-	conns     map[*net.Conn]struct{}
-	done      int32
 	mu        sync.Mutex
 }
 
+// done returns a channel that's closed when the server has finished closing.
+func (srv *Server) done() chan struct{} {
+	srv.mu.Lock()
+	defer srv.mu.Unlock()
+	return srv.doneLocked()
+}
+
+func (srv *Server) doneLocked() chan struct{} {
+	if srv.doneChan == nil {
+		srv.doneChan = make(chan struct{})
+	}
+	return srv.doneChan
+}
+
+// tryFinishShutdown closes srv.done() if there are no active listeners or requests.
+func (srv *Server) tryFinishShutdown() {
+	srv.mu.Lock()
+	defer srv.mu.Unlock()
+	if len(srv.listeners) == 0 && len(srv.conns) == 0 {
+		done := srv.doneLocked()
+		select {
+		case <-done:
+		default:
+			close(done)
+		}
+	}
+}
+
+// Close immediately closes all active net.Listeners and connections.
+// For a graceful shutdown, use Shutdown.
+func (srv *Server) Close() error {
+	if !atomic.CompareAndSwapInt32(&srv.closed, 0, 1) {
+		return ErrServerClosed
+	}
+
+	// Close active listeners and connections.
+	srv.mu.Lock()
+	for _, cancel := range srv.listeners {
+		cancel()
+	}
+	for _, cancel := range srv.conns {
+		cancel()
+	}
+	srv.mu.Unlock()
+
+	select {
+	case <-srv.done():
+		return nil
+	}
+}
+
+// Shutdown gracefully shuts down the server without interrupting any
+// active connections. Shutdown works by first closing all open
+// listeners and then waiting indefinitely for connections
+// to close and then shut down.
+// If the provided context expires before the shutdown is complete,
+// Shutdown returns the context's error.
+//
+// When Shutdown is called, Serve and ListenAndServer immediately
+// return ErrServerClosed. Make sure the program doesn't exit and
+// waits instead for Shutdown to return.
+//
+// Once Shutdown has been called on a server, it may not be reused;
+// future calls to methods such as Serve will return ErrServerClosed.
+func (srv *Server) Shutdown(ctx context.Context) error {
+	if !atomic.CompareAndSwapInt32(&srv.closed, 0, 1) {
+		return ErrServerClosed
+	}
+
+	// Close active listeners.
+	srv.mu.Lock()
+	for _, cancel := range srv.listeners {
+		cancel()
+	}
+	srv.mu.Unlock()
+
+	// Wait for active connections to finish.
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-srv.done():
+		return nil
+	}
+}
+
 // ListenAndServe listens for requests at the server's configured address.
 // ListenAndServe listens on the TCP network address srv.Addr and then calls
 // Serve to handle requests on incoming connections.