Repository: tomarrell/lbadd
Pull request #117 (2020-05-06 15:22:35)

Add context to Send and Receive


Patch:
diff --git a/internal/network/errors.go b/internal/network/errors.go
index 789b3fb..5efa7a9 100644
--- a/internal/network/errors.go
+++ b/internal/network/errors.go
@@ -12,4 +12,7 @@ const (
 	// ErrClosed indicates, that the component is already closed, and it cannot
 	// be used anymore.
 	ErrClosed Error = "already closed"
+	// ErrTimeout indicates, that a the operation took longer than allowed.
+	// Maybe there was a deadline from a context.
+	ErrTimeout Error = "timeout"
 )
diff --git a/internal/network/server.go b/internal/network/server.go
index 0643412..99b659b 100644
--- a/internal/network/server.go
+++ b/internal/network/server.go
@@ -1,6 +1,7 @@
 package network
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"net"
@@ -47,10 +48,10 @@ type Conn interface {
 	// Send sends the given payload to the remote part of this connection. The
 	// message will not be chunked, and can be read with a single call to
 	// Conn.Receive.
-	Send([]byte) error
+	Send(context.Context, []byte) error
 	// Receive reads a whole message and returns it in a byte slice. A message
 	// is a byte slice that was sent with a single call to Conn.Send.
-	Receive() ([]byte, error)
+	Receive(context.Context) ([]byte, error)
 }
 
 // ID describes an identifier that is used for connections. An ID has to be
