Add option to run functions on server shutdown

src/Servers.jl

@@ -74,10 +74,29 @@ struct Server{S <: Union{SSLConfig, Nothing}, I <: Base.IOServer}
     server::I
     hostname::String
     hostport::String
+    on_shutdown::Any
 end
 
 Base.isopen(s::Server) = isopen(s.server)
-Base.close(s::Server) = close(s.server)
+Base.close(s::Server) = (shutdown(s.on_shutdown); close(s.server))
+
+"""
+    shutdown(fns::Vector{<:Function})
+    shutdown(fn::Function)
+    shutdown(::Nothing)
+
+Runs function(s) in `on_shutdown` field of `Server` when
+`Server` is closed.
+"""
+shutdown(fns::Vector{<:Function}) = foreach(shutdown, fns)
+shutdown(::Nothing) = nothing
+function shutdown(fn::Function)
+    try
+        fn()
+    catch e
+        @error "shutdown function $fn failed" exception=(e, catch_backtrace())
+    end
+end
 
 Sockets.accept(s::Server{Nothing}) = accept(s.server)::TCPSocket
 Sockets.accept(s::Server{SSLConfig}) = getsslcontext(accept(s.server), s.ssl)
@@ -118,6 +137,10 @@ Optional keyword arguments:
     allowed per client IP address; excess connections are immediately closed.
     e.g. 5//1.
  - `verbose::Bool=false`, log connection information to `stdout`.
+ - `on_shutdown::Union{Function, Vector{<:Function}, Nothing}=nothing`, one or
+    more functions to be run if the server is closed (for example by an
+    `InterruptException`). Note, shutdown function(s) will not run if a
+    `IOServer` object is supplied and closed by `close(server)`.
 
 e.g.
 ```julia
@@ -203,7 +226,8 @@ function listen(f,
                 rate_limit::Union{Rational{Int}, Nothing}=nothing,
                 reuse_limit::Int=nolimit,
                 readtimeout::Int=0,
-                verbose::Bool=false)
+                verbose::Bool=false,
+                on_shutdown::Union{Function, Vector{<:Function}, Nothing}=nothing)
 
     inet = getinet(host, port)
     if server !== nothing
@@ -230,7 +254,7 @@ function listen(f,
         x -> f(x) && check_rate_limit(x, rate_limit)
     end
 
-    s = Server(sslconfig, tcpserver, string(host), string(port))
+    s = Server(sslconfig, tcpserver, string(host), string(port), on_shutdown)
     return listenloop(f, s, tcpisvalid, connection_count,
                          reuse_limit, readtimeout, verbose)
 end

test/server.jl

@@ -168,4 +168,28 @@ end
     @test_skip HTTP.hasheader(HTTP.get("http://httpbin.org/redirect-to?url=https://httpbin.org/response-headers?Authorization=auth"), "Authorization")
 end # @testset
 
+@testset "on_shutdown" begin
+    @test HTTP.Servers.shutdown(nothing) === nothing
+
+    IOserver = Sockets.listen(Sockets.InetAddr(parse(IPAddr, "127.0.0.1"), 8052))
+
+    # Shutdown adds 1
+    TEST_COUNT = Ref(0)
+    shutdown_add() = TEST_COUNT[] += 1
+    server = HTTP.Servers.Server(nothing, IOserver, "host", "port", shutdown_add)
+    close(server)
+
+    # Shutdown adds 1, performed twice
+    @test TEST_COUNT[] == 1
+    server = HTTP.Servers.Server(nothing, IOserver, "host", "port", [shutdown_add, shutdown_add])
+    close(server)
+    @test TEST_COUNT[] == 3
+
+    # First shutdown function errors, second adds 1
+    shutdown_throw() = throw(ErrorException("Broken"))
+    server = HTTP.Servers.Server(nothing, IOserver, "host", "port", [shutdown_throw, shutdown_add])
+    @test_logs (:error, r"shutdown function .* failed") close(server)
+    @test TEST_COUNT[] == 4
+end # @testset
+
 end # module