github.com/MetalBlockchain/metalgo@v1.11.9/snow/networking/handler/message_queue_test.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package handler
     5  
     6  import (
     7  	"context"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/prometheus/client_golang/prometheus"
    12  	"github.com/stretchr/testify/require"
    13  	"go.uber.org/mock/gomock"
    14  
    15  	"github.com/MetalBlockchain/metalgo/ids"
    16  	"github.com/MetalBlockchain/metalgo/message"
    17  	"github.com/MetalBlockchain/metalgo/proto/pb/p2p"
    18  	"github.com/MetalBlockchain/metalgo/snow/networking/tracker"
    19  	"github.com/MetalBlockchain/metalgo/snow/validators"
    20  	"github.com/MetalBlockchain/metalgo/utils/constants"
    21  	"github.com/MetalBlockchain/metalgo/utils/logging"
    22  )
    23  
    24  func TestQueue(t *testing.T) {
    25  	ctrl := gomock.NewController(t)
    26  	require := require.New(t)
    27  	cpuTracker := tracker.NewMockTracker(ctrl)
    28  	vdrs := validators.NewManager()
    29  	vdr1ID, vdr2ID := ids.GenerateTestNodeID(), ids.GenerateTestNodeID()
    30  	require.NoError(vdrs.AddStaker(constants.PrimaryNetworkID, vdr1ID, nil, ids.Empty, 1))
    31  	require.NoError(vdrs.AddStaker(constants.PrimaryNetworkID, vdr2ID, nil, ids.Empty, 1))
    32  	mIntf, err := NewMessageQueue(
    33  		logging.NoLog{},
    34  		constants.PrimaryNetworkID,
    35  		vdrs,
    36  		cpuTracker,
    37  		"",
    38  		prometheus.NewRegistry(),
    39  	)
    40  	require.NoError(err)
    41  	u := mIntf.(*messageQueue)
    42  	currentTime := time.Now()
    43  	u.clock.Set(currentTime)
    44  
    45  	msg1 := Message{
    46  		InboundMessage: message.InboundPullQuery(
    47  			ids.Empty,
    48  			0,
    49  			time.Second,
    50  			ids.GenerateTestID(),
    51  			0,
    52  			vdr1ID,
    53  		),
    54  		EngineType: p2p.EngineType_ENGINE_TYPE_UNSPECIFIED,
    55  	}
    56  
    57  	// Push then pop should work regardless of usage when there are no other
    58  	// messages on [u.msgs]
    59  	cpuTracker.EXPECT().Usage(vdr1ID, gomock.Any()).Return(0.1).Times(1)
    60  	u.Push(context.Background(), msg1)
    61  	require.Equal(1, u.nodeToUnprocessedMsgs[vdr1ID])
    62  	require.Equal(1, u.Len())
    63  	_, gotMsg1, ok := u.Pop()
    64  	require.True(ok)
    65  	require.Empty(u.nodeToUnprocessedMsgs)
    66  	require.Zero(u.Len())
    67  	require.Equal(msg1, gotMsg1)
    68  
    69  	cpuTracker.EXPECT().Usage(vdr1ID, gomock.Any()).Return(0.0).Times(1)
    70  	u.Push(context.Background(), msg1)
    71  	require.Equal(1, u.nodeToUnprocessedMsgs[vdr1ID])
    72  	require.Equal(1, u.Len())
    73  	_, gotMsg1, ok = u.Pop()
    74  	require.True(ok)
    75  	require.Empty(u.nodeToUnprocessedMsgs)
    76  	require.Zero(u.Len())
    77  	require.Equal(msg1, gotMsg1)
    78  
    79  	cpuTracker.EXPECT().Usage(vdr1ID, gomock.Any()).Return(1.0).Times(1)
    80  	u.Push(context.Background(), msg1)
    81  	require.Equal(1, u.nodeToUnprocessedMsgs[vdr1ID])
    82  	require.Equal(1, u.Len())
    83  	_, gotMsg1, ok = u.Pop()
    84  	require.True(ok)
    85  	require.Empty(u.nodeToUnprocessedMsgs)
    86  	require.Zero(u.Len())
    87  	require.Equal(msg1, gotMsg1)
    88  
    89  	cpuTracker.EXPECT().Usage(vdr1ID, gomock.Any()).Return(0.0).Times(1)
    90  	u.Push(context.Background(), msg1)
    91  	require.Equal(1, u.nodeToUnprocessedMsgs[vdr1ID])
    92  	require.Equal(1, u.Len())
    93  	_, gotMsg1, ok = u.Pop()
    94  	require.True(ok)
    95  	require.Empty(u.nodeToUnprocessedMsgs)
    96  	require.Zero(u.Len())
    97  	require.Equal(msg1, gotMsg1)
    98  
    99  	// Push msg1 from vdr1ID
   100  	u.Push(context.Background(), msg1)
   101  	require.Equal(1, u.nodeToUnprocessedMsgs[vdr1ID])
   102  	require.Equal(1, u.Len())
   103  
   104  	msg2 := Message{
   105  		InboundMessage: message.InboundPullQuery(
   106  			ids.Empty,
   107  			0,
   108  			time.Second,
   109  			ids.GenerateTestID(),
   110  			0,
   111  			vdr2ID,
   112  		),
   113  		EngineType: p2p.EngineType_ENGINE_TYPE_UNSPECIFIED,
   114  	}
   115  
   116  	// Push msg2 from vdr2ID
   117  	u.Push(context.Background(), msg2)
   118  	require.Equal(2, u.Len())
   119  	require.Equal(1, u.nodeToUnprocessedMsgs[vdr2ID])
   120  	// Set vdr1's usage to 99% and vdr2's to .01
   121  	cpuTracker.EXPECT().Usage(vdr1ID, gomock.Any()).Return(.99).Times(2)
   122  	cpuTracker.EXPECT().Usage(vdr2ID, gomock.Any()).Return(.01).Times(1)
   123  	// Pop should return msg2 first because vdr1 has exceeded it's portion of CPU time
   124  	_, gotMsg2, ok := u.Pop()
   125  	require.True(ok)
   126  	require.Equal(1, u.Len())
   127  	require.Equal(msg2, gotMsg2)
   128  	_, gotMsg1, ok = u.Pop()
   129  	require.True(ok)
   130  	require.Equal(msg1, gotMsg1)
   131  	require.Empty(u.nodeToUnprocessedMsgs)
   132  	require.Zero(u.Len())
   133  
   134  	// u is now empty
   135  	// Non-validators should be able to put messages onto [u]
   136  	nonVdrNodeID1, nonVdrNodeID2 := ids.GenerateTestNodeID(), ids.GenerateTestNodeID()
   137  	msg3 := Message{
   138  		InboundMessage: message.InboundPullQuery(ids.Empty, 0, 0, ids.Empty, 0, nonVdrNodeID1),
   139  		EngineType:     p2p.EngineType_ENGINE_TYPE_UNSPECIFIED,
   140  	}
   141  	msg4 := Message{
   142  		InboundMessage: message.InboundPushQuery(ids.Empty, 0, 0, nil, 0, nonVdrNodeID2),
   143  		EngineType:     p2p.EngineType_ENGINE_TYPE_UNSPECIFIED,
   144  	}
   145  	u.Push(context.Background(), msg3)
   146  	u.Push(context.Background(), msg4)
   147  	u.Push(context.Background(), msg1)
   148  	require.Equal(3, u.Len())
   149  
   150  	// msg1 should get popped first because nonVdrNodeID1 and nonVdrNodeID2
   151  	// exceeded their limit
   152  	cpuTracker.EXPECT().Usage(nonVdrNodeID1, gomock.Any()).Return(.34).Times(1)
   153  	cpuTracker.EXPECT().Usage(nonVdrNodeID2, gomock.Any()).Return(.34).Times(2)
   154  	cpuTracker.EXPECT().Usage(vdr1ID, gomock.Any()).Return(0.0).Times(1)
   155  
   156  	// u.msgs is [msg3, msg4, msg1]
   157  	_, gotMsg1, ok = u.Pop()
   158  	require.True(ok)
   159  	require.Equal(msg1, gotMsg1)
   160  	// u.msgs is [msg3, msg4]
   161  	cpuTracker.EXPECT().Usage(nonVdrNodeID1, gomock.Any()).Return(.51).Times(2)
   162  	_, gotMsg4, ok := u.Pop()
   163  	require.True(ok)
   164  	require.Equal(msg4, gotMsg4)
   165  	// u.msgs is [msg3]
   166  	_, gotMsg3, ok := u.Pop()
   167  	require.True(ok)
   168  	require.Equal(msg3, gotMsg3)
   169  	require.Zero(u.Len())
   170  }