src/dma/vendor/github.com/go-redis/redis/internal/pool/conn.go

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