mirror of https://github.com/ehang-io/nps
40 lines
515 B
Go
40 lines
515 B
Go
|
package grace
|
||
|
|
|
||
|
|
import (
|
||
|
|
"errors"
|
||
|
|
"net"
|
||
|
|
"sync"
|
||
|
|
)
|
||
|
|
|
||
|
|
type graceConn struct {
|
||
|
|
net.Conn
|
||
|
|
server *Server
|
||
|
|
m sync.Mutex
|
||
|
|
closed bool
|
||
|
|
}
|
||
|
|
|
||
|
|
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.m.Lock()
|
||
|
|
if c.closed {
|
||
|
|
c.m.Unlock()
|
||
|
|
return
|
||
|
|
}
|
||
|
|
c.server.wg.Done()
|
||
|
|
c.closed = true
|
||
|
|
c.m.Unlock()
|
||
|
|
return c.Conn.Close()
|
||
|
|
}
|