X-Git-Url: https://git.giacomolli.uk/onehouse/blobdiff_plain/af9b6f942b1296e915ff49abae3e0e3ee4068e0f..b9a4f3bdd2609a8b178dd42f313c383ccb59049b:/driver/m31/tcp_client.go?ds=inline

diff --git a/driver/m31/tcp_client.go b/driver/m31/tcp_client.go
index 6484b10..7ee6bbe 100644
--- a/driver/m31/tcp_client.go
+++ b/driver/m31/tcp_client.go
@@ -8,11 +8,13 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"os"
+	"sync"
 	"time"
 )
 
 const (
-	connectTimeout = 2 * time.Second
+	connectTimeout = time.Second
 	writeTimeout   = time.Second
 	readTimeout    = 2 * time.Second
 )
@@ -24,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 {
@@ -186,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)
@@ -204,6 +207,10 @@ func (c *tcpClient) exchange(ctx context.Context, function byte, data []byte) (r
 		if !stop() {
 			<-interrupted
 
+			if err == nil || deadlineOnly(err) {
+				err = ctx.Err()
+			}
+
 			err = errors.Join(err, interruptErr, ctx.Err())
 		}
 	}()
@@ -217,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[:])
@@ -265,3 +287,28 @@ func operationDeadline(ctx context.Context, timeout time.Duration) time.Time {
 
 	return deadline
 }
+
+func deadlineOnly(err error) bool {
+	joined, ok := err.(interface{ Unwrap() []error })
+	if ok {
+		found := false
+
+		for _, cause := range joined.Unwrap() {
+			if cause != nil {
+				found = true
+				if !deadlineOnly(cause) {
+					return false
+				}
+			}
+		}
+
+		return found
+	}
+
+	wrapped, ok := err.(interface{ Unwrap() error })
+	if ok {
+		return deadlineOnly(wrapped.Unwrap())
+	}
+
+	return errors.Is(err, os.ErrDeadlineExceeded)
+}