Make sure all requests are complete before shutting down.

Missed a spot in the net/http docs. Calling Server.Shutdown() causes
ListenAndServe to return *immediately*, and when it does main() exits and
the "Upload Success" message gets cut off before it can be fully sent, which
spooks HTTP clients a bit.

To avoid that, you have to wait for the call to Server.Shutdown() to return,
and then signal back to main() that it's now safe to exit.
main
diff 2021-03-24 06:01:31 +00:00
parent 0b02918dfb
commit c1edf992f9
1 changed files with 15 additions and 2 deletions

17
main.go
View File

@ -86,6 +86,8 @@ func getIP() (string, error) {
return localAddr.IP.String(), nil return localAddr.IP.String(), nil
} }
var done = make(chan struct{})
func main() { func main() {
conf, err := getConfig() conf, err := getConfig()
if err != nil { if err != nil {
@ -122,6 +124,12 @@ func main() {
fmt.Println(err) fmt.Println(err)
os.Exit(1) os.Exit(1)
} }
// Wait for the server to finish any transfers, up to 3 seconds
select {
case <-done:
case <-time.After(3*time.Second):
}
} }
func setupDownload(server *http.Server, conf Config) { func setupDownload(server *http.Server, conf Config) {
@ -133,7 +141,7 @@ func setupDownload(server *http.Server, conf Config) {
downloads-- downloads--
if downloads == 0 { if downloads == 0 {
server.Shutdown(context.Background()) go shutdown(server)
} }
}) })
} }
@ -213,6 +221,11 @@ func setupUpload(server *http.Server, conf Config) {
} }
tpl.ExecuteTemplate(w, "UploadMessage", "Upload successful!") tpl.ExecuteTemplate(w, "UploadMessage", "Upload successful!")
server.Shutdown(context.Background()) go shutdown(server)
}) })
} }
func shutdown(server *http.Server) {
server.Shutdown(context.Background())
done <- struct{}{}
}