github.com/cristalhq/netx@v0.0.0-20221116164110-442313ef3309/conn.go (about)

     1  package netx
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"io"
     7  	"net"
     8  	"sync"
     9  )
    10  
    11  // Conn is a stream oriented network connection with i/o operations that are controlled by Contexts
    12  type CtxConn interface {
    13  	net.Conn
    14  	ReadContext(ctx context.Context, b []byte) (n int, err error)
    15  	WriteContext(ctx context.Context, b []byte) (n int, err error)
    16  }
    17  
    18  // Conn is a generic stream-oriented network connection with stats.
    19  //
    20  // Only Read & Write methods contain additional logic for stats.
    21  //
    22  // See: TCPListener or UDPListener how to create it.
    23  type Conn struct {
    24  	net.TCPConn
    25  	stats     *Stats
    26  	closeOnce sync.Once
    27  }
    28  
    29  // ReadContext does same as Read method but with a context.
    30  // This method requires 1 additional goroutine from a worker pool.
    31  func (c *Conn) ReadContext(ctx context.Context, b []byte) (n int, err error) {
    32  	// TODO: bytes pool
    33  	buf := make([]byte, len(b))
    34  	// TODO: chan pool
    35  	ch := make(chan ioResult, 1)
    36  
    37  	// TODO: goroutine pool
    38  	go func() { ch <- newIOResult(c.Read(buf)) }()
    39  
    40  	select {
    41  	case res := <-ch:
    42  		copy(b, buf)
    43  		return res.n, res.err
    44  	case <-ctx.Done():
    45  		return 0, ctx.Err()
    46  	}
    47  }
    48  
    49  // WriteContext does same as Write method but with a context.
    50  // This method requires 1 additional goroutine from a worker pool.
    51  func (c *Conn) WriteContext(ctx context.Context, b []byte) (n int, err error) {
    52  	// TODO: bytes pool
    53  	buf := make([]byte, len(b))
    54  	copy(buf, b)
    55  	// TODO: chan pool
    56  	ch := make(chan ioResult, 1)
    57  
    58  	// TODO: goroutine pool
    59  	go func() { ch <- newIOResult(c.Write(buf)) }()
    60  
    61  	select {
    62  	case res := <-ch:
    63  		return res.n, res.err
    64  	case <-ctx.Done():
    65  		return 0, ctx.Err()
    66  	}
    67  }
    68  
    69  // Read reads data from the connection.
    70  // Read can be made to time out and return an error after a fixed
    71  // time limit; see SetDeadline and SetReadDeadline.
    72  func (c *Conn) Read(p []byte) (int, error) {
    73  	n, err := c.TCPConn.Read(p)
    74  	c.stats.readBytesAdd(n)
    75  	if err != nil && err != io.EOF {
    76  		var ne net.Error
    77  		if errors.As(err, &ne) && ne.Timeout() {
    78  			c.stats.readTimeoutsInc()
    79  		} else {
    80  			c.stats.readErrorsInc()
    81  		}
    82  	}
    83  	return n, err
    84  }
    85  
    86  // Write writes data to the connection.
    87  // Write can be made to time out and return an error after a fixed
    88  // time limit; see SetDeadline and SetWriteDeadline.
    89  func (c *Conn) Write(p []byte) (int, error) {
    90  	n, err := c.TCPConn.Write(p)
    91  	c.stats.writtenBytesAdd(n)
    92  	if err != nil {
    93  		var ne net.Error
    94  		if errors.As(err, &ne) && ne.Timeout() {
    95  			c.stats.writeTimeoutsInc()
    96  		} else {
    97  			c.stats.writeErrorsInc()
    98  		}
    99  	}
   100  	return n, err
   101  }
   102  
   103  func (c *Conn) Close() error {
   104  	var err error
   105  	c.closeOnce.Do(func() {
   106  		err = c.TCPConn.Close()
   107  		c.stats.connsInc()
   108  		if err != nil {
   109  			c.stats.closeErrorsInc()
   110  		}
   111  	})
   112  	return err
   113  }
   114  
   115  type ioResult struct {
   116  	n   int
   117  	err error
   118  }
   119  
   120  func newIOResult(n int, err error) ioResult {
   121  	return ioResult{n: n, err: err}
   122  }