// First: go get github.com/golang/glog | |
// Then : go test -v -vmodule=*=3 -logtostderr netgrace_test.go | |
// This test package shows how to gracefully shutdown a net.Listener. | |
package netgrace | |
import ( | |
"bufio" | |
"fmt" | |
"github.com/golang/glog" | |
"io" | |
"net" | |
"testing" | |
"time" | |
) | |
var ( | |
// A list of connections created by the server. We keep track of them so we | |
// can gracefully shut them down if they are still alive when the server | |
// goes down. | |
conns []net.Conn | |
// The quit channel for the server. If the server detects that this channel | |
// is closed, then it's a signal for it to shutdown as well. | |
quit chan bool | |
// The TCP server we are going to start | |
ln net.Listener | |
) | |
func init() { | |
quit = make(chan bool) | |
conns = make([]net.Conn, 0, 10) | |
} | |
func startServer() (err error) { | |
// When the server exits, we want to clean up after ourselves. | |
// Most likely all the connections have already been nil'ed but just in case. | |
defer func() { | |
glog.Info("Shutting down server") | |
for i, conn := range conns { | |
if conn != nil { | |
glog.Infof("Closing connection #%d", i) | |
conn.Close() | |
} | |
} | |
}() | |
// Start a echo server on port 5388/tcp | |
ln, err = net.Listen("tcp", ":5388") | |
if err != nil { | |
return err | |
} | |
defer ln.Close() | |
for { | |
glog.Info("Listening for connections") | |
// Accept() blocks waiting for new connection. The problem is that when | |
// the listerner is blocked, and if there's never a new connection, then | |
// this goroutine never exits (unless the program exits). We want a way | |
// to gracefully stop this, so that this goroutine doesn't hang around | |
// forever. For example, this is important if you are running tests with | |
// a client and server, where you want the server to exit after a certain | |
// period of time. If you don't somehow trigger the server to exit while | |
// it's blocked on Accept(), then the test will block. | |
// | |
// Our requirement is to Accept() connections. If an error is returned, | |
// we will try to handle it and then continue Accept(). However, we don't | |
// want to continue Accept() if we are told to quit. | |
// | |
// One way to do this is Close() the net.Listener, which forces Accept() | |
// to return with the errClosing error. Then we can check to see if the | |
// error is because net.Listener is closed, or it's some other error. | |
// If net.Listener is closed, then we would quit. But if it's some other | |
// error, maybe we can handle it and then go back to Accept(). | |
// | |
// Unfortunately errClosing is NOT exported from the net package so there's | |
// no way to check if Accept() returned because the net.Listener is closed. | |
// Others have tried to compare the actual error message ("use of closed | |
// network connection"). However that is not the ideal since the actual | |
// error message may change in the future. | |
// | |
// There was a long discussion thread, | |
// https://code.google.com/p/go/issues/detail?id=4373, | |
// that took place in Nov 2012 regarding the possibility of exporting | |
// errClosing. But at the end the Go Authors decided it is not a | |
// prudent thing to do so the issue was tagged as wontfix. | |
// | |
// What we need is a way to check to see if we are being told to quit, | |
// after Accept() returns. The trick here is a quit channel. It's quite | |
// simple actually. When you want the server to quit, first you | |
// close the quit channel, which tells the Accept() goroutine to quit | |
// if the goroutine checks the quit channel. Then you close the | |
// net.Listener which then forces Accept() to return. | |
// | |
// Without further ado, we start accepting connections | |
conn, err := ln.Accept() | |
if err != nil { | |
glog.Error(err.Error()) | |
// When Accept() returns with a non-nill error, we check the quit | |
// channel to see if we should continue or quit. If quit, then we quit. | |
// Otherwise we continue | |
select { | |
case <-quit: | |
return nil | |
default: | |
// thanks to martingx on reddit for noticing I am missing a default | |
// case. without the default case the select will block. | |
} | |
continue | |
} | |
// Now that we have a connection, we add it to the list, and go handle it. | |
conns = append(conns, conn) | |
go handleConnection(conn, len(conns)-1) | |
} | |
return nil | |
} | |
func handleConnection(conn net.Conn, id int) error { | |
// Again, we are cleaning up after ourselves here. At the exit of this function, | |
// we want to close the connection, and set the appropriate slot in the list | |
// to nil so we don't leak any memory. | |
defer func() { | |
glog.Infof("Closing connection #%d", id) | |
conn.Close() | |
conns[id] = nil | |
}() | |
glog.Infof("Starting connection #%d", id) | |
for { | |
// This is effectively an echo server. Reads from the connection and | |
// immediately write it back. If io.Copy() returns, then it's either | |
// because the socket is closed (err == nil), or there's some type of | |
// real error. Either case we return. | |
if _, err := io.Copy(conn, conn); err != nil { | |
glog.Error(err.Error()) | |
return err | |
} | |
return nil | |
} | |
} | |
func Test10Clients(t *testing.T) { | |
defer func() { | |
glog.Infof("Stopping server...") | |
// When we exit this test, we want to make sure we clean up after | |
// ourselves so we don't leave anything behind. In this case, by | |
// closing the quit channel, we are telling the server to stop | |
// accepting new connection. | |
close(quit) | |
// We then close the net.Listener, which will force Accept() to | |
// return if it's blocked waiting for new connections. | |
ln.Close() | |
// The above order matters somewhat. If you ln.Close() first, then | |
// you run the risk of Accept() returning but the quit channel | |
// hasn't been closed. It is not the end of the world, however, | |
// it just means you will likely see quite a few more errors before | |
// the goroutine detects quit is closed. | |
}() | |
go startServer() | |
time.Sleep(time.Second) | |
// In this test, we start 10 clients, each making a single connection | |
// to the server. Then each will write 10 messages to the server, and | |
// read the same 10 messages back. After that the client quits. | |
for i := 0; i < 10; i++ { | |
go func(id int) { | |
defer func() { | |
glog.Infof("Quiting client #%d", id) | |
}() | |
conn, err := net.Dial("tcp", ":5388") | |
if err != nil { | |
glog.Error(err.Error()) | |
return | |
} | |
defer conn.Close() | |
for i := 0; i < 10; i++ { | |
fmt.Fprintf(conn, "client #%d, count %d\n", id, i) | |
res, err := bufio.NewReader(conn).ReadString('\n') | |
if err != nil { | |
glog.Error(err.Error()) | |
return | |
} | |
glog.Infof("Received: %s", res) | |
time.Sleep(100 * time.Millisecond) | |
} | |
}(i) | |
} | |
// We sleep for a couple of seconds, let the clients run their jobs, | |
// then we exit, which triggers the defer function that will shutdown | |
// the server. | |
time.Sleep(2 * time.Second) | |
// So instead of just quiting, we clean up first, well, in the defer block. | |
// This is expecially important if this is a long running program. | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
This comment has been minimized.
Thank you very much for this gist😃 Combined it with this approach