net_timeout.go 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788
  1. package util
  2. import (
  3. "net"
  4. "time"
  5. "github.com/chrislusf/seaweedfs/weed/stats"
  6. )
  7. // Listener wraps a net.Listener, and gives a place to store the timeout
  8. // parameters. On Accept, it will wrap the net.Conn with our own Conn for us.
  9. type Listener struct {
  10. net.Listener
  11. ReadTimeout time.Duration
  12. WriteTimeout time.Duration
  13. }
  14. func (l *Listener) Accept() (net.Conn, error) {
  15. c, err := l.Listener.Accept()
  16. if err != nil {
  17. return nil, err
  18. }
  19. stats.ConnectionOpen()
  20. tc := &Conn{
  21. Conn: c,
  22. ReadTimeout: l.ReadTimeout,
  23. WriteTimeout: l.WriteTimeout,
  24. }
  25. return tc, nil
  26. }
  27. // Conn wraps a net.Conn, and sets a deadline for every read
  28. // and write operation.
  29. type Conn struct {
  30. net.Conn
  31. ReadTimeout time.Duration
  32. WriteTimeout time.Duration
  33. }
  34. func (c *Conn) Read(b []byte) (count int, e error) {
  35. if c.ReadTimeout != 0 {
  36. err := c.Conn.SetReadDeadline(time.Now().Add(c.ReadTimeout))
  37. if err != nil {
  38. return 0, err
  39. }
  40. }
  41. count, e = c.Conn.Read(b)
  42. if e == nil {
  43. stats.BytesIn(int64(count))
  44. }
  45. return
  46. }
  47. func (c *Conn) Write(b []byte) (count int, e error) {
  48. if c.WriteTimeout != 0 {
  49. err := c.Conn.SetWriteDeadline(time.Now().Add(c.WriteTimeout))
  50. if err != nil {
  51. return 0, err
  52. }
  53. }
  54. count, e = c.Conn.Write(b)
  55. if e == nil {
  56. stats.BytesOut(int64(count))
  57. }
  58. return
  59. }
  60. func (c *Conn) Close() error {
  61. err := c.Conn.Close()
  62. if err == nil {
  63. stats.ConnectionClose()
  64. }
  65. return err
  66. }
  67. func NewListener(addr string, timeout time.Duration) (net.Listener, error) {
  68. l, err := net.Listen("tcp", addr)
  69. if err != nil {
  70. return nil, err
  71. }
  72. tl := &Listener{
  73. Listener: l,
  74. ReadTimeout: timeout,
  75. WriteTimeout: timeout,
  76. }
  77. return tl, nil
  78. }