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.

128 lines
2.7 KiB

  1. package util
  2. import (
  3. "github.com/seaweedfs/seaweedfs/weed/glog"
  4. "net"
  5. "time"
  6. "github.com/seaweedfs/seaweedfs/weed/stats"
  7. )
  8. // Listener wraps a net.Listener, and gives a place to store the timeout
  9. // parameters. On Accept, it will wrap the net.Conn with our own Conn for us.
  10. type Listener struct {
  11. net.Listener
  12. ReadTimeout time.Duration
  13. WriteTimeout time.Duration
  14. }
  15. func (l *Listener) Accept() (net.Conn, error) {
  16. c, err := l.Listener.Accept()
  17. if err != nil {
  18. return nil, err
  19. }
  20. stats.ConnectionOpen()
  21. tc := &Conn{
  22. Conn: c,
  23. ReadTimeout: l.ReadTimeout,
  24. WriteTimeout: l.WriteTimeout,
  25. }
  26. return tc, nil
  27. }
  28. // Conn wraps a net.Conn, and sets a deadline for every read
  29. // and write operation.
  30. type Conn struct {
  31. net.Conn
  32. ReadTimeout time.Duration
  33. WriteTimeout time.Duration
  34. isClosed bool
  35. bytesRead int64
  36. bytesWritten int64
  37. }
  38. func (c *Conn) Read(b []byte) (count int, e error) {
  39. if c.ReadTimeout != 0 {
  40. err := c.Conn.SetReadDeadline(time.Now().Add(c.ReadTimeout * time.Duration(c.bytesRead/40000+1)))
  41. if err != nil {
  42. return 0, err
  43. }
  44. }
  45. count, e = c.Conn.Read(b)
  46. if e == nil {
  47. stats.BytesIn(int64(count))
  48. c.bytesRead += int64(count)
  49. }
  50. return
  51. }
  52. func (c *Conn) Write(b []byte) (count int, e error) {
  53. if c.WriteTimeout != 0 {
  54. // minimum 4KB/s
  55. err := c.Conn.SetWriteDeadline(time.Now().Add(c.WriteTimeout * time.Duration(c.bytesWritten/40000+1)))
  56. if err != nil {
  57. return 0, err
  58. }
  59. }
  60. count, e = c.Conn.Write(b)
  61. if e == nil {
  62. stats.BytesOut(int64(count))
  63. c.bytesWritten += int64(count)
  64. }
  65. return
  66. }
  67. func (c *Conn) Close() error {
  68. err := c.Conn.Close()
  69. if err == nil {
  70. if !c.isClosed {
  71. stats.ConnectionClose()
  72. c.isClosed = true
  73. }
  74. }
  75. return err
  76. }
  77. func NewListener(addr string, timeout time.Duration) (ipListner net.Listener, err error) {
  78. listner, err := net.Listen("tcp", addr)
  79. if err != nil {
  80. return
  81. }
  82. ipListner = &Listener{
  83. Listener: listner,
  84. ReadTimeout: timeout,
  85. WriteTimeout: timeout,
  86. }
  87. return
  88. }
  89. func NewIpAndLocalListeners(host string, port int, timeout time.Duration) (ipListner net.Listener, localListener net.Listener, err error) {
  90. listner, err := net.Listen("tcp", JoinHostPort(host, port))
  91. if err != nil {
  92. return
  93. }
  94. ipListner = &Listener{
  95. Listener: listner,
  96. ReadTimeout: timeout,
  97. WriteTimeout: timeout,
  98. }
  99. if host != "localhost" && host != "" && host != "0.0.0.0" && host != "127.0.0.1" && host != "[::]" && host != "[::1]" {
  100. listner, err = net.Listen("tcp", JoinHostPort("localhost", port))
  101. if err != nil {
  102. glog.V(0).Infof("skip starting on %s:%d: %v", host, port, err)
  103. return ipListner, nil, nil
  104. }
  105. localListener = &Listener{
  106. Listener: listner,
  107. ReadTimeout: timeout,
  108. WriteTimeout: timeout,
  109. }
  110. }
  111. return
  112. }