You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

97 lines
1.9 KiB

  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. isClosed bool
  34. bytesRead int64
  35. bytesWritten int64
  36. }
  37. func (c *Conn) Read(b []byte) (count int, e error) {
  38. if c.ReadTimeout != 0 {
  39. err := c.Conn.SetReadDeadline(time.Now().Add(c.ReadTimeout * time.Duration(c.bytesRead/40000+1)))
  40. if err != nil {
  41. return 0, err
  42. }
  43. }
  44. count, e = c.Conn.Read(b)
  45. if e == nil {
  46. stats.BytesIn(int64(count))
  47. c.bytesRead += int64(count)
  48. }
  49. return
  50. }
  51. func (c *Conn) Write(b []byte) (count int, e error) {
  52. if c.WriteTimeout != 0 {
  53. // minimum 4KB/s
  54. err := c.Conn.SetWriteDeadline(time.Now().Add(c.WriteTimeout * time.Duration(c.bytesWritten/40000+1)))
  55. if err != nil {
  56. return 0, err
  57. }
  58. }
  59. count, e = c.Conn.Write(b)
  60. if e == nil {
  61. stats.BytesOut(int64(count))
  62. c.bytesWritten += int64(count)
  63. }
  64. return
  65. }
  66. func (c *Conn) Close() error {
  67. err := c.Conn.Close()
  68. if err == nil {
  69. if !c.isClosed {
  70. stats.ConnectionClose()
  71. c.isClosed = true
  72. }
  73. }
  74. return err
  75. }
  76. func NewListener(addr string, timeout time.Duration) (net.Listener, error) {
  77. l, err := net.Listen("tcp", addr)
  78. if err != nil {
  79. return nil, err
  80. }
  81. tl := &Listener{
  82. Listener: l,
  83. ReadTimeout: timeout,
  84. WriteTimeout: timeout,
  85. }
  86. return tl, nil
  87. }