Use contexts for Conn wrapper

wrappers/conn.go

@@ -1,12 +1,14 @@
 package wrappers
 
 import (
+	"context"
 	"net"
 	"time"
 
 	"go.uber.org/zap"
 
 	"github.com/9seconds/mtg/stats"
+	"github.com/juju/errors"
 )
 
 // ConnPurpose is intended to be identifier of connection purpose. We
@@ -39,8 +41,10 @@ const (
 // Conn is a basic wrapper for net.Conn providing the most low-level
 // logic and management as possible.
 type Conn struct {
-	connID string
 	conn   net.Conn
+	ctx    context.Context
+	cancel context.CancelFunc
+	connID string
 	logger *zap.SugaredLogger
 
 	publicIPv4 net.IP
@@ -48,28 +52,46 @@ type Conn struct {
 }
 
 func (c *Conn) Write(p []byte) (int, error) {
-	c.conn.SetWriteDeadline(time.Now().Add(connTimeoutWrite)) // nolint: errcheck
-	n, err := c.conn.Write(p)
+	select {
+	case <-c.ctx.Done():
+		return 0, errors.Annotate(c.ctx.Err(), "Cannot write because context was closed")
+	default:
+		c.conn.SetWriteDeadline(time.Now().Add(connTimeoutWrite)) // nolint: errcheck
+		n, err := c.conn.Write(p)
+		if err != nil {
+			c.cancel()
+		}
 
-	c.logger.Debugw("Write to stream", "bytes", n, "error", err)
-	stats.EgressTraffic(n)
+		c.logger.Debugw("Write to stream", "bytes", n, "error", err)
+		stats.EgressTraffic(n)
 
-	return n, err
+		return n, err
+	}
 }
 
 func (c *Conn) Read(p []byte) (int, error) {
-	c.conn.SetReadDeadline(time.Now().Add(connTimeoutRead)) // nolint: errcheck
-	n, err := c.conn.Read(p)
+	select {
+	case <-c.ctx.Done():
+		return 0, errors.Annotate(c.ctx.Err(), "Cannot read because context was closed")
+	default:
+		c.conn.SetReadDeadline(time.Now().Add(connTimeoutRead)) // nolint: errcheck
+		n, err := c.conn.Read(p)
+		if err != nil {
+			c.cancel()
+		}
 
-	c.logger.Debugw("Read from stream", "bytes", n, "error", err)
-	stats.IngressTraffic(n)
+		c.logger.Debugw("Read from stream", "bytes", n, "error", err)
+		stats.IngressTraffic(n)
 
-	return n, err
+		return n, err
+	}
 }
 
 // Close closes underlying net.Conn instance.
 func (c *Conn) Close() error {
 	defer c.logger.Debugw("Close connection")
+
+	c.cancel()
 	return c.conn.Close()
 }
 
@@ -100,7 +122,8 @@ func (c *Conn) RemoteAddr() *net.TCPAddr {
 }
 
 // NewConn initializes Conn wrapper for net.Conn.
-func NewConn(conn net.Conn, connID string, purpose ConnPurpose, publicIPv4, publicIPv6 net.IP) StreamReadWriteCloser {
+func NewConn(ctx context.Context, cancel context.CancelFunc, conn net.Conn,
+	connID string, purpose ConnPurpose, publicIPv4, publicIPv6 net.IP) StreamReadWriteCloser {
 	logger := zap.S().With(
 		"connection_id", connID,
 		"local_address", conn.LocalAddr(),
@@ -109,9 +132,11 @@ func NewConn(conn net.Conn, connID string, purpose ConnPurpose, publicIPv4, publ
 	).Named("conn")
 
 	wrapper := Conn{
-		logger:     logger,
-		connID:     connID,
 		conn:       conn,
+		ctx:        ctx,
+		cancel:     cancel,
+		connID:     connID,
+		logger:     logger,
 		publicIPv4: publicIPv4,
 		publicIPv6: publicIPv6,
 	}