open-consul/vendor/github.com/hashicorp/go-connlimit/wrap.go

28 lines
671 B
Go

package connlimit
import "net"
// WrappedConn wraps a net.Conn and free() func returned by Limiter.Accept so
// that when the wrapped connections Close method is called, its free func is
// also called.
type WrappedConn struct {
net.Conn
free func()
}
// Wrap wraps a net.Conn's Close method so free() is called when Close is
// called. Useful when handing off tracked connections to libraries that close
// them.
func Wrap(conn net.Conn, free func()) net.Conn {
return &WrappedConn{
Conn: conn,
free: free,
}
}
// Close frees the tracked connection and closes the underlying net.Conn.
func (w *WrappedConn) Close() error {
w.free()
return w.Conn.Close()
}