github.com/cnotch/ipchub@v1.1.0/network/socket/buffered/conn.go (about)

     1  // Copyright (c) 2019,CAOHONGJU All rights reserved.
     2  // Use of this source code is governed by a MIT-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package buffered
     6  
     7  import (
     8  	"bufio"
     9  	"bytes"
    10  	"net"
    11  	"time"
    12  
    13  	"github.com/kelindar/rate"
    14  )
    15  
    16  const (
    17  	defaultRate       = 50
    18  	defaultBufferSize = 64 * 1024
    19  	minBufferSize     = 8 * 1024
    20  )
    21  
    22  // Conn wraps a net.Conn and provides buffered ability.
    23  type Conn struct {
    24  	socket     net.Conn      // The underlying network connection.
    25  	reader     *bufio.Reader // The buffered reader
    26  	writer     *bytes.Buffer // The buffered write queue.
    27  	limit      *rate.Limiter // The write rate limiter.
    28  	bufferSize int           // The read and write max buffer size
    29  }
    30  
    31  // NewConn creates a new sniffed connection.
    32  func NewConn(c net.Conn, options ...Option) *Conn {
    33  	conn, ok := c.(*Conn)
    34  	if !ok {
    35  		conn = &Conn{
    36  			socket: c,
    37  		}
    38  	}
    39  
    40  	for _, option := range options {
    41  		option.apply(conn)
    42  	}
    43  
    44  	// 设置默认值刷新频率
    45  	if conn.limit == nil {
    46  		conn.limit = rate.New(defaultRate, time.Second)
    47  	}
    48  
    49  	if conn.bufferSize <= 0 {
    50  		conn.bufferSize = defaultBufferSize
    51  	}
    52  
    53  	// 设置IO缓冲对象
    54  	conn.reader = bufio.NewReaderSize(conn.socket, conn.bufferSize)
    55  	conn.writer = bytes.NewBuffer(make([]byte, 0, conn.bufferSize))
    56  	return conn
    57  }
    58  
    59  // Buffered returns the pending buffer size.
    60  func (m *Conn) Buffered() (n int) {
    61  	return m.writer.Len()
    62  }
    63  
    64  // Reader 返回内部的 bufio.Reader
    65  func (m *Conn) Reader() *bufio.Reader {
    66  	return m.reader
    67  }
    68  
    69  // Flush flushes the underlying buffer by writing into the underlying connection.
    70  func (m *Conn) Flush() (n int, err error) {
    71  	if m.Buffered() == 0 {
    72  		return 0, nil
    73  	}
    74  
    75  	// Flush everything and reset the buffer
    76  	n, err = m.writeFull(m.writer.Bytes())
    77  	m.writer.Reset()
    78  	return
    79  }
    80  
    81  // Read reads the block of data from the underlying buffer.
    82  func (m *Conn) Read(p []byte) (int, error) {
    83  	return m.reader.Read(p)
    84  }
    85  
    86  // Write writes the block of data into the underlying buffer.
    87  func (m *Conn) Write(p []byte) (nn int, err error) {
    88  	var n int
    89  	// 没有足够的空间容纳 p
    90  	for len(p) > m.bufferSize-m.Buffered() && err == nil {
    91  		if m.Buffered() == 0 {
    92  			// Large write, empty buffer.
    93  			// Write directly from p to avoid copy.
    94  			n, err = m.socket.Write(p)
    95  		} else {
    96  			// write buffer to full state,and flush
    97  			n, err = m.writer.Write(p[:m.bufferSize-m.writer.Len()])
    98  			_, err = m.Flush()
    99  		}
   100  		nn += n
   101  		p = p[n:]
   102  	}
   103  
   104  	if err != nil {
   105  		return nn, err
   106  	}
   107  
   108  	// 未到达时间频率的间隔,直接写到缓存
   109  	if m.limit.Limit() {
   110  		n, err = m.writer.Write(p)
   111  		return nn + n, err
   112  	}
   113  
   114  	// 缓存中有数据,flush
   115  	if m.Buffered() > 0 {
   116  		n, err = m.writer.Write(p)
   117  		_, err = m.Flush()
   118  		return nn + n, err
   119  	}
   120  
   121  	// 缓存中无数据,直接写避免内存拷贝
   122  	n, err = m.writeFull(p)
   123  	return nn + n, err
   124  
   125  }
   126  
   127  func (m *Conn) writeFull(p []byte) (nn int, err error) {
   128  	var n int
   129  	for len(p) > 0 && err == nil {
   130  		n, err = m.socket.Write(p)
   131  		nn += n
   132  		p = p[n:]
   133  	}
   134  	return nn, err
   135  }
   136  
   137  // Close closes the connection. Any blocked Read or Write operations will be unblocked
   138  // and return errors.
   139  func (m *Conn) Close() error {
   140  	return m.socket.Close()
   141  }
   142  
   143  // LocalAddr returns the local network address.
   144  func (m *Conn) LocalAddr() net.Addr {
   145  	return m.socket.LocalAddr()
   146  }
   147  
   148  // RemoteAddr returns the remote network address.
   149  func (m *Conn) RemoteAddr() net.Addr {
   150  	return m.socket.RemoteAddr()
   151  }
   152  
   153  // SetDeadline sets the read and write deadlines associated
   154  // with the connection. It is equivalent to calling both
   155  // SetReadDeadline and SetWriteDeadline.
   156  func (m *Conn) SetDeadline(t time.Time) error {
   157  	return m.socket.SetDeadline(t)
   158  }
   159  
   160  // SetReadDeadline sets the deadline for future Read calls
   161  // and any currently-blocked Read call.
   162  func (m *Conn) SetReadDeadline(t time.Time) error {
   163  	return m.socket.SetReadDeadline(t)
   164  }
   165  
   166  // SetWriteDeadline sets the deadline for future Write calls
   167  // and any currently-blocked Write call.
   168  func (m *Conn) SetWriteDeadline(t time.Time) error {
   169  	return m.socket.SetWriteDeadline(t)
   170  }
   171  
   172  // Option 配置 Conn 的选项接口
   173  type Option interface {
   174  	apply(*Conn)
   175  }
   176  
   177  // OptionFunc 包装函数以便它满足 Option 接口
   178  type optionFunc func(*Conn)
   179  
   180  func (f optionFunc) apply(c *Conn) {
   181  	f(c)
   182  }
   183  
   184  // FlushRate Conn 写操作的每秒刷新频率
   185  func FlushRate(r int) Option {
   186  	return optionFunc(func(c *Conn) {
   187  		if r < 1 { // 如果不合规,设置成默认值
   188  			r = defaultRate
   189  		}
   190  		c.limit = rate.New(r, time.Second)
   191  	})
   192  }
   193  
   194  // BufferSize Conn 缓冲大小
   195  func BufferSize(bufferSize int) Option {
   196  	return optionFunc(func(c *Conn) {
   197  		if bufferSize < minBufferSize { // 如果不合规,设置成最小值
   198  			bufferSize = minBufferSize
   199  		}
   200  		c.bufferSize = bufferSize
   201  	})
   202  }