use atomic to prevent race in tests

canonical · Aug 26, 2024 · b46df24 · b46df24
1 parent 37edc0b
commit b46df24
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 11 deletions.
diff --git a/internal/protocol/connector.go b/internal/protocol/connector.go
@@ -8,6 +8,7 @@ import (
 	"net"
 	"sort"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/Rican7/retry"
@@ -252,10 +253,10 @@ func Handshake(ctx context.Context, conn net.Conn, version uint64) (*Protocol, e
 		defer conn.SetDeadline(time.Time{})
 	}
 	// Honor context cancellation.
-	canceled := false
+	var canceled int32 = 0
 	stop := utils.AfterFunc(ctx, func() {
 		if ctx.Err() == context.Canceled {
-			canceled = true
+			atomic.SwapInt32(&canceled, 1)
 			// Cancel read and writes by setting deadline in the past.
 			conn.SetDeadline(time.Unix(1, 0))
 		}
@@ -267,7 +268,7 @@ func Handshake(ctx context.Context, conn net.Conn, version uint64) (*Protocol, e
 	// to the network.
 	n, err := conn.Write(protocol)
 	if err != nil {
-		if canceled && errors.Cause(err).(net.Error).Timeout() {
+		if canceled == 1 && errors.Cause(err).(net.Error).Timeout() {
 			return nil, errors.Wrap(err, "write handshake")
 		}
 		return nil, errors.Wrap(err, "write handshake")

diff --git a/internal/protocol/protocol.go b/internal/protocol/protocol.go
@@ -6,6 +6,7 @@ import (
 	"io"
 	"net"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/pkg/errors"
@@ -62,10 +63,10 @@ func (p *Protocol) Call(ctx context.Context, request, response *Message) (err er
 		defer p.conn.SetDeadline(time.Time{})
 	}
 	// Honor context cancellation.
-	canceled := false
+	var canceled int32 = 0
 	stop := utils.AfterFunc(ctx, func() {
 		if ctx.Err() == context.Canceled {
-			canceled = true
+			atomic.SwapInt32(&canceled, 1)
 			// Cancel read and writes by setting deadline in the past.
 			p.conn.SetDeadline(time.Unix(1, 0))
 		}
@@ -75,14 +76,14 @@ func (p *Protocol) Call(ctx context.Context, request, response *Message) (err er
 	desc := requestDesc(request.mtype)
 
 	if err = p.send(request); err != nil {
-		if canceled && errors.Cause(err).(net.Error).Timeout() {
+		if canceled == 1 && errors.Cause(err).(net.Error).Timeout() {
 			return errors.Wrapf(err, "call %s (canceled): send", desc)
 		}
 		return errors.Wrapf(err, "call %s (budget %s): send", desc, budget)
 	}
 
 	if err = p.recv(response); err != nil {
-		if canceled && errors.Cause(err).(net.Error).Timeout() {
+		if canceled == 1 && errors.Cause(err).(net.Error).Timeout() {
 			return errors.Wrapf(err, "call %s (canceled): receive", desc)
 		}
 		return errors.Wrapf(err, "call %s (budget %s): receive", desc, budget)
@@ -112,10 +113,10 @@ func (p *Protocol) Interrupt(ctx context.Context, request *Message, response *Me
 		defer p.conn.SetDeadline(time.Time{})
 	}
 	// Honor context cancellation.
-	canceled := false
+	var canceled int32 = 0
 	stop := utils.AfterFunc(ctx, func() {
 		if ctx.Err() == context.Canceled {
-			canceled = true
+			atomic.SwapInt32(&canceled, 1)
 			// Cancel read and writes by setting deadline in the past.
 			p.conn.SetDeadline(time.Unix(1, 0))
 		}
@@ -127,15 +128,15 @@ func (p *Protocol) Interrupt(ctx context.Context, request *Message, response *Me
 	// TODO: the context cancelation and honoring should happen in the protocol
 	// primitives which can be tested better.
 	if err := p.send(request); err != nil {
-		if canceled && errors.Cause(err).(net.Error).Timeout() {
+		if canceled == 1 && errors.Cause(err).(net.Error).Timeout() {
 			return errors.Wrapf(err, "interrupt request (canceled): send")
 		}
 		return errors.Wrapf(err, "interrupt request (budget %s): send", budget)
 	}
 
 	for {
 		if err := p.recv(response); err != nil {
-			if canceled && errors.Cause(err).(net.Error).Timeout() {
+			if canceled == 1 && errors.Cause(err).(net.Error).Timeout() {
 				return errors.Wrapf(err, "interrupt request (canceled): receive")
 			}
 			return errors.Wrapf(err, "interrupt request (budget %s): receive", budget)