micro/gateway/conn.go

package gateway

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.InTraffic, int64(n))
	return
}

func (c *Conn) Write(b []byte) (n int, err error) {
	n, err = c.conn.Write(b)
	atomic.AddInt64(&c.state.OutTraffic, int64(n))
	return
}

func (c *Conn) Close() error {
	if atomic.CompareAndSwapInt32(&c.exitFlag, 0, 1) {
		atomic.AddInt32(&c.state.Concurrency, -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)
	}
	atomic.AddInt32(&state.Concurrency, 1)
	return conn
}