本站源代码
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.

94 lines
1.8KB

  1. package pool
  2. import (
  3. "net"
  4. "sync/atomic"
  5. "time"
  6. "github.com/go-redis/redis/internal/proto"
  7. )
  8. var noDeadline = time.Time{}
  9. type Conn struct {
  10. netConn net.Conn
  11. rd *proto.Reader
  12. rdLocked bool
  13. wr *proto.Writer
  14. InitedAt time.Time
  15. pooled bool
  16. usedAt atomic.Value
  17. }
  18. func NewConn(netConn net.Conn) *Conn {
  19. cn := &Conn{
  20. netConn: netConn,
  21. }
  22. cn.rd = proto.NewReader(netConn)
  23. cn.wr = proto.NewWriter(netConn)
  24. cn.SetUsedAt(time.Now())
  25. return cn
  26. }
  27. func (cn *Conn) UsedAt() time.Time {
  28. return cn.usedAt.Load().(time.Time)
  29. }
  30. func (cn *Conn) SetUsedAt(tm time.Time) {
  31. cn.usedAt.Store(tm)
  32. }
  33. func (cn *Conn) SetNetConn(netConn net.Conn) {
  34. cn.netConn = netConn
  35. cn.rd.Reset(netConn)
  36. cn.wr.Reset(netConn)
  37. }
  38. func (cn *Conn) setReadTimeout(timeout time.Duration) error {
  39. now := time.Now()
  40. cn.SetUsedAt(now)
  41. if timeout > 0 {
  42. return cn.netConn.SetReadDeadline(now.Add(timeout))
  43. }
  44. return cn.netConn.SetReadDeadline(noDeadline)
  45. }
  46. func (cn *Conn) setWriteTimeout(timeout time.Duration) error {
  47. now := time.Now()
  48. cn.SetUsedAt(now)
  49. if timeout > 0 {
  50. return cn.netConn.SetWriteDeadline(now.Add(timeout))
  51. }
  52. return cn.netConn.SetWriteDeadline(noDeadline)
  53. }
  54. func (cn *Conn) Write(b []byte) (int, error) {
  55. return cn.netConn.Write(b)
  56. }
  57. func (cn *Conn) RemoteAddr() net.Addr {
  58. return cn.netConn.RemoteAddr()
  59. }
  60. func (cn *Conn) WithReader(timeout time.Duration, fn func(rd *proto.Reader) error) error {
  61. _ = cn.setReadTimeout(timeout)
  62. return fn(cn.rd)
  63. }
  64. func (cn *Conn) WithWriter(timeout time.Duration, fn func(wr *proto.Writer) error) error {
  65. _ = cn.setWriteTimeout(timeout)
  66. firstErr := fn(cn.wr)
  67. err := cn.wr.Flush()
  68. if err != nil && firstErr == nil {
  69. firstErr = err
  70. }
  71. return firstErr
  72. }
  73. func (cn *Conn) Close() error {
  74. return cn.netConn.Close()
  75. }
上海开阖软件有限公司 沪ICP备12045867号-1