conn.go

   1 // Deadlined socket read/write. https://github.com/golang/go/issues/16100
   2
   3 package godlighty
   4
   5 import (
   6         "net"
   7         "time"
   8 )
   9
  10 type DeadlinedConn struct {
  11         net.Conn
  12         ReadTimeout              time.Duration
  13         WriteTimeout             time.Duration
  14         ReadThreshold            int64
  15         WriteThreshold           int64
  16         BytesReadFromDeadline    int64
  17         BytesWrittenFromDeadline int64
  18 }
  19
  20 func (c *DeadlinedConn) Read(b []byte) (n int, err error) {
  21         if c.BytesReadFromDeadline > c.ReadThreshold {
  22                 c.BytesReadFromDeadline = 0
  23                 err = c.Conn.SetDeadline(time.Now().Add(c.ReadTimeout))
  24                 if err != nil {
  25                         return
  26                 }
  27         }
  28         n, err = c.Conn.Read(b)
  29         c.BytesReadFromDeadline += int64(n)
  30         return
  31 }
  32
  33 func (c *DeadlinedConn) Write(b []byte) (n int, err error) {
  34         if c.BytesWrittenFromDeadline > c.WriteThreshold {
  35                 c.BytesWrittenFromDeadline = 0
  36                 err = c.Conn.SetWriteDeadline(time.Now().Add(c.WriteTimeout))
  37                 if err != nil {
  38                         return
  39                 }
  40         }
  41         n, err = c.Conn.Write(b)
  42         c.BytesWrittenFromDeadline += int64(n)
  43         return
  44 }
  45
  46 type DeadlinedListener struct {
  47         net.Listener
  48         ReadTimeout  time.Duration
  49         WriteTimeout time.Duration
  50 }
  51
  52 func (l *DeadlinedListener) Accept() (net.Conn, error) {
  53         c, err := l.Listener.Accept()
  54         if err != nil {
  55                 return nil, err
  56         }
  57         return &DeadlinedConn{
  58                 Conn:           c,
  59                 ReadTimeout:    l.ReadTimeout,
  60                 WriteTimeout:   l.WriteTimeout,
  61                 ReadThreshold:  int64((l.ReadTimeout * 1024) / time.Second),
  62                 WriteThreshold: int64((l.WriteTimeout * 1024) / time.Second),
  63         }, nil
  64 }
  65
  66 func DeadlinedListen(
  67         network, addr string,
  68         readTimeout, writeTimeout time.Duration,
  69 ) (net.Listener, error) {
  70         l, err := net.Listen(network, addr)
  71         if err != nil {
  72                 return nil, err
  73         }
  74         return &DeadlinedListener{
  75                 Listener:     l,
  76                 ReadTimeout:  readTimeout,
  77                 WriteTimeout: writeTimeout,
  78         }, nil
  79 }