X-Git-Url: https://git.giacomolli.uk/onehouse/blobdiff_plain/c87607884883cca4e47736db79a11c03f175138f..9ca7956528b2d1720cc0e63ca168fd8518609a28:/driver/m31/tcp_client.go

diff --git a/driver/m31/tcp_client.go b/driver/m31/tcp_client.go
index 972f698..7ee6bbe 100644
--- a/driver/m31/tcp_client.go
+++ b/driver/m31/tcp_client.go
@@ -9,11 +9,12 @@ import (
 	"io"
 	"net"
 	"os"
+	"sync"
 	"time"
 )
 
 const (
-	connectTimeout = 2 * time.Second
+	connectTimeout = time.Second
 	writeTimeout   = time.Second
 	readTimeout    = 2 * time.Second
 )
@@ -25,7 +26,7 @@ type tcpClient struct {
 }
 
 // NewTCPClient supplies Modbus TCP functions 01, 02, 03, 05, 15, and 16 without
-// opening a connection. Connect has a two-second timeout. Writes have a
+// opening a connection. Connect has a one-second timeout. Writes have a
 // one-second timeout, and reads have a two-second timeout. Earlier caller
 // deadlines apply to each operation. The caller owns Close.
 func NewTCPClient() Client {
@@ -187,16 +188,17 @@ func (c *tcpClient) exchange(ctx context.Context, function byte, data []byte) (r
 		return nil, fmt.Errorf("set Modbus write deadline: %w", err)
 	}
 
-	err = c.conn.SetReadDeadline(operationDeadline(ctx, readTimeout))
-	if err != nil {
-		return nil, fmt.Errorf("set Modbus read deadline: %w", err)
-	}
-
 	interrupted := make(chan struct{})
 
-	var interruptErr error
+	var (
+		interruptErr error
+		deadlineMu   sync.Mutex
+	)
 
 	stop := context.AfterFunc(ctx, func() {
+		deadlineMu.Lock()
+		defer deadlineMu.Unlock()
+
 		interruptErr = c.conn.SetDeadline(time.Now())
 
 		close(interrupted)
@@ -222,6 +224,21 @@ func (c *tcpClient) exchange(ctx context.Context, function byte, data []byte) (r
 		return nil, fmt.Errorf("write Modbus request: %w", io.ErrShortWrite)
 	}
 
+	deadlineMu.Lock()
+
+	if ctx.Err() != nil {
+		deadlineMu.Unlock()
+
+		return nil, ctx.Err()
+	}
+
+	err = c.conn.SetReadDeadline(operationDeadline(ctx, readTimeout))
+	deadlineMu.Unlock()
+
+	if err != nil {
+		return nil, fmt.Errorf("set Modbus read deadline: %w", err)
+	}
+
 	var header [7]byte
 
 	_, err = io.ReadFull(c.conn, header[:])