package util import ( "code.google.com/p/weed-fs/go/stats" "net" "time" ) // Listener wraps a net.Listener, and gives a place to store the timeout // parameters. On Accept, it will wrap the net.Conn with our own Conn for us. type Listener struct { net.Listener ReadTimeout time.Duration WriteTimeout time.Duration } func (l *Listener) Accept() (net.Conn, error) { c, err := l.Listener.Accept() if err != nil { return nil, err } stats.ConnectionOpen() tc := &Conn{ Conn: c, ReadTimeout: l.ReadTimeout, WriteTimeout: l.WriteTimeout, } return tc, nil } // Conn wraps a net.Conn, and sets a deadline for every read // and write operation. type Conn struct { net.Conn ReadTimeout time.Duration WriteTimeout time.Duration } func (c *Conn) Read(b []byte) (count int, e error) { err := c.Conn.SetReadDeadline(time.Now().Add(c.ReadTimeout)) if err != nil { return 0, err } count, e = c.Conn.Read(b) if e == nil { stats.BytesIn(int64(count)) } return } func (c *Conn) Write(b []byte) (count int, e error) { err := c.Conn.SetWriteDeadline(time.Now().Add(c.WriteTimeout)) if err != nil { return 0, err } count, e = c.Conn.Write(b) if e == nil { stats.BytesOut(int64(count)) } return } func (c *Conn) Close() error { stats.ConnectionClose() return c.Conn.Close() } func NewListener(addr string, timeout time.Duration) (net.Listener, error) { l, err := net.Listen("tcp", addr) if err != nil { return nil, err } tl := &Listener{ Listener: l, ReadTimeout: timeout, WriteTimeout: timeout, } return tl, nil }