github.com/celestiaorg/celestia-node@v0.15.0-beta.1/share/p2p/peers/pool_test.go (about)

     1  package peers
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/libp2p/go-libp2p/core/peer"
     9  	"github.com/stretchr/testify/require"
    10  )
    11  
    12  func TestPool(t *testing.T) {
    13  	t.Run("add / remove peers", func(t *testing.T) {
    14  		p := newPool(time.Second)
    15  
    16  		peers := []peer.ID{"peer1", "peer1", "peer2", "peer3"}
    17  		// adding same peer twice should not produce copies
    18  		p.add(peers...)
    19  		require.Equal(t, len(peers)-1, p.activeCount)
    20  
    21  		p.remove("peer1", "peer2")
    22  		require.Equal(t, len(peers)-3, p.activeCount)
    23  
    24  		peerID, ok := p.tryGet()
    25  		require.True(t, ok)
    26  		require.Equal(t, peers[3], peerID)
    27  
    28  		p.remove("peer3")
    29  		p.remove("peer3")
    30  		require.Equal(t, 0, p.activeCount)
    31  		_, ok = p.tryGet()
    32  		require.False(t, ok)
    33  	})
    34  
    35  	t.Run("round robin", func(t *testing.T) {
    36  		p := newPool(time.Second)
    37  
    38  		peers := []peer.ID{"peer1", "peer1", "peer2", "peer3"}
    39  		// adding same peer twice should not produce copies
    40  		p.add(peers...)
    41  		require.Equal(t, 3, p.activeCount)
    42  
    43  		peerID, ok := p.tryGet()
    44  		require.True(t, ok)
    45  		require.Equal(t, peer.ID("peer1"), peerID)
    46  
    47  		peerID, ok = p.tryGet()
    48  		require.True(t, ok)
    49  		require.Equal(t, peer.ID("peer2"), peerID)
    50  
    51  		peerID, ok = p.tryGet()
    52  		require.True(t, ok)
    53  		require.Equal(t, peer.ID("peer3"), peerID)
    54  
    55  		peerID, ok = p.tryGet()
    56  		require.True(t, ok)
    57  		require.Equal(t, peer.ID("peer1"), peerID)
    58  
    59  		p.remove("peer2", "peer3")
    60  		require.Equal(t, 1, p.activeCount)
    61  
    62  		// pointer should skip removed items until found active one
    63  		peerID, ok = p.tryGet()
    64  		require.True(t, ok)
    65  		require.Equal(t, peer.ID("peer1"), peerID)
    66  	})
    67  
    68  	t.Run("wait for peer", func(t *testing.T) {
    69  		timeout := time.Second
    70  		shortCtx, cancel := context.WithTimeout(context.Background(), timeout/10)
    71  		t.Cleanup(cancel)
    72  
    73  		longCtx, cancel := context.WithTimeout(context.Background(), timeout)
    74  		t.Cleanup(cancel)
    75  
    76  		p := newPool(time.Second)
    77  		done := make(chan struct{})
    78  
    79  		go func() {
    80  			select {
    81  			case <-p.next(shortCtx):
    82  			case <-shortCtx.Done():
    83  				require.Error(t, shortCtx.Err())
    84  				// unlock longCtx waiter by adding new peer
    85  				p.add("peer1")
    86  			}
    87  		}()
    88  
    89  		go func() {
    90  			defer close(done)
    91  			select {
    92  			case peerID := <-p.next(longCtx):
    93  				require.Equal(t, peer.ID("peer1"), peerID)
    94  			case <-longCtx.Done():
    95  				require.NoError(t, longCtx.Err())
    96  			}
    97  		}()
    98  
    99  		select {
   100  		case <-done:
   101  		case <-longCtx.Done():
   102  			require.NoError(t, longCtx.Err())
   103  		}
   104  	})
   105  
   106  	t.Run("nextIdx got removed", func(t *testing.T) {
   107  		p := newPool(time.Second)
   108  
   109  		peers := []peer.ID{"peer1", "peer2", "peer3"}
   110  		p.add(peers...)
   111  		p.nextIdx = 2
   112  		p.remove(peers[p.nextIdx])
   113  
   114  		// if previous nextIdx was removed, tryGet should iterate until available peer found
   115  		peerID, ok := p.tryGet()
   116  		require.True(t, ok)
   117  		require.Equal(t, peers[0], peerID)
   118  	})
   119  
   120  	t.Run("cleanup", func(t *testing.T) {
   121  		p := newPool(time.Second)
   122  		p.cleanupThreshold = 3
   123  
   124  		peers := []peer.ID{"peer1", "peer2", "peer3", "peer4", "peer5"}
   125  		p.add(peers...)
   126  		require.Equal(t, len(peers), p.activeCount)
   127  
   128  		// point to last element that will be removed, to check how pointer will be updated
   129  		p.nextIdx = len(peers) - 1
   130  
   131  		// remove some, but not trigger cleanup yet
   132  		p.remove(peers[3:]...)
   133  		require.Equal(t, len(peers)-2, p.activeCount)
   134  		require.Equal(t, len(peers), len(p.statuses))
   135  
   136  		// trigger cleanup
   137  		p.remove(peers[2])
   138  		require.Equal(t, len(peers)-3, p.activeCount)
   139  		require.Equal(t, len(peers)-3, len(p.statuses))
   140  
   141  		// nextIdx pointer should be updated after next tryGet
   142  		p.tryGet()
   143  		require.Equal(t, 1, p.nextIdx)
   144  	})
   145  
   146  	t.Run("cooldown blocks get", func(t *testing.T) {
   147  		ttl := time.Second / 10
   148  		p := newPool(ttl)
   149  
   150  		peerID := peer.ID("peer1")
   151  		p.add(peerID)
   152  
   153  		_, ok := p.tryGet()
   154  		require.True(t, ok)
   155  
   156  		p.putOnCooldown(peerID)
   157  		// item should be unavailable
   158  		_, ok = p.tryGet()
   159  		require.False(t, ok)
   160  
   161  		ctx, cancel := context.WithTimeout(context.Background(), ttl*5)
   162  		defer cancel()
   163  		select {
   164  		case <-p.next(ctx):
   165  		case <-ctx.Done():
   166  			t.Fatal("item should be already available")
   167  		}
   168  	})
   169  
   170  	t.Run("put on cooldown removed item should be noop", func(t *testing.T) {
   171  		p := newPool(time.Second)
   172  		p.cleanupThreshold = 3
   173  
   174  		peerID := peer.ID("peer1")
   175  		p.add(peerID)
   176  
   177  		p.remove(peerID)
   178  		p.cleanup()
   179  		p.putOnCooldown(peerID)
   180  
   181  		_, ok := p.tryGet()
   182  		require.False(t, ok)
   183  	})
   184  }