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.
80 lines
1.5 KiB
80 lines
1.5 KiB
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
|
|
}
|