From 1afb49f628cef84536a0d83732f95b0a581b1381 Mon Sep 17 00:00:00 2001
From: David Cowden <dcow@smallstep.com>
Date: Tue, 16 Jun 2020 22:55:59 -0700
Subject: [PATCH] server: Add ShutdownAndWait
ShutdownAndWait mirrors the net/http Shutdown API allowing the user to
wait until the server is idle or until the provided context is canceled.
---
server.go | 22 ++++++++++++++++++++++
1 file changed, 22 insertions(+)
diff --git a/server.go b/server.go
index 906e9b5..a0c6c4e 100644
--- a/server.go
+++ b/server.go
@@ -277,6 +277,28 @@ func (srv *Server) Shutdown() error {
return nil
}
+// Shutdown the server and wait for all connections to drain gracefully. If the
+// provided context is canceled return the context's error.
+func (srv *Server) ShutdownAndWait(ctx context.Context) error {
+ err := srv.Shutdown()
+ if err != nil && err != ErrServerClosed {
+ return err
+ }
+
+ idle := make(chan struct{}, 1)
+ go func() {
+ srv.Idle.Wait()
+ close(idle)
+ }()
+
+ select {
+ case <-ctx.Done():
+ return ctx.Err()
+ case <-idle:
+ return nil
+ }
+}
+
// CloseCalled reports whether the server is Closed or not. A closed server
// does not allow new connections to be tracked.
func (srv *Server) CloseCalled() bool {
--
GitLab