diff --git a/grace/conn.go b/grace/conn.go index c7b35b80..6807e1ac 100644 --- a/grace/conn.go +++ b/grace/conn.go @@ -1,13 +1,28 @@ package grace -import "net" +import ( + "errors" + "net" +) type graceConn struct { net.Conn server *Server } -func (c graceConn) Close() error { +func (c graceConn) Close() (err error) { + defer func() { + if r := recover(); r != nil { + switch x := r.(type) { + case string: + err = errors.New(x) + case error: + err = x + default: + err = errors.New("Unknown panic") + } + } + }() c.server.wg.Done() return c.Conn.Close() }