Comments

yarpc · Nov 10, 2016 · a8ab052 · a8ab052
1 parent 29e77a3
commit a8ab052
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 28 deletions.
diff --git a/transport/peer/peerlist/single.go b/transport/peer/peerlist/single.go
@@ -3,45 +3,57 @@ package peerlist
 import (
 	"context"
 
+	"sync"
+
 	"go.uber.org/yarpc/transport"
 	"go.uber.org/yarpc/transport/internal/errors"
-
-	"github.com/uber-go/atomic"
 )
 
 type single struct {
-	peerID  transport.PeerIdentifier
-	peer    transport.Peer
-	agent   transport.PeerAgent
-	started atomic.Bool
+	lock sync.RWMutex
+
+	initialPeerID transport.PeerIdentifier
+	peer          transport.Peer
+	agent         transport.Agent
+	started       bool
 }
 
 // NewSingle creates a static PeerList with a single Peer
-func NewSingle(pi transport.PeerIdentifier, agent transport.PeerAgent) transport.PeerList {
+func NewSingle(pi transport.PeerIdentifier, agent transport.Agent) transport.PeerList {
 	return &single{
-		peerID: pi,
-		agent:  agent,
+		initialPeerID: pi,
+		agent:         agent,
+		started:       false,
 	}
 }
 
 func (pl *single) Start() error {
-	if pl.started.Swap(true) {
+	pl.lock.Lock()
+	defer pl.lock.Unlock()
+	if pl.started {
 		return errors.ErrPeerListAlreadyStarted("single")
 	}
-	peer, err := pl.agent.RetainPeer(pl.peerID, pl)
+	pl.started = true
+
+	peer, err := pl.agent.RetainPeer(pl.initialPeerID, pl)
 	if err != nil {
-		pl.started.Swap(false)
+		pl.started = false
 		return err
 	}
 	pl.peer = peer
 	return nil
 }
 
 func (pl *single) Stop() error {
-	if !pl.started.Swap(false) {
+	pl.lock.Lock()
+	defer pl.lock.Unlock()
+
+	if !pl.started {
 		return errors.ErrPeerListNotStarted("single")
 	}
-	err := pl.agent.ReleasePeer(pl.peerID, pl)
+	pl.started = false
+
+	err := pl.agent.ReleasePeer(pl.initialPeerID, pl)
 	if err != nil {
 		return err
 	}
@@ -51,11 +63,14 @@ func (pl *single) Stop() error {
 }
 
 func (pl *single) ChoosePeer(context.Context, *transport.Request) (transport.Peer, error) {
-	if !pl.started.Load() {
+	pl.lock.RLock()
+	defer pl.lock.RUnlock()
+
+	if !pl.started {
 		return nil, errors.ErrPeerListNotStarted("single")
 	}
 	return pl.peer, nil
 }
 
-// NotifyPending when the number of Pending requests changes
+// NotifyStatusChanged when the Peer status changes
 func (pl *single) NotifyStatusChanged(transport.Peer) {}
diff --git a/transport/peer/peerlist/single_test.go b/transport/peer/peerlist/single_test.go
@@ -25,11 +25,11 @@ func TestSingle(t *testing.T) {
 	type testStruct struct {
 		msg                   string
 		pid                   transport.PeerIdentifier
-		agent                 *transporttest.MockPeerAgent
+		agent                 *transporttest.MockAgent
 		appliedFunc           func(*single) error
 		expectedPeerID        transport.PeerIdentifier
 		expectedPeer          transport.Peer
-		expectedAgent         transport.PeerAgent
+		expectedAgent         transport.Agent
 		expectedStarted       bool
 		expectedErr           error
 		expectedChooseResults []expectedChooseResult
@@ -38,7 +38,7 @@ func TestSingle(t *testing.T) {
 		func() (s testStruct) {
 			s.msg = "setup"
 			s.pid = transporttest.NewMockPeerIdentifier(mockCtrl)
-			s.agent = transporttest.NewMockPeerAgent(mockCtrl)
+			s.agent = transporttest.NewMockAgent(mockCtrl)
 
 			s.appliedFunc = func(pl *single) error {
 				return nil
@@ -52,7 +52,7 @@ func TestSingle(t *testing.T) {
 		func() (s testStruct) {
 			s.msg = "stop before start"
 			s.pid = transporttest.NewMockPeerIdentifier(mockCtrl)
-			s.agent = transporttest.NewMockPeerAgent(mockCtrl)
+			s.agent = transporttest.NewMockAgent(mockCtrl)
 
 			s.appliedFunc = func(pl *single) error {
 				return pl.Stop()
@@ -67,7 +67,7 @@ func TestSingle(t *testing.T) {
 		func() (s testStruct) {
 			s.msg = "choose before start"
 			s.pid = transporttest.NewMockPeerIdentifier(mockCtrl)
-			s.agent = transporttest.NewMockPeerAgent(mockCtrl)
+			s.agent = transporttest.NewMockAgent(mockCtrl)
 
 			s.appliedFunc = func(pl *single) error {
 				return nil
@@ -85,7 +85,7 @@ func TestSingle(t *testing.T) {
 		func() (s testStruct) {
 			s.msg = "start and choose"
 			s.pid = transporttest.NewMockPeerIdentifier(mockCtrl)
-			s.agent = transporttest.NewMockPeerAgent(mockCtrl)
+			s.agent = transporttest.NewMockAgent(mockCtrl)
 
 			s.expectedPeer = transporttest.NewMockPeer(mockCtrl)
 			s.agent.EXPECT().RetainPeer(s.pid, gomock.Any()).Return(s.expectedPeer, nil)
@@ -106,7 +106,7 @@ func TestSingle(t *testing.T) {
 		func() (s testStruct) {
 			s.msg = "start with agent error"
 			s.pid = transporttest.NewMockPeerIdentifier(mockCtrl)
-			s.agent = transporttest.NewMockPeerAgent(mockCtrl)
+			s.agent = transporttest.NewMockAgent(mockCtrl)
 
 			s.expectedErr = fmt.Errorf("test error")
 			s.agent.EXPECT().RetainPeer(s.pid, gomock.Any()).Return(nil, s.expectedErr)
@@ -123,7 +123,7 @@ func TestSingle(t *testing.T) {
 		func() (s testStruct) {
 			s.msg = "start twice"
 			s.pid = transporttest.NewMockPeerIdentifier(mockCtrl)
-			s.agent = transporttest.NewMockPeerAgent(mockCtrl)
+			s.agent = transporttest.NewMockAgent(mockCtrl)
 
 			s.expectedPeer = transporttest.NewMockPeer(mockCtrl)
 			s.agent.EXPECT().RetainPeer(s.pid, gomock.Any()).Return(s.expectedPeer, nil)
@@ -142,7 +142,7 @@ func TestSingle(t *testing.T) {
 		func() (s testStruct) {
 			s.msg = "start stop"
 			s.pid = transporttest.NewMockPeerIdentifier(mockCtrl)
-			s.agent = transporttest.NewMockPeerAgent(mockCtrl)
+			s.agent = transporttest.NewMockAgent(mockCtrl)
 
 			peer := transporttest.NewMockPeer(mockCtrl)
 			s.agent.EXPECT().RetainPeer(s.pid, gomock.Any()).Return(peer, nil)
@@ -166,7 +166,7 @@ func TestSingle(t *testing.T) {
 		func() (s testStruct) {
 			s.msg = "start stop release failure"
 			s.pid = transporttest.NewMockPeerIdentifier(mockCtrl)
-			s.agent = transporttest.NewMockPeerAgent(mockCtrl)
+			s.agent = transporttest.NewMockAgent(mockCtrl)
 
 			s.expectedPeer = transporttest.NewMockPeer(mockCtrl)
 			s.agent.EXPECT().RetainPeer(s.pid, gomock.Any()).Return(s.expectedPeer, nil)
@@ -196,9 +196,9 @@ func TestSingle(t *testing.T) {
 
 		assert.Equal(t, tt.expectedErr, err, tt.msg)
 		assert.Equal(t, tt.expectedAgent, pl.agent, tt.msg)
-		assert.Equal(t, tt.expectedPeerID, pl.peerID, tt.msg)
+		assert.Equal(t, tt.expectedPeerID, pl.initialPeerID, tt.msg)
 		assert.Equal(t, tt.expectedPeer, pl.peer, tt.msg)
-		assert.Equal(t, tt.expectedStarted, pl.started.Load(), tt.msg)
+		assert.Equal(t, tt.expectedStarted, pl.started, tt.msg)
 
 		for _, expectedResult := range tt.expectedChooseResults {
 			peer, err := pl.ChoosePeer(context.Background(), &transport.Request{})