diff --git a/internal/network/tcp_conn.go b/internal/network/tcp_conn.go
index cffc93a..1302144 100644
--- a/internal/network/tcp_conn.go
+++ b/internal/network/tcp_conn.go
@@ -4,6 +4,8 @@ import (
 	"encoding/binary"
 	"fmt"
 	"net"
+	"sync"
+	"time"
 
 	"golang.org/x/net/context"
 	"golang.org/x/sync/errgroup"
@@ -20,20 +22,29 @@ var (
 var _ Conn = (*tcpConn)(nil)
 
 type tcpConn struct {
-	id         ID
+	id     ID
+	closed bool
+
+	readLock   sync.Mutex
+	writeLock  sync.Mutex
 	underlying net.Conn
-	closed     bool
 }
 
 // DialTCP dials to the given address, assuming a TCP network. The returned Conn
 // is ready to use.
-func DialTCP(addr string) (Conn, error) {
-	conn, err := net.Dial("tcp", addr)
+func DialTCP(ctx context.Context, addr string) (Conn, error) {
+	// dial the remote endpoint
+	var d net.Dialer
+	conn, err := d.DialContext(ctx, "tcp", addr)
 	if err != nil {
 		return nil, fmt.Errorf("dial tcp: %w", err)
 	}
+
+	// create a new connection object
 	tcpConn := newTCPConn(conn)
-	myID, err := tcpConn.Receive()
+
+	// receive the connection ID from the remote endpoint and apply it
+	myID, err := tcpConn.Receive(ctx)
 	if err != nil {
 		_ = tcpConn.Close()
 		return nil, fmt.Errorf("receive ID: %w", err)
@@ -44,6 +55,8 @@ func DialTCP(addr string) (Conn, error) {
 		return nil, fmt.Errorf("parse ID: %w", err)
 	}
 	tcpConn.id = parsedID
+
+	// return the connection object
 	return tcpConn, nil
 }
 
@@ -60,52 +73,132 @@ func (c *tcpConn) ID() ID {
 	return c.id
 }
 
-func (c *tcpConn) Send(payload []byte) error {
+func (c *tcpConn) Send(ctx context.Context, payload []byte) error {
 	if c.closed {
 		return ErrClosed
 	}
 
-	var frameSize [frameSizeBytes]byte
-	byteOrder.PutUint32(frameSize[:], uint32(len(payload)))
-
-	n, err := c.underlying.Write(frameSize[:])
-	if err != nil {
-		return fmt.Errorf("write size: %w", err)
-	}
-	if n != frameSizeBytes {
-		return fmt.Errorf("write bytes: written %v of %v size bytes", n, len(payload))
+	c.writeLock.Lock()
+	defer c.writeLock.Unlock()
+
+	if deadline, ok := ctx.Deadline(); ok {
+		// Set the write deadline on the underlying connection according to the
+		// given context. This write deadline applies to the whole function, so
+		// we only set it once here. On the next write-call, it will be set
+		// again, or will be reset in the else block, to not keep an old
+		// deadline.
+		_ = c.underlying.SetWriteDeadline(deadline)
+	} else {
+		_ = c.underlying.SetWriteDeadline(time.Time{}) // remove the write deadline
 	}
 
-	n, err = c.underlying.Write(payload)
-	if err != nil {
-		return fmt.Errorf("write payload: %w", err)
-	}
-	if n != len(payload) {
-		return fmt.Errorf("write bytes: written %v of %v payload bytes", n, len(payload))
+	select {
+	case err := <-c.sendAsync(payload):
+		return err
+	case <-ctx.Done():
+		return ErrTimeout
 	}
-	return nil
 }
 
-func (c *tcpConn) Receive() ([]byte, error) {
-	var frameSizeB [frameSizeBytes]byte
-	n, err := c.underlying.Read(frameSizeB[:])
-	if err != nil {
-		return nil, fmt.Errorf("read frame size: %w", err)
-	}
-	if n != frameSizeBytes {
-		return nil, fmt.Errorf("read only %v frame size bytes of %v expected", n, frameSizeBytes)
+func (c *tcpConn) sendAsync(payload []byte) chan error {
+	result := make(chan error)
+	go func() {
+		var frameSize [frameSizeBytes]byte
+		byteOrder.PutUint32(frameSize[:], uint32(len(payload)))
+
+		n, err := c.underlying.Write(frameSize[:])
+		if err != nil {
+			// if the error is a timeout, yield a plain timeout error
+			if netErr, ok := err.(*net.OpError); ok && netErr.Timeout() {
+				result <- ErrTimeout
+				return
+			}
+			result <- fmt.Errorf("write size: %w", err)
+			return
+		}
+		if n != frameSizeBytes {
+			result <- fmt.Errorf("write bytes: written %v of %v size bytes", n, len(payload))
+			return
+		}
+
+		n, err = c.underlying.Write(payload)
+		if err != nil {
+			result <- fmt.Errorf("write payload: %w", err)
+			return
+		}
+		if n != len(payload) {
+			result <- fmt.Errorf("write bytes: written %v of %v payload bytes", n, len(payload))
+			return
+		}
+
+		result <- nil
+	}()
+	return result
+}
+
+func (c *tcpConn) Receive(ctx context.Context) ([]byte, error) {
+	if c.closed {
+		return nil, ErrClosed
 	}
 
-	frameSize := byteOrder.Uint32(frameSizeB[:])
-	frameData := make([]byte, frameSize)
-	n, err = c.underlying.Read(frameData)
-	if err != nil {
-		return nil, fmt.Errorf("read frame payload: %w", err)
+	c.readLock.Lock()
+	defer c.readLock.Unlock()
+
+	if deadline, ok := ctx.Deadline(); ok {
+		// Set the read deadline on the underlying connection according to the
+		// given context. This read deadline applies to the whole function, so
+		// we only set it once here. On the next read-call, it will be set
+		// again, or will be reset in the else block, to not keep an old
+		// deadline.
+		_ = c.underlying.SetReadDeadline(deadline)
+	} else {
+		_ = c.underlying.SetReadDeadline(time.Time{}) // remove the read deadline
 	}
-	if n != int(frameSize) {
-		return nil, fmt.Errorf("read only %v frame payload bytes of %v expected", n, frameSize)
+
+	select {
+	case res := <-c.receiveAsync():
+		if err, ok := res.(error); ok {
+			return nil, err
+		}
+		return res.([]byte), nil
+	case <-ctx.Done():
+		return nil, ErrTimeout
 	}
-	return frameData, nil
+}
+
+func (c *tcpConn) receiveAsync() chan interface{} {
+	result := make(chan interface{})
+	go func() {
+		var frameSizeB [frameSizeBytes]byte
+		n, err := c.underlying.Read(frameSizeB[:])
+		if err != nil {
+			// if the error is a timeout, yield a plain timeout error
+			if netErr, ok := err.(*net.OpError); ok && netErr.Timeout() {
+				result <- ErrTimeout
+				return
+			}
+			result <- fmt.Errorf("read frame size: %w", err)
+			return
+		}
+		if n != frameSizeBytes {
+			result <- fmt.Errorf("read only %v frame size bytes of %v expected", n, frameSizeBytes)
+			return
+		}
+
+		frameSize := byteOrder.Uint32(frameSizeB[:])
+		frameData := make([]byte, frameSize)
+		n, err = c.underlying.Read(frameData)
+		if err != nil {
+			result <- fmt.Errorf("read frame payload: %w", err)
+			return
+		}
+		if n != int(frameSize) {
+			result <- fmt.Errorf("read only %v frame payload bytes of %v expected", n, frameSize)
+			return
+		}
+		result <- frameData
+	}()
+	return result
 }
 
 func (c *tcpConn) Close() error {
diff --git a/internal/network/tcp_server.go b/internal/network/tcp_server.go
index 7faec1f..097c097 100644
--- a/internal/network/tcp_server.go
+++ b/internal/network/tcp_server.go
@@ -3,6 +3,7 @@ package network
 import (
 	"fmt"
 	"net"
+	"time"
 
 	"github.com/rs/zerolog"
 	"golang.org/x/net/context"
@@ -96,18 +97,27 @@ func (s *tcpServer) handleIncomingConnections() {
 			continue
 		}
 
-		tcpConn := newTCPConn(conn)
-		err = tcpConn.Send(tcpConn.id.Bytes())
-		if err != nil {
-			s.log.Error().
-				Err(err).
-				Msg("send ID")
-			_ = tcpConn.Close()
-			continue
-		}
+		go s.handleIncomingNetConn(conn)
+	}
+}
 
-		if s.onConnect != nil {
-			go s.onConnect(tcpConn)
-		}
+func (s *tcpServer) handleIncomingNetConn(conn net.Conn) {
+	tcpConn := newTCPConn(conn)
+
+	ctx := context.Background()
+	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	defer cancel()
+
+	err := tcpConn.Send(ctx, tcpConn.id.Bytes())
+	if err != nil {
+		s.log.Error().
+			Err(err).
+			Msg("send ID")
+		_ = tcpConn.Close()
+		return
+	}
+
+	if s.onConnect != nil {
+		s.onConnect(tcpConn)
 	}
 }


Test patch:
diff --git a/internal/network/example_test.go b/internal/network/example_test.go
index c62bbae..e9a6c0a 100644
--- a/internal/network/example_test.go
+++ b/internal/network/example_test.go
@@ -1,6 +1,7 @@
 package network_test
 
 import (
+	"context"
 	"fmt"
 	"log"
 
@@ -11,9 +12,11 @@ import (
 func ExampleServer() {
 	// When using, please don't ignore all the errors as we do here.
 
+	ctx := context.Background()
+
 	srv := network.NewTCPServer(zerolog.Nop()) // or whatever server is available
 	srv.OnConnect(func(conn network.Conn) {
-		_ = conn.Send([]byte("Hello, World!"))
+		_ = conn.Send(ctx, []byte("Hello, World!"))
 	})
 	go func() {
 		if err := srv.Open(":59513"); err != nil {
@@ -23,11 +26,11 @@ func ExampleServer() {
 
 	<-srv.Listening() // wait for the server to come up
 
-	client, _ := network.DialTCP(":59513")
+	client, _ := network.DialTCP(ctx, ":59513")
 	defer func() {
 		_ = client.Close()
 	}()
-	received, _ := client.Receive()
+	received, _ := client.Receive(ctx)
 	fmt.Println(string(received))
 	// Output: Hello, World!
 }
diff --git a/internal/network/tcp_conn_test.go b/internal/network/tcp_conn_test.go
index c45ac6c..7ad6ae2 100644
--- a/internal/network/tcp_conn_test.go
+++ b/internal/network/tcp_conn_test.go
@@ -1,16 +1,21 @@
 package network
 
 import (
+	"context"
 	"net"
 	"strconv"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 )
 
 func TestTCPConnSendReceive(t *testing.T) {
 	assert := assert.New(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
 	conn1, conn2 := net.Pipe()
 	tcpConn1, tcpConn2 := newTCPConn(conn1), newTCPConn(conn2)
 
@@ -20,12 +25,12 @@ func TestTCPConnSendReceive(t *testing.T) {
 	wg.Add(1)
 	go func() {
 		var err error
-		recv, err = tcpConn2.Receive()
+		recv, err = tcpConn2.Receive(ctx)
 		assert.NoError(err)
 		wg.Done()
 	}()
 
-	err := tcpConn1.Send(payload)
+	err := tcpConn1.Send(ctx, payload)
 	assert.NoError(err)
 
 	wg.Wait()
@@ -34,6 +39,9 @@ func TestTCPConnSendReceive(t *testing.T) {
 
 func TestDialTCP(t *testing.T) {
 	assert := assert.New(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
 	payload := []byte("Hello, World!")
 
 	lis, err := net.Listen("tcp", ":0")
@@ -46,18 +54,43 @@ func TestDialTCP(t *testing.T) {
 
 		tcpConn := newTCPConn(conn)
 		srvConnID = tcpConn.ID().String()
-		assert.NoError(tcpConn.Send(tcpConn.ID().Bytes()))
-		assert.NoError(tcpConn.Send(payload))
+		assert.NoError(tcpConn.Send(ctx, tcpConn.ID().Bytes()))
+		assert.NoError(tcpConn.Send(ctx, payload))
 	}()
 
 	port := lis.Addr().(*net.TCPAddr).Port
 
-	conn, err := DialTCP(":" + strconv.Itoa(port))
+	conn, err := DialTCP(ctx, ":"+strconv.Itoa(port))
 	assert.NoError(err)
 	defer func() { assert.NoError(conn.Close()) }()
 	assert.Equal(srvConnID, conn.ID().String())
 
-	recv, err := conn.Receive()
+	recv, err := conn.Receive(ctx)
 	assert.NoError(err)
 	assert.Equal(payload, recv)
 }
+
+func TestTCPConnWriteContext(t *testing.T) {
+	assert := assert.New(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
+	defer cancel()
+
+	conn1, conn2 := net.Pipe()
+	tcpConn1, _ := newTCPConn(conn1), newTCPConn(conn2)
+
+	err := tcpConn1.Send(ctx, []byte("Hello")) // will not be able to write within 10ms, because noone is reading
+	assert.Equal(ErrTimeout, err)
+}
+
+func TestTCPConnReadContext(t *testing.T) {
+	assert := assert.New(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
+	defer cancel()
+
+	conn1, conn2 := net.Pipe()
+	tcpConn1, _ := newTCPConn(conn1), newTCPConn(conn2)
+
+	data, err := tcpConn1.Receive(ctx) // will not be able to receive within 10ms, because noone is writing
+	assert.Equal(ErrTimeout, err)
+	assert.Nil(data)
+}