73 lines
1.3 KiB
Go
73 lines
1.3 KiB
Go
package entry
|
|
|
|
import (
|
|
"net"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
type Conn struct {
|
|
buf []byte
|
|
conn net.Conn
|
|
exitFlag int32
|
|
state *State
|
|
}
|
|
|
|
func (c *Conn) Read(b []byte) (n int, err error) {
|
|
var m int
|
|
if len(c.buf) > 0 {
|
|
if len(b) >= len(c.buf) {
|
|
m = copy(b[:], c.buf[:])
|
|
c.buf = c.buf[m:]
|
|
}
|
|
}
|
|
n, err = c.conn.Read(b[m:])
|
|
n += m
|
|
atomic.AddInt64(&c.state.Traffic.In, int64(n))
|
|
return
|
|
}
|
|
|
|
func (c *Conn) Write(b []byte) (n int, err error) {
|
|
n, err = c.conn.Write(b)
|
|
atomic.AddInt64(&c.state.Traffic.Out, int64(n))
|
|
return
|
|
}
|
|
|
|
func (c *Conn) Close() error {
|
|
if atomic.CompareAndSwapInt32(&c.exitFlag, 0, 1) {
|
|
atomic.AddInt32(&c.state.Concurrency, -1)
|
|
atomic.AddInt64(&c.state.Request.Processed, 1)
|
|
return c.conn.Close()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (c *Conn) LocalAddr() net.Addr {
|
|
return c.conn.LocalAddr()
|
|
}
|
|
|
|
func (c *Conn) RemoteAddr() net.Addr {
|
|
return c.conn.RemoteAddr()
|
|
}
|
|
|
|
func (c *Conn) SetDeadline(t time.Time) error {
|
|
return c.conn.SetDeadline(t)
|
|
}
|
|
|
|
func (c *Conn) SetReadDeadline(t time.Time) error {
|
|
return c.conn.SetReadDeadline(t)
|
|
}
|
|
|
|
func (c *Conn) SetWriteDeadline(t time.Time) error {
|
|
return c.conn.SetWriteDeadline(t)
|
|
}
|
|
|
|
func wrapConn(c net.Conn, state *State, buf []byte) net.Conn {
|
|
conn := &Conn{conn: c, buf: buf, state: state}
|
|
if buf != nil {
|
|
conn.buf = make([]byte, len(buf))
|
|
copy(conn.buf, buf)
|
|
}
|
|
return conn
|
|
}
|