github.com/uber/kraken@v0.1.4/lib/torrent/scheduler/connstate/state_test.go (about)

     1  // Copyright (c) 2016-2019 Uber Technologies, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  package connstate
    15  
    16  import (
    17  	"testing"
    18  	"time"
    19  
    20  	"github.com/andres-erbsen/clock"
    21  	"github.com/stretchr/testify/require"
    22  	"go.uber.org/zap"
    23  
    24  	"github.com/uber/kraken/core"
    25  	"github.com/uber/kraken/lib/torrent/networkevent"
    26  	"github.com/uber/kraken/lib/torrent/scheduler/conn"
    27  	"github.com/uber/kraken/lib/torrent/storage"
    28  )
    29  
    30  func testState(config Config, clk clock.Clock) *State {
    31  	return New(config, clk, core.PeerIDFixture(), networkevent.NewTestProducer(), zap.NewNop().Sugar())
    32  }
    33  
    34  func TestStateBlacklist(t *testing.T) {
    35  	require := require.New(t)
    36  
    37  	config := Config{
    38  		BlacklistDuration: 30 * time.Second,
    39  	}
    40  	clk := clock.NewMock()
    41  	s := testState(config, clk)
    42  
    43  	p := core.PeerIDFixture()
    44  	h := core.InfoHashFixture()
    45  
    46  	require.NoError(s.Blacklist(p, h))
    47  	require.True(s.Blacklisted(p, h))
    48  	require.Error(s.Blacklist(p, h))
    49  
    50  	clk.Add(config.BlacklistDuration + 1)
    51  
    52  	require.False(s.Blacklisted(p, h))
    53  	require.NoError(s.Blacklist(p, h))
    54  }
    55  
    56  func TestStateBlacklistSnapshot(t *testing.T) {
    57  	require := require.New(t)
    58  
    59  	config := Config{
    60  		BlacklistDuration: 30 * time.Second,
    61  	}
    62  	clk := clock.NewMock()
    63  	s := testState(config, clk)
    64  
    65  	p := core.PeerIDFixture()
    66  	h := core.InfoHashFixture()
    67  
    68  	require.NoError(s.Blacklist(p, h))
    69  
    70  	expected := []BlacklistedConn{{p, h, config.BlacklistDuration}}
    71  	require.Equal(expected, s.BlacklistSnapshot())
    72  }
    73  
    74  func TestStateClearBlacklist(t *testing.T) {
    75  	require := require.New(t)
    76  
    77  	s := testState(Config{}, clock.NewMock())
    78  
    79  	h := core.InfoHashFixture()
    80  
    81  	var peers []core.PeerID
    82  	for i := 0; i < 10; i++ {
    83  		p := core.PeerIDFixture()
    84  		peers = append(peers, p)
    85  		require.NoError(s.Blacklist(p, h))
    86  		require.True(s.Blacklisted(p, h))
    87  	}
    88  
    89  	s.ClearBlacklist(h)
    90  
    91  	for _, p := range peers {
    92  		require.False(s.Blacklisted(p, h))
    93  	}
    94  }
    95  
    96  func TestStateAddPendingPreventsDuplicates(t *testing.T) {
    97  	require := require.New(t)
    98  
    99  	s := testState(Config{}, clock.New())
   100  
   101  	p := core.PeerIDFixture()
   102  	h := core.InfoHashFixture()
   103  
   104  	require.NoError(s.AddPending(p, h, nil))
   105  
   106  	require.Equal(ErrConnAlreadyPending, s.AddPending(p, h, nil))
   107  }
   108  
   109  func TestStateAddPendingReservesCapacity(t *testing.T) {
   110  	require := require.New(t)
   111  
   112  	config := Config{
   113  		MaxOpenConnectionsPerTorrent: 10,
   114  	}
   115  	s := testState(config, clock.New())
   116  
   117  	h := core.InfoHashFixture()
   118  
   119  	for i := 0; i < config.MaxOpenConnectionsPerTorrent; i++ {
   120  		require.NoError(s.AddPending(core.PeerIDFixture(), h, nil))
   121  	}
   122  	require.Equal(ErrTorrentAtCapacity, s.AddPending(core.PeerIDFixture(), h, nil))
   123  }
   124  
   125  func TestStateDeletePendingAllowsFutureAddPending(t *testing.T) {
   126  	require := require.New(t)
   127  
   128  	s := testState(Config{}, clock.New())
   129  
   130  	p := core.PeerIDFixture()
   131  	h := core.InfoHashFixture()
   132  
   133  	require.NoError(s.AddPending(p, h, nil))
   134  	s.DeletePending(p, h)
   135  	require.NoError(s.AddPending(p, h, nil))
   136  }
   137  
   138  func TestStateDeletePendingFreesCapacity(t *testing.T) {
   139  	require := require.New(t)
   140  
   141  	s := testState(Config{MaxOpenConnectionsPerTorrent: 1}, clock.New())
   142  
   143  	h := core.InfoHashFixture()
   144  	p1 := core.PeerIDFixture()
   145  	p2 := core.PeerIDFixture()
   146  
   147  	require.NoError(s.AddPending(p1, h, nil))
   148  	require.Equal(ErrTorrentAtCapacity, s.AddPending(p2, h, nil))
   149  	s.DeletePending(p1, h)
   150  	require.NoError(s.AddPending(p2, h, nil))
   151  }
   152  
   153  func TestStateMovePendingToActivePreventsFuturePending(t *testing.T) {
   154  	require := require.New(t)
   155  
   156  	s := testState(Config{}, clock.New())
   157  
   158  	c, cleanup := conn.Fixture()
   159  	defer cleanup()
   160  
   161  	require.NoError(s.AddPending(c.PeerID(), c.InfoHash(), nil))
   162  	require.NoError(s.MovePendingToActive(c))
   163  	require.Equal(ErrConnAlreadyActive, s.AddPending(c.PeerID(), c.InfoHash(), nil))
   164  }
   165  
   166  func TestStateMovePendingToActiveRejectsNonPendingConns(t *testing.T) {
   167  	require := require.New(t)
   168  
   169  	s := testState(Config{}, clock.New())
   170  
   171  	c, cleanup := conn.Fixture()
   172  	defer cleanup()
   173  
   174  	require.Equal(ErrInvalidActiveTransition, s.MovePendingToActive(c))
   175  
   176  	require.NoError(s.AddPending(c.PeerID(), c.InfoHash(), nil))
   177  	require.NoError(s.MovePendingToActive(c))
   178  	require.Equal(ErrInvalidActiveTransition, s.MovePendingToActive(c))
   179  }
   180  
   181  func TestStateMovePendingToActiveRejectsClosedConns(t *testing.T) {
   182  	require := require.New(t)
   183  
   184  	s := testState(Config{}, clock.New())
   185  
   186  	c, cleanup := conn.Fixture()
   187  	defer cleanup()
   188  
   189  	require.NoError(s.AddPending(c.PeerID(), c.InfoHash(), nil))
   190  	c.Close()
   191  	require.Equal(ErrConnClosed, s.MovePendingToActive(c))
   192  }
   193  
   194  func TestStateDeleteActiveFreesCapacity(t *testing.T) {
   195  	require := require.New(t)
   196  
   197  	s := testState(Config{MaxOpenConnectionsPerTorrent: 1}, clock.New())
   198  
   199  	c, cleanup := conn.Fixture()
   200  	defer cleanup()
   201  
   202  	p2 := core.PeerIDFixture()
   203  
   204  	require.NoError(s.AddPending(c.PeerID(), c.InfoHash(), nil))
   205  	require.NoError(s.MovePendingToActive(c))
   206  	require.Equal(ErrTorrentAtCapacity, s.AddPending(p2, c.InfoHash(), nil))
   207  	s.DeleteActive(c)
   208  	require.NoError(s.AddPending(p2, c.InfoHash(), nil))
   209  }
   210  
   211  func TestStateDeleteActiveNoopsWhenConnIsNotActive(t *testing.T) {
   212  	require := require.New(t)
   213  
   214  	s := testState(Config{MaxOpenConnectionsPerTorrent: 1}, clock.New())
   215  
   216  	c, cleanup := conn.Fixture()
   217  	defer cleanup()
   218  
   219  	require.NoError(s.AddPending(core.PeerIDFixture(), c.InfoHash(), nil))
   220  
   221  	s.DeleteActive(c)
   222  
   223  	require.Equal(ErrTorrentAtCapacity, s.AddPending(core.PeerIDFixture(), c.InfoHash(), nil))
   224  }
   225  
   226  func TestStateActiveConns(t *testing.T) {
   227  	require := require.New(t)
   228  
   229  	s := testState(Config{}, clock.New())
   230  
   231  	conns := make(map[core.PeerID]*conn.Conn)
   232  	for i := 0; i < 10; i++ {
   233  		c, cleanup := conn.Fixture()
   234  		defer cleanup()
   235  
   236  		conns[c.PeerID()] = c
   237  
   238  		require.NoError(s.AddPending(c.PeerID(), c.InfoHash(), nil))
   239  		require.NoError(s.MovePendingToActive(c))
   240  	}
   241  
   242  	result := s.ActiveConns()
   243  	require.Len(result, len(conns))
   244  	for _, c := range result {
   245  		require.Equal(conns[c.PeerID()], c)
   246  	}
   247  
   248  	for _, c := range conns {
   249  		s.DeleteActive(c)
   250  	}
   251  	require.Empty(s.ActiveConns())
   252  }
   253  
   254  func TestStateSaturated(t *testing.T) {
   255  	require := require.New(t)
   256  
   257  	s := testState(Config{MaxOpenConnectionsPerTorrent: 10}, clock.New())
   258  
   259  	info := storage.TorrentInfoFixture(1, 1)
   260  
   261  	var conns []*conn.Conn
   262  	for i := 0; i < 10; i++ {
   263  		c, _, cleanup := conn.PipeFixture(conn.Config{}, info)
   264  		defer cleanup()
   265  
   266  		require.NoError(s.AddPending(c.PeerID(), info.InfoHash(), nil))
   267  		conns = append(conns, c)
   268  	}
   269  
   270  	// Pending conns do not count towards saturated.
   271  	require.False(s.Saturated(info.InfoHash()))
   272  
   273  	for i := 0; i < 9; i++ {
   274  		require.NoError(s.MovePendingToActive(conns[i]))
   275  		require.False(s.Saturated(info.InfoHash()))
   276  	}
   277  
   278  	// Adding 10th conn should mean we're saturated.
   279  	require.NoError(s.MovePendingToActive(conns[9]))
   280  	require.True(s.Saturated(info.InfoHash()))
   281  
   282  	// Removing one should mean we're no longer saturated.
   283  	s.DeleteActive(conns[5])
   284  	require.False(s.Saturated(info.InfoHash()))
   285  }
   286  
   287  func TestMaxMutualConns(t *testing.T) {
   288  	require := require.New(t)
   289  
   290  	mutualConnLimit := 5
   291  	s := testState(Config{
   292  		MaxMutualConnections: mutualConnLimit, MaxOpenConnectionsPerTorrent: 20}, clock.New())
   293  
   294  	neighbors := make([]core.PeerID, 10)
   295  	h := core.InfoHashFixture()
   296  	for i := 0; i < 10; i++ {
   297  		peerID := core.PeerIDFixture()
   298  		neighbors[i] = peerID
   299  		require.NoError(s.AddPending(peerID, h, nil))
   300  	}
   301  	require.Equal(s.AddPending(core.PeerIDFixture(), h, neighbors), ErrTooManyMutualConns)
   302  	require.Equal(s.AddPending(core.PeerIDFixture(), h, neighbors[:mutualConnLimit+1]), ErrTooManyMutualConns)
   303  	require.NoError(s.AddPending(core.PeerIDFixture(), h, neighbors[:mutualConnLimit]))
   304  }