go.uber.org/yarpc@v1.72.1/peer/x/peerheap/list_test.go (about)

     1  // Copyright (c) 2022 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package peerheap
    22  
    23  import (
    24  	"context"
    25  	"sort"
    26  	"testing"
    27  	"time"
    28  
    29  	"github.com/golang/mock/gomock"
    30  	"github.com/stretchr/testify/assert"
    31  	"go.uber.org/multierr"
    32  	"go.uber.org/yarpc/api/peer"
    33  	. "go.uber.org/yarpc/api/peer/peertest"
    34  	intyarpcerrors "go.uber.org/yarpc/internal/yarpcerrors"
    35  	"go.uber.org/yarpc/yarpcerrors"
    36  )
    37  
    38  func newNotRunningError(err error) error {
    39  	return intyarpcerrors.AnnotateWithInfo(yarpcerrors.FromError(err), "%s peer list is not running", "peer heap")
    40  
    41  }
    42  
    43  func TestPeerHeapList(t *testing.T) {
    44  	type testStruct struct {
    45  		msg string
    46  
    47  		// StartWaitTimeout is how long the list will block in starting in Update calls
    48  		startWaitTimeout time.Duration
    49  
    50  		// PeerIDs that will be returned from the transport's OnRetain with "Available" status
    51  		retainedAvailablePeerIDs []string
    52  
    53  		// PeerIDs that will be returned from the transport's OnRetain with "Unavailable" status
    54  		retainedUnavailablePeerIDs []string
    55  
    56  		// PeerIDs that will be released from the transport
    57  		releasedPeerIDs []string
    58  
    59  		// PeerIDs that will return "retainErr" from the transport's OnRetain function
    60  		errRetainedPeerIDs []string
    61  		retainErr          error
    62  
    63  		// PeerIDs that will return "releaseErr" from the transport's OnRelease function
    64  		errReleasedPeerIDs []string
    65  		releaseErr         error
    66  
    67  		// A list of actions that will be applied on the PeerList
    68  		peerListActions []PeerListAction
    69  
    70  		// PeerIDs expected to be in the PeerList's "Available" list after the actions have been applied
    71  		expectedAvailablePeers []string
    72  
    73  		// PeerIDs expected to be in the PeerList's "Unavailable" list after the actions have been applied
    74  		expectedUnavailablePeers []string
    75  
    76  		// Boolean indicating whether the PeerList is "running" after the actions have been applied
    77  		expectedRunning bool
    78  	}
    79  	tests := []testStruct{
    80  		{
    81  			msg:                      "setup",
    82  			retainedAvailablePeerIDs: []string{"1"},
    83  			expectedAvailablePeers:   []string{"1"},
    84  			peerListActions: []PeerListAction{
    85  				StartAction{},
    86  				UpdateAction{AddedPeerIDs: []string{"1"}},
    87  			},
    88  			expectedRunning: true,
    89  		},
    90  		{
    91  			msg:                        "setup with disconnected",
    92  			retainedAvailablePeerIDs:   []string{"1"},
    93  			retainedUnavailablePeerIDs: []string{"2"},
    94  			peerListActions: []PeerListAction{
    95  				StartAction{},
    96  				UpdateAction{AddedPeerIDs: []string{"1", "2"}},
    97  			},
    98  			expectedAvailablePeers:   []string{"1"},
    99  			expectedUnavailablePeers: []string{"2"},
   100  			expectedRunning:          true,
   101  		},
   102  		{
   103  			msg:                      "start",
   104  			retainedAvailablePeerIDs: []string{"1"},
   105  			expectedAvailablePeers:   []string{"1"},
   106  			peerListActions: []PeerListAction{
   107  				StartAction{},
   108  				UpdateAction{AddedPeerIDs: []string{"1"}},
   109  				ChooseAction{
   110  					ExpectedPeer: "1",
   111  				},
   112  			},
   113  			expectedRunning: true,
   114  		},
   115  		{
   116  			msg:                        "start stop",
   117  			retainedAvailablePeerIDs:   []string{"1", "2", "3", "4", "5", "6"},
   118  			retainedUnavailablePeerIDs: []string{"7", "8", "9"},
   119  			releasedPeerIDs:            []string{"1", "2", "3", "4", "5", "6", "7", "8", "9"},
   120  			peerListActions: []PeerListAction{
   121  				StartAction{},
   122  				UpdateAction{AddedPeerIDs: []string{"1", "2", "3", "4", "5", "6", "7", "8", "9"}},
   123  				StopAction{},
   124  				ChooseAction{
   125  					ExpectedErr:         newNotRunningError(yarpcerrors.FailedPreconditionErrorf("could not wait for instance to start running: current state is \"stopped\"")),
   126  					InputContextTimeout: 10 * time.Millisecond,
   127  				},
   128  			},
   129  			expectedRunning: false,
   130  		},
   131  		{
   132  			msg:                      "start many and choose",
   133  			retainedAvailablePeerIDs: []string{"1", "2", "3", "4", "5", "6"},
   134  			expectedAvailablePeers:   []string{"1", "2", "3", "4", "5", "6"},
   135  			peerListActions: []PeerListAction{
   136  				StartAction{},
   137  				UpdateAction{AddedPeerIDs: []string{"1", "2", "3", "4", "5", "6"}},
   138  				ChooseAction{ExpectedPeer: "1"},
   139  				ChooseAction{ExpectedPeer: "2"},
   140  				ChooseAction{ExpectedPeer: "3"},
   141  				ChooseAction{ExpectedPeer: "4"},
   142  				ChooseAction{ExpectedPeer: "5"},
   143  				ChooseAction{ExpectedPeer: "6"},
   144  				ChooseAction{ExpectedPeer: "1"},
   145  			},
   146  			expectedRunning: true,
   147  		},
   148  		{
   149  			msg:                      "assure start is idempotent",
   150  			retainedAvailablePeerIDs: []string{"1"},
   151  			expectedAvailablePeers:   []string{"1"},
   152  			peerListActions: []PeerListAction{
   153  				StartAction{},
   154  				UpdateAction{AddedPeerIDs: []string{"1"}},
   155  				StartAction{},
   156  				StartAction{},
   157  				ChooseAction{
   158  					ExpectedPeer: "1",
   159  				},
   160  			},
   161  			expectedRunning: true,
   162  		},
   163  		{
   164  			msg:                      "stop no start",
   165  			retainedAvailablePeerIDs: []string{},
   166  			releasedPeerIDs:          []string{},
   167  			peerListActions: []PeerListAction{
   168  				StopAction{},
   169  				UpdateAction{
   170  					AddedPeerIDs: []string{"1"},
   171  					ExpectedErr:  newNotRunningError(yarpcerrors.FailedPreconditionErrorf("could not wait for instance to start running: current state is \"stopped\"")),
   172  				},
   173  			},
   174  			expectedRunning: false,
   175  		},
   176  		{
   177  			msg:                "update retain error",
   178  			errRetainedPeerIDs: []string{"1"},
   179  			retainErr:          peer.ErrInvalidPeerType{},
   180  			peerListActions: []PeerListAction{
   181  				StartAction{},
   182  				UpdateAction{AddedPeerIDs: []string{"1"}, ExpectedErr: peer.ErrInvalidPeerType{}},
   183  			},
   184  			expectedRunning: true,
   185  		},
   186  		{
   187  			msg:                      "update retain multiple errors",
   188  			retainedAvailablePeerIDs: []string{"2"},
   189  			errRetainedPeerIDs:       []string{"1", "3"},
   190  			retainErr:                peer.ErrInvalidPeerType{},
   191  			peerListActions: []PeerListAction{
   192  				StartAction{},
   193  				UpdateAction{
   194  					AddedPeerIDs: []string{"1", "2", "3"},
   195  					ExpectedErr:  multierr.Combine(peer.ErrInvalidPeerType{}, peer.ErrInvalidPeerType{}),
   196  				},
   197  			},
   198  			expectedAvailablePeers: []string{"2"},
   199  			expectedRunning:        true,
   200  		},
   201  		{
   202  			msg:                      "start stop release error",
   203  			retainedAvailablePeerIDs: []string{"1"},
   204  			errReleasedPeerIDs:       []string{"1"},
   205  			releaseErr:               peer.ErrTransportHasNoReferenceToPeer{},
   206  			peerListActions: []PeerListAction{
   207  				StartAction{},
   208  				UpdateAction{AddedPeerIDs: []string{"1"}},
   209  				StopAction{
   210  					ExpectedErr: peer.ErrTransportHasNoReferenceToPeer{},
   211  				},
   212  			},
   213  			expectedRunning: false,
   214  		},
   215  		{
   216  			msg:                      "assure stop is idempotent",
   217  			retainedAvailablePeerIDs: []string{"1"},
   218  			errReleasedPeerIDs:       []string{"1"},
   219  			releaseErr:               peer.ErrTransportHasNoReferenceToPeer{},
   220  			peerListActions: []PeerListAction{
   221  				StartAction{},
   222  				UpdateAction{AddedPeerIDs: []string{"1"}},
   223  				ConcurrentAction{
   224  					Actions: []PeerListAction{
   225  						StopAction{
   226  							ExpectedErr: peer.ErrTransportHasNoReferenceToPeer{},
   227  						},
   228  						StopAction{
   229  							ExpectedErr: peer.ErrTransportHasNoReferenceToPeer{},
   230  						},
   231  						StopAction{
   232  							ExpectedErr: peer.ErrTransportHasNoReferenceToPeer{},
   233  						},
   234  					},
   235  				},
   236  			},
   237  			expectedRunning: false,
   238  		},
   239  		{
   240  			msg:                      "start stop release multiple errors",
   241  			retainedAvailablePeerIDs: []string{"1", "2", "3"},
   242  			releasedPeerIDs:          []string{"2"},
   243  			errReleasedPeerIDs:       []string{"1", "3"},
   244  			releaseErr:               peer.ErrTransportHasNoReferenceToPeer{},
   245  			peerListActions: []PeerListAction{
   246  				StartAction{},
   247  				UpdateAction{AddedPeerIDs: []string{"1", "2", "3"}},
   248  				StopAction{
   249  					ExpectedErr: multierr.Combine(
   250  						peer.ErrTransportHasNoReferenceToPeer{},
   251  						peer.ErrTransportHasNoReferenceToPeer{},
   252  					),
   253  				},
   254  			},
   255  			expectedRunning: false,
   256  		},
   257  		{
   258  			msg: "choose before start",
   259  			peerListActions: []PeerListAction{
   260  				ChooseAction{
   261  					ExpectedErr:         newNotRunningError(yarpcerrors.FailedPreconditionErrorf("context finished while waiting for instance to start: context deadline exceeded")),
   262  					InputContextTimeout: 10 * time.Millisecond,
   263  				},
   264  				ChooseAction{
   265  					ExpectedErr:         newNotRunningError(yarpcerrors.FailedPreconditionErrorf("context finished while waiting for instance to start: context deadline exceeded")),
   266  					InputContextTimeout: 10 * time.Millisecond,
   267  				},
   268  			},
   269  			expectedRunning: false,
   270  		},
   271  		{
   272  			msg:                      "update before start",
   273  			startWaitTimeout:         time.Second,
   274  			retainedAvailablePeerIDs: []string{"1"},
   275  			expectedAvailablePeers:   []string{"1"},
   276  			peerListActions: []PeerListAction{
   277  				ConcurrentAction{
   278  					Actions: []PeerListAction{
   279  						UpdateAction{AddedPeerIDs: []string{"1"}},
   280  						StartAction{},
   281  					},
   282  					Wait: 20 * time.Millisecond,
   283  				},
   284  			},
   285  			expectedRunning: true,
   286  		},
   287  		{
   288  			msg:              "update timeout before start",
   289  			startWaitTimeout: 30 * time.Millisecond,
   290  			peerListActions: []PeerListAction{
   291  				ConcurrentAction{
   292  					Actions: []PeerListAction{
   293  						UpdateAction{
   294  							AddedPeerIDs: []string{"1"},
   295  							ExpectedErr:  newNotRunningError(yarpcerrors.FailedPreconditionErrorf("context finished while waiting for instance to start: context deadline exceeded")),
   296  						},
   297  						StartAction{},
   298  					},
   299  					Wait: 50 * time.Millisecond,
   300  				},
   301  			},
   302  			expectedRunning: true,
   303  		},
   304  		{
   305  			msg: "start choose no peers",
   306  			peerListActions: []PeerListAction{
   307  				StartAction{},
   308  				ChooseAction{
   309  					InputContextTimeout: 20 * time.Millisecond,
   310  					ExpectedErr:         newUnavailableError(context.DeadlineExceeded),
   311  				},
   312  			},
   313  			expectedRunning: true,
   314  		},
   315  		{
   316  			msg:                      "start then add",
   317  			retainedAvailablePeerIDs: []string{"1", "2"},
   318  			expectedAvailablePeers:   []string{"1", "2"},
   319  			peerListActions: []PeerListAction{
   320  				StartAction{},
   321  				UpdateAction{AddedPeerIDs: []string{"1"}},
   322  				UpdateAction{AddedPeerIDs: []string{"2"}},
   323  				ChooseAction{ExpectedPeer: "1"},
   324  				ChooseAction{ExpectedPeer: "2"},
   325  				ChooseAction{ExpectedPeer: "1"},
   326  			},
   327  			expectedRunning: true,
   328  		},
   329  		{
   330  			msg:                      "start remove",
   331  			retainedAvailablePeerIDs: []string{"1", "2"},
   332  			expectedAvailablePeers:   []string{"2"},
   333  			releasedPeerIDs:          []string{"1"},
   334  			peerListActions: []PeerListAction{
   335  				StartAction{},
   336  				UpdateAction{AddedPeerIDs: []string{"1", "2"}},
   337  				UpdateAction{RemovedPeerIDs: []string{"1"}},
   338  				ChooseAction{ExpectedPeer: "2"},
   339  			},
   340  			expectedRunning: true,
   341  		},
   342  		{
   343  			msg:                      "start add many and remove many",
   344  			retainedAvailablePeerIDs: []string{"1", "2", "3-r", "4-r", "5-a-r", "6-a-r", "7-a", "8-a"},
   345  			releasedPeerIDs:          []string{"3-r", "4-r", "5-a-r", "6-a-r"},
   346  			expectedAvailablePeers:   []string{"1", "2", "7-a", "8-a"},
   347  			peerListActions: []PeerListAction{
   348  				StartAction{},
   349  				UpdateAction{AddedPeerIDs: []string{"1", "2", "3-r", "4-r"}},
   350  				UpdateAction{
   351  					AddedPeerIDs: []string{"5-a-r", "6-a-r", "7-a", "8-a"},
   352  				},
   353  				UpdateAction{
   354  					RemovedPeerIDs: []string{"5-a-r", "6-a-r", "3-r", "4-r"},
   355  				},
   356  				ChooseAction{ExpectedPeer: "1"},
   357  				ChooseAction{ExpectedPeer: "2"},
   358  				ChooseAction{ExpectedPeer: "7-a"},
   359  				ChooseAction{ExpectedPeer: "8-a"},
   360  				ChooseAction{ExpectedPeer: "1"},
   361  			},
   362  			expectedRunning: true,
   363  		},
   364  		{
   365  			msg:                      "add retain error",
   366  			retainedAvailablePeerIDs: []string{"1", "2"},
   367  			expectedAvailablePeers:   []string{"1", "2"},
   368  			errRetainedPeerIDs:       []string{"3"},
   369  			retainErr:                peer.ErrInvalidPeerType{},
   370  			peerListActions: []PeerListAction{
   371  				StartAction{},
   372  				UpdateAction{AddedPeerIDs: []string{"1", "2"}},
   373  				UpdateAction{
   374  					AddedPeerIDs: []string{"3"},
   375  					ExpectedErr:  peer.ErrInvalidPeerType{},
   376  				},
   377  				ChooseAction{ExpectedPeer: "1"},
   378  				ChooseAction{ExpectedPeer: "2"},
   379  				ChooseAction{ExpectedPeer: "1"},
   380  			},
   381  			expectedRunning: true,
   382  		},
   383  		{
   384  			msg:                      "add duplicate peer",
   385  			retainedAvailablePeerIDs: []string{"1", "2"},
   386  			expectedAvailablePeers:   []string{"1", "2"},
   387  			peerListActions: []PeerListAction{
   388  				StartAction{},
   389  				UpdateAction{AddedPeerIDs: []string{"1", "2"}},
   390  				UpdateAction{
   391  					AddedPeerIDs: []string{"2"},
   392  					ExpectedErr:  peer.ErrPeerAddAlreadyInList("2"),
   393  				},
   394  				ChooseAction{ExpectedPeer: "1"},
   395  				ChooseAction{ExpectedPeer: "2"},
   396  				ChooseAction{ExpectedPeer: "1"},
   397  			},
   398  			expectedRunning: true,
   399  		},
   400  		{
   401  			msg:                      "remove peer not in list",
   402  			retainedAvailablePeerIDs: []string{"1", "2"},
   403  			expectedAvailablePeers:   []string{"1", "2"},
   404  			peerListActions: []PeerListAction{
   405  				StartAction{},
   406  				UpdateAction{AddedPeerIDs: []string{"1", "2"}},
   407  				UpdateAction{
   408  					RemovedPeerIDs: []string{"3"},
   409  					ExpectedErr:    peer.ErrPeerRemoveNotInList("3"),
   410  				},
   411  				ChooseAction{ExpectedPeer: "1"},
   412  				ChooseAction{ExpectedPeer: "2"},
   413  				ChooseAction{ExpectedPeer: "1"},
   414  			},
   415  			expectedRunning: true,
   416  		},
   417  		{
   418  			msg:                      "remove release error",
   419  			retainedAvailablePeerIDs: []string{"1", "2"},
   420  			errReleasedPeerIDs:       []string{"2"},
   421  			releaseErr:               peer.ErrTransportHasNoReferenceToPeer{},
   422  			expectedAvailablePeers:   []string{"1"},
   423  			peerListActions: []PeerListAction{
   424  				StartAction{},
   425  				UpdateAction{AddedPeerIDs: []string{"1", "2"}},
   426  				UpdateAction{
   427  					RemovedPeerIDs: []string{"2"},
   428  					ExpectedErr:    peer.ErrTransportHasNoReferenceToPeer{},
   429  				},
   430  				ChooseAction{ExpectedPeer: "1"},
   431  				ChooseAction{ExpectedPeer: "1"},
   432  			},
   433  			expectedRunning: true,
   434  		},
   435  		// Flaky CI
   436  		// {
   437  		// 	msg: "block until add",
   438  		// 	retainedAvailablePeerIDs: []string{"1"},
   439  		// 	expectedAvailablePeers:   []string{"1"},
   440  		// 	peerListActions: []PeerListAction{
   441  		// 		StartAction{},
   442  		// 		ConcurrentAction{
   443  		// 			Actions: []PeerListAction{
   444  		// 				ChooseAction{
   445  		// 					InputContextTimeout: 200 * time.Millisecond,
   446  		// 					ExpectedPeer:        "1",
   447  		// 				},
   448  		// 				UpdateAction{AddedPeerIDs: []string{"1"}},
   449  		// 			},
   450  		// 			Wait: 20 * time.Millisecond,
   451  		// 		},
   452  		// 		ChooseAction{ExpectedPeer: "1"},
   453  		// 	},
   454  		// 	expectedRunning: true,
   455  		// },
   456  		// {
   457  		// 	msg: "multiple blocking until add",
   458  		// 	retainedAvailablePeerIDs: []string{"1"},
   459  		// 	expectedAvailablePeers:   []string{"1"},
   460  		// 	peerListActions: []PeerListAction{
   461  		// 		StartAction{},
   462  		// 		ConcurrentAction{
   463  		// 			Actions: []PeerListAction{
   464  		// 				ChooseAction{
   465  		// 					InputContextTimeout: 200 * time.Millisecond,
   466  		// 					ExpectedPeer:        "1",
   467  		// 				},
   468  		// 				ChooseAction{
   469  		// 					InputContextTimeout: 200 * time.Millisecond,
   470  		// 					ExpectedPeer:        "1",
   471  		// 				},
   472  		// 				ChooseAction{
   473  		// 					InputContextTimeout: 200 * time.Millisecond,
   474  		// 					ExpectedPeer:        "1",
   475  		// 				},
   476  		// 				UpdateAction{AddedPeerIDs: []string{"1"}},
   477  		// 			},
   478  		// 			Wait: 10 * time.Millisecond,
   479  		// 		},
   480  		// 		ChooseAction{ExpectedPeer: "1"},
   481  		// 	},
   482  		// 	expectedRunning: true,
   483  		// },
   484  		{
   485  			msg:                      "block but added too late",
   486  			retainedAvailablePeerIDs: []string{"1"},
   487  			expectedAvailablePeers:   []string{"1"},
   488  			peerListActions: []PeerListAction{
   489  				StartAction{},
   490  				ConcurrentAction{
   491  					Actions: []PeerListAction{
   492  						ChooseAction{
   493  							InputContextTimeout: 10 * time.Millisecond,
   494  							ExpectedErr:         newUnavailableError(context.DeadlineExceeded),
   495  						},
   496  						UpdateAction{AddedPeerIDs: []string{"1"}},
   497  					},
   498  					Wait: 20 * time.Millisecond,
   499  				},
   500  				ChooseAction{ExpectedPeer: "1"},
   501  			},
   502  			expectedRunning: true,
   503  		},
   504  		// Flaky CI
   505  		// {
   506  		// 	msg: "block until new peer after removal of only peer",
   507  		// 	retainedAvailablePeerIDs: []string{"1", "2"},
   508  		// 	releasedPeerIDs:          []string{"1"},
   509  		// 	expectedAvailablePeers:   []string{"2"},
   510  		// 	peerListActions: []PeerListAction{
   511  		// 		StartAction{},
   512  		// 		UpdateAction{AddedPeerIDs: []string{"1"}},
   513  		// 		UpdateAction{RemovedPeerIDs: []string{"1"}},
   514  		// 		ConcurrentAction{
   515  		// 			Actions: []PeerListAction{
   516  		// 				ChooseAction{
   517  		// 					InputContextTimeout: 200 * time.Millisecond,
   518  		// 					ExpectedPeer:        "2",
   519  		// 				},
   520  		// 				UpdateAction{AddedPeerIDs: []string{"2"}},
   521  		// 			},
   522  		// 			Wait: 20 * time.Millisecond,
   523  		// 		},
   524  		// 		ChooseAction{ExpectedPeer: "2"},
   525  		// 	},
   526  		// 	expectedRunning: true,
   527  		// },
   528  		{
   529  			msg: "no blocking with no context deadline",
   530  			peerListActions: []PeerListAction{
   531  				StartAction{},
   532  				ChooseAction{
   533  					InputContext: context.Background(),
   534  					ExpectedErr:  _noContextDeadlineError,
   535  				},
   536  			},
   537  			expectedRunning: true,
   538  		},
   539  		{
   540  			msg:                        "add unavailable peer",
   541  			retainedAvailablePeerIDs:   []string{"1"},
   542  			retainedUnavailablePeerIDs: []string{"2"},
   543  			expectedAvailablePeers:     []string{"1"},
   544  			expectedUnavailablePeers:   []string{"2"},
   545  			peerListActions: []PeerListAction{
   546  				StartAction{},
   547  				UpdateAction{AddedPeerIDs: []string{"1"}},
   548  				UpdateAction{AddedPeerIDs: []string{"2"}},
   549  				ChooseAction{
   550  					ExpectedPeer:        "1",
   551  					InputContextTimeout: 20 * time.Millisecond,
   552  				},
   553  				ChooseAction{
   554  					ExpectedPeer:        "1",
   555  					InputContextTimeout: 20 * time.Millisecond,
   556  				},
   557  			},
   558  			expectedRunning: true,
   559  		},
   560  		{
   561  			msg:                        "remove unavailable peer",
   562  			retainedUnavailablePeerIDs: []string{"1"},
   563  			releasedPeerIDs:            []string{"1"},
   564  			peerListActions: []PeerListAction{
   565  				StartAction{},
   566  				UpdateAction{AddedPeerIDs: []string{"1"}},
   567  				UpdateAction{RemovedPeerIDs: []string{"1"}},
   568  				ChooseAction{
   569  					InputContextTimeout: 10 * time.Millisecond,
   570  					ExpectedErr:         newUnavailableError(context.DeadlineExceeded),
   571  				},
   572  			},
   573  			expectedRunning: true,
   574  		},
   575  		{
   576  			msg:                        "notify peer is now available",
   577  			retainedUnavailablePeerIDs: []string{"1"},
   578  			expectedAvailablePeers:     []string{"1"},
   579  			peerListActions: []PeerListAction{
   580  				StartAction{},
   581  				UpdateAction{AddedPeerIDs: []string{"1"}},
   582  				ChooseAction{
   583  					InputContextTimeout: 10 * time.Millisecond,
   584  					ExpectedErr:         newUnavailableError(context.DeadlineExceeded),
   585  				},
   586  				NotifyStatusChangeAction{PeerID: "1", NewConnectionStatus: peer.Available},
   587  				ChooseAction{ExpectedPeer: "1"},
   588  			},
   589  			expectedRunning: true,
   590  		},
   591  		{
   592  			msg:                      "notify peer is still available",
   593  			retainedAvailablePeerIDs: []string{"1"},
   594  			expectedAvailablePeers:   []string{"1"},
   595  			peerListActions: []PeerListAction{
   596  				StartAction{},
   597  				UpdateAction{AddedPeerIDs: []string{"1"}},
   598  				ChooseAction{ExpectedPeer: "1"},
   599  				NotifyStatusChangeAction{PeerID: "1", NewConnectionStatus: peer.Available},
   600  				ChooseAction{ExpectedPeer: "1"},
   601  			},
   602  			expectedRunning: true,
   603  		},
   604  		{
   605  			msg:                      "notify peer is now unavailable",
   606  			retainedAvailablePeerIDs: []string{"1"},
   607  			expectedUnavailablePeers: []string{"1"},
   608  			peerListActions: []PeerListAction{
   609  				StartAction{},
   610  				UpdateAction{AddedPeerIDs: []string{"1"}},
   611  				ChooseAction{ExpectedPeer: "1"},
   612  				NotifyStatusChangeAction{PeerID: "1", NewConnectionStatus: peer.Unavailable},
   613  				ChooseAction{
   614  					InputContextTimeout: 10 * time.Millisecond,
   615  					ExpectedErr:         newUnavailableError(context.DeadlineExceeded),
   616  				},
   617  			},
   618  			expectedRunning: true,
   619  		},
   620  		{
   621  			msg:                        "notify peer is still unavailable",
   622  			retainedUnavailablePeerIDs: []string{"1"},
   623  			expectedUnavailablePeers:   []string{"1"},
   624  			peerListActions: []PeerListAction{
   625  				StartAction{},
   626  				UpdateAction{AddedPeerIDs: []string{"1"}},
   627  				NotifyStatusChangeAction{PeerID: "1", NewConnectionStatus: peer.Unavailable},
   628  				ChooseAction{
   629  					InputContextTimeout: 10 * time.Millisecond,
   630  					ExpectedErr:         newUnavailableError(context.DeadlineExceeded),
   631  				},
   632  			},
   633  			expectedRunning: true,
   634  		},
   635  		{
   636  			msg:                      "notify invalid peer",
   637  			retainedAvailablePeerIDs: []string{"1"},
   638  			releasedPeerIDs:          []string{"1"},
   639  			peerListActions: []PeerListAction{
   640  				StartAction{},
   641  				UpdateAction{AddedPeerIDs: []string{"1"}},
   642  				UpdateAction{RemovedPeerIDs: []string{"1"}},
   643  				NotifyStatusChangeAction{PeerID: "1", NewConnectionStatus: peer.Available},
   644  			},
   645  			expectedRunning: true,
   646  		},
   647  		{
   648  			// v: Available, u: Unavailable, a: Added, r: Removed
   649  			msg:                        "notify peer stress test",
   650  			retainedAvailablePeerIDs:   []string{"1v", "2va", "3vau", "4var", "5vaur"},
   651  			retainedUnavailablePeerIDs: []string{"6u", "7ua", "8uav", "9uar", "10uavr"},
   652  			releasedPeerIDs:            []string{"4var", "5vaur", "9uar", "10uavr"},
   653  			expectedAvailablePeers:     []string{"1v", "2va", "8uav"},
   654  			expectedUnavailablePeers:   []string{"3vau", "6u", "7ua"},
   655  			peerListActions: []PeerListAction{
   656  				StartAction{},
   657  				UpdateAction{AddedPeerIDs: []string{"1v", "6u"}},
   658  
   659  				// Added Peers
   660  				UpdateAction{
   661  					AddedPeerIDs: []string{"2va", "3vau", "4var", "5vaur", "7ua", "8uav", "9uar", "10uavr"},
   662  				},
   663  
   664  				ChooseMultiAction{ExpectedPeers: []string{
   665  					"1v", "2va", "3vau", "4var", "5vaur",
   666  					"1v", "2va", "3vau", "4var", "5vaur",
   667  				}},
   668  
   669  				// Change Status to Unavailable
   670  				NotifyStatusChangeAction{PeerID: "3vau", NewConnectionStatus: peer.Unavailable},
   671  				NotifyStatusChangeAction{PeerID: "5vaur", NewConnectionStatus: peer.Unavailable},
   672  
   673  				ChooseMultiAction{ExpectedPeers: []string{"1v", "2va", "4var"}},
   674  				ChooseMultiAction{ExpectedPeers: []string{"1v", "2va", "4var"}},
   675  
   676  				// Change Status to Available
   677  				NotifyStatusChangeAction{PeerID: "8uav", NewConnectionStatus: peer.Available},
   678  				NotifyStatusChangeAction{PeerID: "10uavr", NewConnectionStatus: peer.Available},
   679  
   680  				ChooseMultiAction{ExpectedPeers: []string{"8uav", "10uavr"}}, // realign
   681  				ChooseMultiAction{ExpectedPeers: []string{"1v", "2va", "4var", "8uav", "10uavr"}},
   682  				ChooseMultiAction{ExpectedPeers: []string{"1v", "2va", "4var", "8uav", "10uavr"}},
   683  
   684  				// Remove Peers
   685  				UpdateAction{
   686  					RemovedPeerIDs: []string{"4var", "5vaur", "9uar", "10uavr"},
   687  				},
   688  
   689  				ChooseMultiAction{ExpectedPeers: []string{"1v", "2va", "8uav"}},
   690  				ChooseMultiAction{ExpectedPeers: []string{"1v", "2va", "8uav"}},
   691  			},
   692  			expectedRunning: true,
   693  		},
   694  		// Flaky CI
   695  		// {
   696  		// 	msg: "block until notify available",
   697  		// 	retainedUnavailablePeerIDs: []string{"1"},
   698  		// 	expectedAvailablePeers:     []string{"1"},
   699  		// 	peerListActions: []PeerListAction{
   700  		// 		StartAction{},
   701  		// 		UpdateAction{AddedPeerIDs: []string{"1"}},
   702  		// 		ConcurrentAction{
   703  		// 			Actions: []PeerListAction{
   704  		// 				ChooseAction{
   705  		// 					InputContextTimeout: 200 * time.Millisecond,
   706  		// 					ExpectedPeer:        "1",
   707  		// 				},
   708  		// 				NotifyStatusChangeAction{PeerID: "1", NewConnectionStatus: peer.Available},
   709  		// 			},
   710  		// 			Wait: 20 * time.Millisecond,
   711  		// 		},
   712  		// 		ChooseAction{ExpectedPeer: "1"},
   713  		// 	},
   714  		// 	expectedRunning: true,
   715  		// },
   716  	}
   717  
   718  	for _, tt := range tests {
   719  		t.Run(tt.msg, func(t *testing.T) {
   720  			mockCtrl := gomock.NewController(t)
   721  			defer mockCtrl.Finish()
   722  
   723  			transport := NewMockTransport(mockCtrl)
   724  
   725  			// Healthy Transport Retain/Release
   726  			peerMap := ExpectPeerRetains(
   727  				transport,
   728  				tt.retainedAvailablePeerIDs,
   729  				tt.retainedUnavailablePeerIDs,
   730  			)
   731  			ExpectPeerReleases(transport, tt.releasedPeerIDs, nil)
   732  
   733  			// Unhealthy Transport Retain/Release
   734  			ExpectPeerRetainsWithError(transport, tt.errRetainedPeerIDs, tt.retainErr)
   735  			ExpectPeerReleases(transport, tt.errReleasedPeerIDs, tt.releaseErr)
   736  
   737  			var opts []HeapOption
   738  			if tt.startWaitTimeout != 0 {
   739  				opts = append(opts, StartupWait(tt.startWaitTimeout))
   740  			}
   741  			pl := New(transport, opts...)
   742  
   743  			deps := ListActionDeps{
   744  				Peers: peerMap,
   745  			}
   746  			ApplyPeerListActions(t, pl, tt.peerListActions, deps)
   747  
   748  			var availablePeers []string
   749  			var unavailablePeers []string
   750  			for _, ps := range pl.byScore.peers {
   751  				if ps.status.ConnectionStatus == peer.Available {
   752  					availablePeers = append(availablePeers, ps.id.Identifier())
   753  				} else if ps.status.ConnectionStatus == peer.Unavailable {
   754  					unavailablePeers = append(unavailablePeers, ps.id.Identifier())
   755  				}
   756  			}
   757  			sort.Strings(availablePeers)
   758  			sort.Strings(unavailablePeers)
   759  
   760  			assert.Equal(t, availablePeers, tt.expectedAvailablePeers, "incorrect available peers")
   761  			assert.Equal(t, unavailablePeers, tt.expectedUnavailablePeers, "incorrect unavailable peers")
   762  			assert.Equal(t, tt.expectedRunning, pl.IsRunning(), "Peer list should match expected final running state")
   763  		})
   764  	}
   765  }